Passing invariant masses and virtualities from phasespace to matrix element

madgraph/iolibs/template_files/madmatrix/umami.cc

@@ -320,6 +320,12 @@ extern "C"
         case UMAMI_IN_DIAGRAM_INDEX:
           diagram_in = static_cast<const int*>( input );
           break;
+        case UMAMI_IN_INVARIANT_COUNT:
+        case UMAMI_IN_INVARIANT_PIDS_AND_MASKS:
+        case UMAMI_IN_INVARIANT_MASSES:
+        case UMAMI_IN_INVARIANT_VIRTUALITIES:
+          // accept, but ignore externally supplied invariants
+          break;
         default:
           return UMAMI_ERROR_UNSUPPORTED_INPUT;
       }

madgraph/iolibs/template_files/mg7/madevent.py

@@ -1060,6 +1060,9 @@ def build_multichannel_phasespace(self) -> PhaseSpace:
                     invariant_power=self.process.run_card["phasespace"]["invariant_power"],
                     permutations=chan_permutations,
                     leptonic=self.process.leptonic,
+                    return_invariants=self.process.run_card["phasespace"][
+                        "pass_invariants_to_matrix_element"
+                    ],
                 )
                 prefix = f"subproc{self.subproc_id}.channel{channel_id}"
                 if topo_count > 1:
@@ -1114,6 +1117,9 @@ def build_flat_phasespace(self) -> PhaseSpace:
             mode=self.t_channel_mode(self.process.run_card["phasespace"]["flat_mode"]),
             cuts=self.cuts,
             leptonic=self.process.leptonic,
+            return_invariants=self.process.run_card["phasespace"][
+                "pass_invariants_to_matrix_element"
+            ],
         )
         prefix = f"subproc{self.subproc_id}.flat"
         discrete_before, discrete_after = self.build_discrete(
@@ -1377,37 +1383,62 @@ def build_integrands(
             flavor_remap.append(flav["index"])
             flavor_factors.append(len(flav["options"]))
             flavor_mirror.append(flav["mirror"])
-        if self.matrix_element:
-            matrix_element = ms.MatrixElement(
-                self.matrix_element,
-                ms.Integrand.matrix_element_inputs,
-                ms.Integrand.matrix_element_outputs,
-                True,
-            )
-        else:
-            matrix_element = ms.MatrixElement(
-                0xBADCAFE,
-                self.particle_count,
-                ms.Integrand.matrix_element_inputs,
-                ms.Integrand.matrix_element_outputs,
-                self.meta["diagram_count"],
-                True,
-            )
+        pass_invariants = self.process.run_card["phasespace"][
+            "pass_invariants_to_matrix_element"
+        ]
+        matrix_element_inputs = list(ms.Integrand.matrix_element_inputs)
+        if pass_invariants:
+            matrix_element_inputs += [
+                ms.MatrixElement.invariant_count_in,
+                ms.MatrixElement.invariant_pids_and_masks_in,
+                ms.MatrixElement.invariant_masses_in,
+                ms.MatrixElement.invariant_virtualities_in,
+            ]
+
         pdf_grid = None if self.process.leptonic else self.process.pdf_grid
         pdf_arg = None if self.process.leptonic else ms.CachedPdf()
-        cross_section = ms.DifferentialCrossSection(
-            matrix_element=matrix_element,
-            cm_energy=self.process.e_cm,
-            running_coupling=None,
-            energy_scale=ms.CachedScale(),
-            pid_options=flavors,
-            pdf1=pdf_arg,
-            pdf2=pdf_arg,
-            input_momentum_fraction=True,
-        )
+
+        def build_cross_section(invariant_count: int) -> ms.DifferentialCrossSection:
+            if self.matrix_element:
+                matrix_element = ms.MatrixElement(
+                    self.matrix_element,
+                    matrix_element_inputs,
+                    ms.Integrand.matrix_element_outputs,
+                    True,
+                    invariant_count,
+                )
+            else:
+                matrix_element = ms.MatrixElement(
+                    0xBADCAFE,
+                    self.particle_count,
+                    matrix_element_inputs,
+                    ms.Integrand.matrix_element_outputs,
+                    self.meta["diagram_count"],
+                    True,
+                    invariant_count,
+                )
+            return ms.DifferentialCrossSection(
+                matrix_element=matrix_element,
+                cm_energy=self.process.e_cm,
+                running_coupling=None,
+                energy_scale=ms.CachedScale(),
+                pid_options=flavors,
+                pdf1=pdf_arg,
+                pdf2=pdf_arg,
+                input_momentum_fraction=True,
+            )
+
+        # only build per-integrand cross sections if invariants are passed
+        shared_cross_section = None if pass_invariants else build_cross_section(0)
+
         partial_weights = self.process.run_card["generation"]["systematics"]
         integrands = []
         for channel in phasespace.channels:
+            cross_section = (
+                shared_cross_section
+                if shared_cross_section is not None
+                else build_cross_section(channel.phasespace_mapping.invariant_count())
+            )
             integrands.append(ms.Integrand(
                 channel.phasespace_mapping,
                 cross_section,

madgraph/iolibs/template_files/mg7/run_card.toml

@@ -64,6 +64,7 @@ flat_mode = %(phasespace.flat_mode)s # options: propagator, rambo, chili
 simplified_channel_count = %(phasespace.simplified_channel_count)s
 invariant_power = %(phasespace.invariant_power)s
 bw_cutoff = %(phasespace.bw_cutoff)s
+pass_invariants_to_matrix_element = %(phasespace.pass_invariants_to_matrix_element)s
 
 [multiparticles]
 $multiparticles

madgraph/various/RunCardLO_to_MG7_mapping.md

@@ -130,7 +130,7 @@ Cuts that are **not representable** in the current MG7 cut engine ([x] unless no
 `generation.freeze_max_weight_after`, `generation.max_overweight_truncation`,
 `generation.cut_efficiency_threshold`, `generation.max_cut_repetitions`,
 all of `[vegas]`, `phasespace.{mode,t_channel,flat_mode,invariant_power,
-simplified_channel_count,decays}`, all of `[madnis]`.
+simplified_channel_count,decays,pass_invariants_to_matrix_element}`, all of `[madnis]`.
 
 ---
 

madgraph/various/banner.py

@@ -6505,6 +6505,7 @@ def default_setup(self):
         self.add_toml_param('phasespace', 'simplified_channel_count', 10)
         self.add_toml_param('phasespace', 'invariant_power', 0.7)
         self.add_toml_param('phasespace', 'bw_cutoff', 15)
+        self.add_toml_param('phasespace', 'pass_invariants_to_matrix_element', False)
 
         # ----------------------------- [madnis] -----------------------
         self.add_toml_param('madnis', 'enable', False)

madspace/include/madspace/compgraphs/function_builder_mixin.inc

@@ -71,6 +71,10 @@ Value sub(Value in1, Value in2) {
     return instruction("sub", {in1, in2})[0];
 }
 
+Value neg(Value in) {
+    return instruction("neg", {in})[0];
+}
+
 Value mul(Value in1, Value in2) {
     return instruction("mul", {in1, in2})[0];
 }
@@ -389,29 +393,29 @@ std::array<Value, 2> uniform_invariant_inverse(Value s, Value s_min, Value s_max
     return {output_vector[0], output_vector[1]};
 }
 
-std::array<Value, 2> breit_wigner_invariant(Value r, Value mass, Value width, Value s_min, Value s_max) {
+std::array<Value, 3> breit_wigner_invariant(Value r, Value mass, Value width, Value s_min, Value s_max) {
     auto output_vector = instruction("breit_wigner_invariant", {r, mass, width, s_min, s_max});
-    return {output_vector[0], output_vector[1]};
+    return {output_vector[0], output_vector[1], output_vector[2]};
 }
 
 std::array<Value, 2> breit_wigner_invariant_inverse(Value s, Value mass, Value width, Value s_min, Value s_max) {
     auto output_vector = instruction("breit_wigner_invariant_inverse", {s, mass, width, s_min, s_max});
     return {output_vector[0], output_vector[1]};
 }
 
-std::array<Value, 2> stable_invariant(Value r, Value mass, Value s_min, Value s_max) {
+std::array<Value, 3> stable_invariant(Value r, Value mass, Value s_min, Value s_max) {
     auto output_vector = instruction("stable_invariant", {r, mass, s_min, s_max});
-    return {output_vector[0], output_vector[1]};
+    return {output_vector[0], output_vector[1], output_vector[2]};
 }
 
 std::array<Value, 2> stable_invariant_inverse(Value s, Value mass, Value s_min, Value s_max) {
     auto output_vector = instruction("stable_invariant_inverse", {s, mass, s_min, s_max});
     return {output_vector[0], output_vector[1]};
 }
 
-std::array<Value, 2> stable_invariant_nu(Value r, Value mass, Value nu, Value s_min, Value s_max) {
+std::array<Value, 3> stable_invariant_nu(Value r, Value mass, Value nu, Value s_min, Value s_max) {
     auto output_vector = instruction("stable_invariant_nu", {r, mass, nu, s_min, s_max});
-    return {output_vector[0], output_vector[1]};
+    return {output_vector[0], output_vector[1], output_vector[2]};
 }
 
 std::array<Value, 2> stable_invariant_nu_inverse(Value s, Value mass, Value nu, Value s_min, Value s_max) {
@@ -591,6 +595,10 @@ Value permute_momenta(Value momenta, Value permutations, Value index) {
     return instruction("permute_momenta", {momenta, permutations, index})[0];
 }
 
+Value permute_bits(Value input, Value permutations, Value index) {
+    return instruction("permute_bits", {input, permutations, index})[0];
+}
+
 Value gather(Value index, Value choices) {
     return instruction("gather", {index, choices})[0];
 }