Merge branch 'main' into heidelberg26

Author: theoheimel

madgraph/core/helas_objects.py

@@ -3758,7 +3758,7 @@ def default_setup(self):
         self['has_mirror_process'] = False
         self['allowed_flavors'] = [] # list of all allowed flavors for the process
         self['allowed_flavors_with_iden'] = [] # list of all allowed flavors for the process but grouped by identical matrix-element
-        self['allowed_flavors_unique'] = []
+        self['allowed_flavors_with_iden_sign'] = [] # list of all allowed flavors for the process but grouped by identical matrix-element
 
     def filter(self, name, value):
         """Filter for valid diagram property values."""
@@ -5136,12 +5136,12 @@ def get_external_nhel(self):
         return allowed_helicity
 
 
-    def get_external_flavors(self, all_perm=False, preserve_sign=False):
+    def get_external_flavors(self, all_perm=False, return_sign=False):
         """If merged particles are used, determine the list of possible flavor that are not zero """
 
         if self['allowed_flavors']:
-            if preserve_sign:
-                return self['allowed_flavors_sign']
+            if return_sign:
+                return self['allowed_flavors'], self['allowed_flavors_sign']
             else:
                 return self['allowed_flavors']
         pdgs=[]
@@ -5150,12 +5150,8 @@ def get_external_flavors(self, all_perm=False, preserve_sign=False):
         external_number=1
         for wf in external_wfs:
             if wf.get('number_external')==external_number:
-                pdgs.append(
-                    wf.get('particle').get_pdg_code()
-                    if external_number > 2
-                    else wf.get('particle').get_anti_pdg_code()
-                )
                 external_number=external_number+1
+                pdgs.append(wf.get('particle').get_pdg_code())
 
         to_map = {}
         model = self.get('processes')[0].get('model')
@@ -5164,7 +5160,6 @@ def get_external_flavors(self, all_perm=False, preserve_sign=False):
 
         flavor_list = []
         pdg_list = []
-    
         restricted_flavor = [None]*len(external_wfs)
         for i,wf in enumerate(external_wfs):
             if wf.get('flavor'):
@@ -5222,43 +5217,38 @@ def get_external_flavors(self, all_perm=False, preserve_sign=False):
 
         self['allowed_flavors'] = flavor_list
         self['allowed_flavors_sign'] = pdg_list
-        return pdg_list if preserve_sign else flavor_list
+        if return_sign:
+            return self['allowed_flavors'], self['allowed_flavors_sign']
+        else:
+            return self['allowed_flavors']
 
-    def get_external_flavors_with_iden(self, return_pdgs=False):
+    def get_external_flavors_with_iden(self, return_sign=False):
         if self['allowed_flavors_with_iden']:
-            return self['allowed_flavors_with_iden']
+            if return_sign:
+                return self['allowed_flavors_with_iden'], self['allowed_flavors_with_iden_sign']
+            else:
+                return self['allowed_flavors_with_iden']
 
         model = self.get('processes')[0].get('model')
-        all_flv = self.get_external_flavors()
+        all_flv, all_flv_sign = self.get_external_flavors(return_sign=True)
         map_all_flv = {}
-        for i, flv1 in  enumerate(all_flv):
+        map_all_flv_sign = {}
+        for i, (flv1, flv1_sign) in  enumerate(zip(all_flv, all_flv_sign)):
             coup = self.get_coupling_for_flv(flv1, model)
             if coup in map_all_flv:
                 map_all_flv[coup].append(flv1)
+                map_all_flv_sign[coup].append(flv1_sign)
             else:
                 map_all_flv[coup] = [flv1]
+                map_all_flv_sign[coup] = [flv1_sign]
 
+        self['allowed_flavors_with_iden_sign'] = map_all_flv_sign.values()
         self['allowed_flavors_with_iden'] = map_all_flv.values()
-        return self['allowed_flavors_with_iden']
-
-    def get_external_flavors_unique(self):
-        if self['allowed_flavors_unique']:
-            return self['allowed_flavors_unique']
-
-        model = self.get('processes')[0].get('model')
-        all_flv = self.get_external_flavors(preserve_sign=True)
-        map_all_flv = {}
-        for i, flv1 in  enumerate(all_flv):
-            coup = self.get_coupling_for_flv(flv1, model)
-            key = (coup, flv1[0], flv1[1])
-            if key in map_all_flv:
-                map_all_flv[key].append(flv1)
-            else:
-                map_all_flv[key] = [flv1]
+        if return_sign:
+            return self['allowed_flavors_with_iden'], self['allowed_flavors_with_iden_sign']
+        else:
+            return self['allowed_flavors_with_iden']
 
-        self['allowed_flavors_unique'] = list(map_all_flv.values())
-        return self['allowed_flavors_unique']
-    
     def check_flavor(self, real_pdgs, model, debug=False):
         """check if any feynman diagram is compatible with the pdg codes replaced by the real_pdgs"""
         HelasDiagram.done_flavor = []

madgraph/iolibs/export_cpp.py

@@ -915,7 +915,7 @@ def get_process_class_definitions(self, write=True):
                 self.matrix_elements[0].get("diagrams")
             )
             replace_dict['nflav'] = len(
-                self.matrix_elements[0].get_external_flavors_unique()
+                self.matrix_elements[0].get_external_flavors_with_iden()
             )
 
         else:
@@ -1193,9 +1193,9 @@ def get_sigmaKin_lines(self, color_amplitudes, write=True):
                 return replace_dict
 
     def get_flavor_table(self, matrix_element):
-        flavors = matrix_element.get_external_flavors_unique()
-        flavor_multipliers = "{" + ",".join(str(len(f)) for f in flavors) + "}"
-
+        print(list(matrix_element.get_external_flavors()))
+        flavors = list(matrix_element.get_external_flavors_with_iden())
+        print(flavors)
         flavor_dict = {
             1: 0, 2: 1, 3: 2, 4: 3, # quarks
             11: 0, 13: 1, 15: 2,    # charged leptons
@@ -1211,7 +1211,6 @@ def get_flavor_table(self, matrix_element):
         ext_count = len(flavors[0][0])
         return f"""
         static const int flavor_table[{flavor_count}][{ext_count}] = {full_flavor_table};
-        static const int flavor_multipliers[{flavor_count}] = {flavor_multipliers};
         """
 
     def get_all_sigmaKin_lines(self, color_amplitudes, class_name):

madgraph/iolibs/export_mg7.py

@@ -28,8 +28,18 @@ def get_subprocess_info(matrix_element, proc_dir, lib_me_path):
     )
     diagrams = amplitude.get("diagrams")
     helas_diagrams = matrix_element.get("diagrams")
-    all_flavors = matrix_element.get_external_flavors_unique()
-
+    all_flavors, all_flavors_sign = matrix_element.get_external_flavors_with_iden(return_sign=True)
+    all_flavors_same_initial = []
+    all_flavors_indices = []
+    for i, flavors in enumerate(all_flavors_sign):
+        flv_dict = defaultdict(list)
+        for flv in flavors:
+            flv_dict[(flv[0], flv[1])].append(flv)
+        indices = []
+        for flv in flv_dict.values():
+            indices.append(len(all_flavors_same_initial))
+            all_flavors_same_initial.append((i, flv))
+        all_flavors_indices.append(indices)
 
     color_basis = matrix_element.get("color_basis")
     if color_basis:
@@ -59,9 +69,10 @@ def get_subprocess_info(matrix_element, proc_dir, lib_me_path):
 
         helas_diagram = helas_diagrams[diagram_index]
         active_flavors = [
-            i
-            for i, flavors in enumerate(all_flavors)
-            if helas_diagram.check_flavor(flavors[0], model)
+            flav_id
+            for indices, flavors in zip(all_flavors_indices, all_flavors)
+            if helas_diagram.check_flavor([flv for flv in flavors[0]], model)
+            for flav_id in indices
         ]
 
         diagram = diagrams[diagram_index]
@@ -120,9 +131,9 @@ def get_subprocess_info(matrix_element, proc_dir, lib_me_path):
         color_flows = [[
             [[color_flow_dict[leg.get("number")][i] for i in [0, 1]] for leg in legs]
             for color_flow_dict in color_flow_dicts
-        ]] * len(all_flavors) #TODO: this is wrong for multiple flavors!!!
+        ]] * len(all_flavors_same_initial) #TODO: this is wrong for multiple flavors!!!
     else:
-        color_flows = [[[[0, 0]] * n_external]] * len(all_flavors) #TODO: this is wrong for multiple flavors!!!
+        color_flows = [[[[0, 0]] * n_external]] * len(all_flavors_same_initial) #TODO: this is wrong for multiple flavors!!!
 
     # We need the both particle and antiparticle wf_ids, since the identity
     # depends on the direction of the wf.
@@ -146,7 +157,10 @@ def get_subprocess_info(matrix_element, proc_dir, lib_me_path):
                 sign = -1 if part_id < 0 else 1
                 pdg_color_types[sign * pdg] = sign * model.get_particle(part_id).get_color()
 
-    flavors = [{"index": i, "options": options} for i, options in enumerate(all_flavors)]
+    flavors = [
+        {"index": index, "options": options}
+        for index, options in all_flavors_same_initial
+    ]
     return {
         "incoming": incoming,
         "outgoing": outgoing,

madgraph/iolibs/template_files/cpp_process_sigmaKin_function.inc

@@ -8,7 +8,6 @@ const int denominator = %(den_factors)s;
 %(flavor_table)s
 
 const int* flavor = &flavor_table[flavor_id][0];
-int multiplier = flavor_multipliers[flavor_id];
 
 ntry[flavor_id]++;
 
@@ -50,4 +49,4 @@ if (sum_hel[flavor_id] == 0 || ntry[flavor_id] < 10){
     }
 }
 
-return multiplier * matrix_element / (denominator * broken_sym(flavor));
+return matrix_element / (denominator * broken_sym(flavor));