Merging more integration channels within and across subprocesses

madgraph/iolibs/export_cpp.py

@@ -27,6 +27,7 @@
 import shutil
 import subprocess
 import json
+from collections import defaultdict
 
 import madgraph.core.base_objects as base_objects
 import madgraph.core.color_algebra as color
@@ -1544,9 +1545,7 @@ def get_sigmaKin_lines(self, color_amplitudes, write=True):
                 return replace_dict
 
     def get_flavor_table(self, matrix_element):
-        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
@@ -3199,6 +3198,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.me_lib_format = args[1].get("me_lib_format", None)
         self.process_info = []
+        self.merged_subprocesses = defaultdict(list)
 
     def generate_subprocess_directory(
         self, matrix_element, cpp_helas_call_writer, proc_number=None
@@ -3234,7 +3234,11 @@ def generate_subprocess_directory(
             plot.draw()
 
         me_lib_path = self.me_lib_format.format(process_id = proc_dir_name)
-        self.process_info.append(process_exporter_mg7.get_subprocess_info(dirpath, me_lib_path))
+        subproc_info, diagram_tags, subproc_class = process_exporter_mg7.get_subprocess_info(dirpath, me_lib_path)
+        self.merged_subprocesses[subproc_class].append(
+            (len(self.process_info), diagram_tags)
+        )
+        self.process_info.append(subproc_info)
 
     def copy_template(self, model):
         super().copy_template(model)
@@ -3258,13 +3262,96 @@ def copy_template(self, model):
                 )
             os.chmod(madnis_bin, 0o755)
 
+    def get_merged_info(self):
+        merged_subproc_info = []
+        for subprocesses in self.merged_subprocesses.values():
+            channels = []
+            flavors = []
+            subproc_indices = []
+            unique_diagram_tags = []
+            unique_diagrams = []
+            for subproc_index_in_group, (subproc_index, diagram_tags) in enumerate(
+                subprocesses
+            ):
+                subproc_indices.append(subproc_index)
+                subproc_info = self.process_info[subproc_index]
+                flavor_offset = len(flavors)
+                flavors.extend(
+                    {
+                        "subprocess": subproc_index_in_group,
+                        "flavor": ps_flavor,
+                    }
+                    for ps_flavor, flavor in enumerate(
+                        subproc_info["flavors"]
+                    )
+                )
+
+                for chan_info, chan_tags in zip(subproc_info["channels"], diagram_tags):
+                    chan_index = len(channels)
+
+                    same_diags = []
+                    for tag in chan_tags:
+                        try:
+                            index = unique_diagram_tags.index(tag)
+                            chan_index, diag_index = unique_diagrams[index]
+                            same_diags.append(diag_index)
+                        except ValueError:
+                            same_diags.append(None)
+
+                    if chan_index == len(channels):
+                        channels.append(
+                            {
+                                "subprocess": subproc_index,
+                                "channel": chan_index,
+                                "diagrams": [],
+                            }
+                        )
+
+                    diagrams = channels[chan_index]["diagrams"]
+                    for diag_info, same_diag, tag in zip(
+                        chan_info["diagrams"], same_diags, chan_tags
+                    ):
+                        if same_diag is None:
+                            unique_diagram_tags.append(tag)
+                            diag_index = len(diagrams)
+                            unique_diagrams.append((chan_index, diag_index))
+                            diagrams.append(
+                                {
+                                    "diagram": [-1] * len(subprocesses),
+                                    "permutation": diag_info["permutation"],
+                                    "active_flavors": [],
+                                }
+                            )
+                        else:
+                            diag_index = same_diag
+                        diag_dict = diagrams[diag_index]
+                        diag_dict["diagram"][subproc_index_in_group] = diag_info["diagram"]
+                        diag_dict["active_flavors"].extend(
+                            flavor_offset + flav for flav in diag_info["active_flavors"]
+                        )
+
+            merged_subproc_info.append({
+                "incoming": subproc_info["incoming"],
+                "outgoing": subproc_info["outgoing"],
+                "subprocesses": subproc_indices,
+                "channels": channels,
+                "flavors": flavors,
+            })
+        return merged_subproc_info
+
     def finalize(self, matrix_elements=None, history='', *args, **kwargs):
         file_name = os.path.normpath(os.path.join(
             self.dir_path, "SubProcesses", "subprocesses.json"
         ))
         with open(file_name, 'w') as f:
             json.dump(self.process_info, f)
 
+        merged_file_name = os.path.normpath(os.path.join(
+            self.dir_path, "SubProcesses", "merged_subprocesses.json"
+        ))
+        with open(merged_file_name, 'w') as f:
+            json.dump(self.get_merged_info(), f)
+
         # Generate Cards/run_card.toml from the template, filling in
         # process-dependent defaults (mirrors the LO run_card.dat logic).
         self.create_run_card(matrix_elements, history)

madgraph/iolibs/export_mg7.py

@@ -4,6 +4,50 @@
 
 from madgraph.various.diagram_symmetry import find_symmetry, IdentifySGConfigTag
 from madgraph.iolibs import export_cpp
+from madgraph.iolibs.group_subprocs import IdentifyConfigTag
+from madgraph.core.diagram_generation import DiagramTag
+
+class IdentifyJetTag(IdentifyConfigTag):
+    """ Like IndentifyConfigTag, but ignores spin and color """
+
+    @staticmethod
+    def link_from_leg(leg, model):
+        (leg_num1, _, mass, width, _), leg_num2 = super(
+            IdentifyJetTag, IdentifyJetTag
+        ).link_from_leg(leg, model)[0]
+        return [((leg_num1, mass, width), leg_num2)]
+
+    @staticmethod
+    def vertex_id_from_vertex(vertex, last_vertex, model, ninitial):
+        vertex = super(IdentifyJetTag, IdentifyJetTag).vertex_id_from_vertex(
+            vertex, last_vertex, model, ninitial
+        )
+        if len(vertex) == 1:
+            return ((0,),)
+        (_, mass, width), _ = vertex
+        return ((mass, width), 0)
+
+
+class IdentifySGJetTag(IdentifySGConfigTag):
+    """ Like IndentifySGConfigTag, but ignores spin, color and charge """
+
+    @staticmethod
+    def link_from_leg(leg, model):
+        (state, _, _, _, mass, width), leg_num = super(
+            IdentifySGJetTag, IdentifySGJetTag
+        ).link_from_leg(leg, model)[0]
+        return [((state, mass, width), leg_num)]
+
+    @staticmethod
+    def vertex_id_from_vertex(vertex, last_vertex, model, ninitial):
+        vertex = super(IdentifySGJetTag, IdentifySGJetTag).vertex_id_from_vertex(
+            vertex, last_vertex, model, ninitial
+        )
+        if len(vertex) == 1:
+            return ((0,),)
+        (_, mass, width, qcd, onshell), = vertex
+        return ((mass, width, qcd, onshell),)
+
 
 class OneProcessExporterMG7(export_cpp.OneProcessExporterCPP):
 
@@ -13,15 +57,25 @@ def __init__(self, matrix_element, cpp_helas_call_writer):
         self.name = f"P{matrix_element.get('processes')[0].shell_string()}"
         self.model = self.matrix_element.get("processes")[0].get("model")
         self.amplitude = self.matrix_element.get("base_amplitude")
-        self.sym_indices, self.sym_perms, _ = find_symmetry(
-            self.matrix_element, lambda diag: IdentifySGConfigTag(diag, self.model)
-        )
+        merge_jets = False
+        if merge_jets:
+            self.sym_indices, self.sym_perms, _ = find_symmetry(
+                self.matrix_element,
+                lambda diag: IdentifySGJetTag(diag, self.model),
+                skip_identical_check=True
+            )
+        else:
+            self.sym_indices, self.sym_perms, _ = find_symmetry(
+                self.matrix_element, lambda diag: IdentifySGConfigTag(diag, self.model)
+            )
+
         self.diagrams = self.amplitude.get("diagrams")
         self.helas_diagrams = self.matrix_element.get("diagrams")
         self.all_flavors, self.all_flavors_pdgs = self.matrix_element.get_external_flavors_with_iden(return_pdgs=True)
         self.process = self.amplitude.get("process")
         self.legs = self.process.get("legs_with_decays")
         self.color_basis = self.matrix_element.get("color_basis")
+        self.set_subprocess_class()
         self.set_topology()
         self.set_flavor_indices()
         self.set_active_flavors()
@@ -30,6 +84,25 @@ def __init__(self, matrix_element, cpp_helas_call_writer):
     def generate_process_files(self):
         super().generate_process_files()
 
+    def set_subprocess_class(self):
+        is_parts = [
+            self.model.get_particle(l.get("id"))
+            for l in self.process.get("legs")
+            if not l.get("state")
+        ]
+        fs_parts = [
+            self.model.get_particle(l.get("id"))
+            for l in self.process.get("legs")
+            if l.get("state")
+        ]
+        self.subprocess_class = (
+            tuple(
+                (p.get("mass"), l.get("onshell"))
+                for (p, l) in zip(is_parts + fs_parts, self.process.get("legs"))
+            ),
+            self.process.get("id"),
+        )
+
     def set_topology(self):
         self.edge_names = {}
         self.incoming = [None] * 2
@@ -76,15 +149,12 @@ def set_channels_colors_map(self):
                 for diag_tuple in self.color_basis[col_basis_elem]:
                     diag_jamps[diag_tuple[0]].append(ijamp)
 
-        sym_indices, sym_perms, _ = find_symmetry(
-            self.matrix_element, lambda diag: IdentifySGConfigTag(diag, self.model)
-        )
-
         self.channels = []
-        self.channel_indices = []
-        for diagram_index, (sym_index, sym_perm) in enumerate(zip(sym_indices, sym_perms)):
+        channel_indices = []
+        self.diagram_tags = []
+        for diagram_index, (sym_index, sym_perm) in enumerate(zip(self.sym_indices, self.sym_perms)):
             if sym_index == 0:
-                self.channel_indices.append(-1)
+                channel_indices.append(-1)
                 continue
 
             active_colors = diag_jamps[diagram_index] if self.color_basis else [0]
@@ -93,19 +163,23 @@ def set_channels_colors_map(self):
                 raise RuntimeError(
                     f"no valid flavor configurations found for diagram {diagram_index+1}"
                 )
+            diagram = self.diagrams[diagram_index]
             if sym_index < 0:
-                self.channels[self.channel_indices[-sym_index - 1]]["diagrams"].append(
+                chan_index = channel_indices[-sym_index - 1]
+                self.diagram_tags[chan_index].append(
+                    IdentifyJetTag(diagram, self.model),
+                )
+                self.channels[chan_index]["diagrams"].append(
                     {
                         "diagram": diagram_index,
                         "permutation": sym_perm,
                         "active_flavors": active_flavors,
                         "active_colors": active_colors,
                     }
                 )
-                self.channel_indices.append(-1)
+                channel_indices.append(-1)
                 continue
 
-            diagram = self.diagrams[diagram_index]
             vertices = []
             propagators = []
             on_shell_propagators = []
@@ -129,7 +203,9 @@ def set_channels_colors_map(self):
                         on_shell_propagators.append(prop_index)
                 vertices.append(vertex_props)
 
-            self.channel_indices.append(len(self.channels))
+            chan_index = len(self.channels)
+            self.diagram_tags.append([IdentifyJetTag(diagram, self.model)])
+            channel_indices.append(chan_index)
             self.channels.append(
                 {
                     "propagators": propagators,
@@ -211,15 +287,20 @@ def get_subprocess_info(self, proc_dir, lib_me_path):
             }
             for index, options in self.all_flavors_same_initial
         ]
-        return {
-            "incoming": self.incoming,
-            "outgoing": self.outgoing,
-            "channels": self.channels,
-            "me_path": lib_me_path,
-            "path": proc_dir,
-            "flavors": flavors,
-            "color_flows": color_flows,
-            "pdg_color_types": pdg_color_types,
-            "diagram_count": len(self.diagrams),
-            "helicities": list(self.matrix_element.get_helicity_matrix()),
-        }
+
+        return (
+            {
+                "incoming": self.incoming,
+                "outgoing": self.outgoing,
+                "channels": self.channels,
+                "me_path": lib_me_path,
+                "path": proc_dir,
+                "flavors": flavors,
+                "color_flows": color_flows,
+                "pdg_color_types": pdg_color_types,
+                "diagram_count": len(self.diagrams),
+                "helicities": list(self.matrix_element.get_helicity_matrix()),
+            },
+            self.diagram_tags,
+            self.subprocess_class,
+        )