MLM merging

madgraph/iolibs/template_files/mg7/madevent.py

@@ -275,18 +275,22 @@ def init_beam(self) -> None:
             "half_transverse_mass": ms.EnergyScale.half_transverse_mass,
             "partonic_energy": ms.EnergyScale.partonic_energy,
         }
-        if beam_args["dynamical_scale_choice"] in dynamical_scales:
+        if beam_args["dynamical_scale_choice"] == "mlm":
+            self.mlm_clustering = True
+            self.scale_kwargs = {}
+        elif beam_args["dynamical_scale_choice"] in dynamical_scales:
             dynamical_scale_type = dynamical_scales[beam_args["dynamical_scale_choice"]]
+            self.mlm_clustering = False
+            self.scale_kwargs = dict(
+                dynamical_scale_type=dynamical_scale_type,
+                ren_scale_fixed=beam_args["fixed_ren_scale"],
+                fact_scale_fixed=beam_args["fixed_fact_scale"],
+                ren_scale=beam_args["ren_scale"],
+                fact_scale1=beam_args["fact_scale1"],
+                fact_scale2=beam_args["fact_scale2"],
+            )
         else:
             raise ValueError("Unknown dynamical scale choice")
-        self.scale_kwargs = dict(
-            dynamical_scale_type=dynamical_scale_type,
-            ren_scale_fixed=beam_args["fixed_ren_scale"],
-            fact_scale_fixed=beam_args["fixed_fact_scale"],
-            ren_scale=beam_args["ren_scale"],
-            fact_scale1=beam_args["fact_scale1"],
-            fact_scale2=beam_args["fact_scale2"],
-        )
 
         pdf_set = beam_args["pdf"]
         if PDF_PATH is None:
@@ -851,9 +855,20 @@ def __init__(self, process: MadgraphProcess, meta: dict, subproc_id: int):
             else None
         )
 
-        self.scale = ms.EnergyScale(
-            particle_count=self.particle_count, **self.process.scale_kwargs
-        )
+        if self.process.mlm_clustering:
+            mc_data = self.build_multi_channel_data()
+            self.scale = ms.EnergyScale(ms.MLMClustering(
+                topologies = [topo[0] for topo in mc_data.topologies],
+                permutations = mc_data.permutations,
+                diagram_indices = mc_data.diagram_indices,
+                bw_cutoff = self.process.run_card["phasespace"]["bw_cutoff"],
+                jet_radius = self.process.run_card["beam"]["jet_radius"],
+                hadronic = not self.process.leptonic,
+            ))
+        else:
+            self.scale = ms.EnergyScale(
+                particle_count=self.particle_count, **self.process.scale_kwargs
+            )
 
         if self.process.run_card["run"]["dummy_matrix_element"]:
             self.matrix_element = None
@@ -1291,18 +1306,27 @@ def build_integrands(
             flavor_remap.append(flav["index"])
             flavor_factors.append(len(flav["options"]))
             flavor_mirror.append(flav["mirror"])
+        if self.scale.is_mlm():
+            me_inputs = [
+                ms.MatrixElement.diagram_in
+                if inp == ms.MatrixElement.random_diagram_in
+                else inp
+                for inp in ms.Integrand.matrix_element_inputs
+            ]
+        else:
+            me_inputs = ms.Integrand.matrix_element_inputs
         if self.matrix_element:
             matrix_element = ms.MatrixElement(
                 self.matrix_element,
-                ms.Integrand.matrix_element_inputs,
+                me_inputs,
                 ms.Integrand.matrix_element_outputs,
                 True,
             )
         else:
             matrix_element = ms.MatrixElement(
                 0xBADCAFE,
                 self.particle_count,
-                ms.Integrand.matrix_element_inputs,
+                me_inputs,
                 ms.Integrand.matrix_element_outputs,
                 self.meta["diagram_count"],
                 True,

madgraph/iolibs/template_files/mg7/run_card.toml

@@ -11,7 +11,7 @@ simd_vector_size = -1
 cpu_thread_pool_size = -1
 gpu_thread_pool_size = 1
 combine_thread_pool_size = -1
-output_format = "compact_npy" # options: compact_npy, lhe_npy, lhe
+output_format = "lhe" # options: compact_npy, lhe_npy, lhe
 verbosity = "pretty" # options: silent, pretty, log
 dummy_matrix_element = false
 save_gridpack = false
@@ -26,8 +26,9 @@ fixed_fact_scale = true
 ren_scale = 91.188
 fact_scale1 = 91.188
 fact_scale2 = 91.188
-# options: transverse_energy, transverse_mass, half_transverse_mass, partonic_energy
+# options: transverse_energy, transverse_mass, half_transverse_mass, partonic_energy, mlm
 dynamical_scale_choice = "half_transverse_mass"
+jet_radius = 0.4
 
 [generation]
 events = 100000

madspace/CMakeLists.txt

@@ -211,6 +211,7 @@ add_library(
     src/phasespace/channel_weight_network.cpp
     src/phasespace/pdf.cpp
     src/phasespace/matrix_element.cpp
+    src/phasespace/mlm_clustering.cpp
     src/phasespace/cross_section.cpp
     src/phasespace/scale.cpp
     include/madspace/util.hpp
@@ -266,6 +267,7 @@ add_library(
     include/madspace/phasespace/channel_weight_network.hpp
     include/madspace/phasespace/pdf.hpp
     include/madspace/phasespace/matrix_element.hpp
+    include/madspace/phasespace/mlm_clustering.hpp
     include/madspace/phasespace/cross_section.hpp
     include/madspace/phasespace/scale.hpp
 )
@@ -300,6 +302,7 @@ set(
     src/kernels/rambo.hpp
     src/kernels/cuts.hpp
     src/kernels/chili.hpp
+    src/kernels/mlm.hpp
     src/kernels/nn.hpp
     src/kernels/observables.hpp
     src/kernels/discrete.hpp

madspace/include/madspace/compgraphs/function_builder_mixin.inc

@@ -481,6 +481,16 @@ Value scale_partonic_energy(Value momenta) {
     return instruction("scale_partonic_energy", {momenta})[0];
 }
 
+std::array<Value, 5> mlm_clustering_hadronic(Value momenta, Value random, Value state_machine, Value external_masses, Value bw_masses, Value bw_widths, Value bw_cutoffs, Value jet_radius) {
+    auto output_vector = instruction("mlm_clustering_hadronic", {momenta, random, state_machine, external_masses, bw_masses, bw_widths, bw_cutoffs, jet_radius});
+    return {output_vector[0], output_vector[1], output_vector[2], output_vector[3], output_vector[4]};
+}
+
+std::array<Value, 5> mlm_clustering_leptonic(Value momenta, Value random, Value state_machine, Value external_masses, Value bw_masses, Value bw_widths, Value bw_cutoffs, Value jet_radius) {
+    auto output_vector = instruction("mlm_clustering_leptonic", {momenta, random, state_machine, external_masses, bw_masses, bw_widths, bw_cutoffs, jet_radius});
+    return {output_vector[0], output_vector[1], output_vector[2], output_vector[3], output_vector[4]};
+}
+
 std::array<Value, 2> chili_forward(Value r, Value e_cm, Value m_out, Value pt_min, Value y_max) {
     auto output_vector = instruction("chili_forward", {r, e_cm, m_out, pt_min, y_max});
     return {output_vector[0], output_vector[1]};

madspace/include/madspace/compgraphs/opcode_mixin.inc

@@ -106,52 +106,54 @@ scale_transverse_energy = 104,
 scale_transverse_mass = 105,
 scale_half_transverse_mass = 106,
 scale_partonic_energy = 107,
-chili_forward = 108,
-chili_inverse = 109,
-matrix_element = 110,
-collect_channel_weights = 111,
-interpolate_pdf = 112,
-interpolate_alpha_s = 113,
-matmul = 114,
-relu = 115,
-leaky_relu = 116,
-elu = 117,
-gelu = 118,
-sigmoid = 119,
-softplus = 120,
-rqs_reshape = 121,
-rqs_find_bin = 122,
-rqs_forward = 123,
-rqs_inverse = 124,
-softmax = 125,
-softmax_prior = 126,
-sample_discrete = 127,
-sample_discrete_inverse = 128,
-sample_discrete_probs = 129,
-sample_discrete_probs_inverse = 130,
-discrete_histogram = 131,
-permute_momenta = 132,
-gather = 133,
-gather_int = 134,
-gather_vector = 135,
-select_int = 136,
-select = 137,
-select_vector = 138,
-argsort = 139,
-quantile = 140,
-one_hot = 141,
-madnis_abs_weight = 142,
-madnis_softclip = 143,
-madnis_variance = 144,
-madnis_single_channel_variance = 145,
-madnis_multi_channel_variance = 146,
-nonzero = 147,
-batch_gather = 148,
-batch_scatter = 149,
-random = 150,
-random_int = 151,
-unweight = 152,
-vegas_forward = 153,
-vegas_inverse = 154,
-vegas_histogram = 155,
-histogram = 156
+mlm_clustering_hadronic = 108,
+mlm_clustering_leptonic = 109,
+chili_forward = 110,
+chili_inverse = 111,
+matrix_element = 112,
+collect_channel_weights = 113,
+interpolate_pdf = 114,
+interpolate_alpha_s = 115,
+matmul = 116,
+relu = 117,
+leaky_relu = 118,
+elu = 119,
+gelu = 120,
+sigmoid = 121,
+softplus = 122,
+rqs_reshape = 123,
+rqs_find_bin = 124,
+rqs_forward = 125,
+rqs_inverse = 126,
+softmax = 127,
+softmax_prior = 128,
+sample_discrete = 129,
+sample_discrete_inverse = 130,
+sample_discrete_probs = 131,
+sample_discrete_probs_inverse = 132,
+discrete_histogram = 133,
+permute_momenta = 134,
+gather = 135,
+gather_int = 136,
+gather_vector = 137,
+select_int = 138,
+select = 139,
+select_vector = 140,
+argsort = 141,
+quantile = 142,
+one_hot = 143,
+madnis_abs_weight = 144,
+madnis_softclip = 145,
+madnis_variance = 146,
+madnis_single_channel_variance = 147,
+madnis_multi_channel_variance = 148,
+nonzero = 149,
+batch_gather = 150,
+batch_scatter = 151,
+random = 152,
+random_int = 153,
+unweight = 154,
+vegas_forward = 155,
+vegas_inverse = 156,
+vegas_histogram = 157,
+histogram = 158

madspace/include/madspace/compgraphs/type.hpp

@@ -118,6 +118,9 @@ const Type batch_four_vec{DataType::dt_float, batch_size, {4}};
 inline Type batch_float_array(int count) {
     return {DataType::dt_float, batch_size, {count}};
 }
+inline Type batch_int_array(int count) {
+    return {DataType::dt_int, batch_size, {count}};
+}
 inline Type batch_four_vec_array(int count) {
     return {DataType::dt_float, batch_size, {count, 4}};
 }

madspace/include/madspace/driver/channel_generator.hpp

@@ -104,6 +104,7 @@ class ChannelEventGenerator {
         int color_index, helicity_index, diagram_index, flavor_index;
         int ren_scale, alpha_qcd;
         int x1, fact_scale1, x2, fact_scale2, partial_weight_product;
+        int cluster_scales;
         int random, rest;
     };
 

madspace/include/madspace/driver/io.hpp

@@ -115,7 +115,7 @@ class ParticleRecord {
     // clustering data
     UnalignedRef<double> cluster_scale() { return &_data[_offsets[0] + 0]; }
 
-    void from_lhe_particle(const LHEParticle& particle) {
+    void from_lhe_particle(const LHEParticle& particle, bool has_clustering = false) {
         lhe_pdg_id() = particle.pdg_id;
         lhe_status_code() = particle.status_code;
         lhe_mother1() = particle.mother1;
@@ -129,6 +129,9 @@ class ParticleRecord {
         lhe_mass() = particle.mass;
         lhe_lifetime() = particle.lifetime;
         lhe_spin() = particle.spin;
+        if (has_clustering) {
+            cluster_scale() = particle.cluster_scale;
+        }
     }
 
 private:

madspace/include/madspace/driver/lhe_output.hpp

@@ -50,6 +50,7 @@ struct LHEParticle {
     double px, py, pz, energy, mass;
     double lifetime;
     double spin;
+    double cluster_scale;
 };
 
 struct LHEEvent {
@@ -60,6 +61,7 @@ struct LHEEvent {
     double alpha_qed;
     double alpha_qcd;
     std::vector<LHEParticle> particles;
+    bool has_cluster_scales;
 
     void format_to(std::string& buffer) const;
 };

madspace/include/madspace/phasespace.hpp

@@ -16,6 +16,7 @@
 #include "phasespace/luminosity.hpp"
 #include "phasespace/madnis.hpp"
 #include "phasespace/matrix_element.hpp"
+#include "phasespace/mlm_clustering.hpp"
 #include "phasespace/mlp.hpp"
 #include "phasespace/multichannel.hpp"
 #include "phasespace/observable.hpp"

madspace/include/madspace/phasespace/mlm_clustering.hpp

@@ -0,0 +1,33 @@
+#pragma once
+
+#include "madspace/phasespace/base.hpp"
+#include "madspace/phasespace/topology.hpp"
+
+namespace madspace {
+
+class MLMClustering : public FunctionGenerator {
+public:
+    MLMClustering(
+        std::vector<Topology> topologies,
+        nested_vector3<std::size_t> permutations,
+        nested_vector2<std::size_t> diagram_indices,
+        double bw_cutoff = 15,
+        double jet_radius = 0.4,
+        bool hadronic = true
+    );
+
+private:
+    NamedVector<Value> build_function_impl(
+        FunctionBuilder& fb, const NamedVector<Value>& args
+    ) const override;
+
+    std::vector<me_int_t> _cluster_state_machine;
+    std::vector<double> _external_masses;
+    std::vector<double> _bw_masses;
+    std::vector<double> _bw_widths;
+    double _bw_cutoff;
+    double _jet_radius;
+    bool _hadronic;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/scale.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "madspace/phasespace/base.hpp"
+#include "madspace/phasespace/mlm_clustering.hpp"
 
 namespace madspace {
 
@@ -36,6 +37,9 @@ class EnergyScale : public FunctionGenerator {
         double fact_scale1,
         double fact_scale2
     );
+    EnergyScale(const MLMClustering& clustering);
+
+    bool is_mlm() const { return _clustering.has_value(); }
 
 private:
     NamedVector<Value> build_function_impl(
@@ -48,6 +52,7 @@ class EnergyScale : public FunctionGenerator {
     double _ren_scale;
     double _fact_scale1;
     double _fact_scale2;
+    std::optional<MLMClustering> _clustering;
 };
 
 } // namespace madspace