diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 57c6f73c9..e52957daa 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/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,10 +1306,19 @@ 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, ) @@ -1302,7 +1326,7 @@ def build_integrands( 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, diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 83af64539..47418315b 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/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 diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 203c36d90..df2e82c48 100644 --- a/madspace/CMakeLists.txt +++ b/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 diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc index 360d8db14..99845b03e 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/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 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 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 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]}; diff --git a/madspace/include/madspace/compgraphs/opcode_mixin.inc b/madspace/include/madspace/compgraphs/opcode_mixin.inc index ba664afd0..d1034b835 100644 --- a/madspace/include/madspace/compgraphs/opcode_mixin.inc +++ b/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 diff --git a/madspace/include/madspace/compgraphs/type.hpp b/madspace/include/madspace/compgraphs/type.hpp index f0d8ada0b..787e20b21 100644 --- a/madspace/include/madspace/compgraphs/type.hpp +++ b/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}}; } diff --git a/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 957037d15..dae774f2d 100644 --- a/madspace/include/madspace/driver/channel_generator.hpp +++ b/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; }; diff --git a/madspace/include/madspace/driver/io.hpp b/madspace/include/madspace/driver/io.hpp index 783de70da..1ab7f211d 100644 --- a/madspace/include/madspace/driver/io.hpp +++ b/madspace/include/madspace/driver/io.hpp @@ -115,7 +115,7 @@ class ParticleRecord { // clustering data UnalignedRef 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: diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index 5451d6f25..f29c612f7 100644 --- a/madspace/include/madspace/driver/lhe_output.hpp +++ b/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 particles; + bool has_cluster_scales; void format_to(std::string& buffer) const; }; diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index ca6cb3cf8..ef927428a 100644 --- a/madspace/include/madspace/phasespace.hpp +++ b/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" diff --git a/madspace/include/madspace/phasespace/mlm_clustering.hpp b/madspace/include/madspace/phasespace/mlm_clustering.hpp new file mode 100644 index 000000000..aa7ce7163 --- /dev/null +++ b/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 topologies, + nested_vector3 permutations, + nested_vector2 diagram_indices, + double bw_cutoff = 15, + double jet_radius = 0.4, + bool hadronic = true + ); + +private: + NamedVector build_function_impl( + FunctionBuilder& fb, const NamedVector& args + ) const override; + + std::vector _cluster_state_machine; + std::vector _external_masses; + std::vector _bw_masses; + std::vector _bw_widths; + double _bw_cutoff; + double _jet_radius; + bool _hadronic; +}; + +} // namespace madspace diff --git a/madspace/include/madspace/phasespace/scale.hpp b/madspace/include/madspace/phasespace/scale.hpp index b53235821..f95cfbed3 100644 --- a/madspace/include/madspace/phasespace/scale.hpp +++ b/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 build_function_impl( @@ -48,6 +52,7 @@ class EnergyScale : public FunctionGenerator { double _ren_scale; double _fact_scale1; double _fact_scale2; + std::optional _clustering; }; } // namespace madspace diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index cc8055381..ab9d85aa6 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -2113,6 +2113,94 @@ scale_partonic_energy: type: [float] desc: +mlm_clustering_hadronic: + inputs: + - name: momenta + type: [float, n, 4] + desc: + - name: random + type: [float] + desc: + - name: state_machine + type: [int, single, m] + desc: + - name: external_masses + type: [float, n] + desc: + - name: bw_masses + type: [float, single, k] + desc: + - name: bw_widths + type: [float, single, k] + desc: + - name: bw_cutoffs + type: [float, single] + desc: + - name: jet_radius + type: [float, single] + desc: + outputs: + - name: ren_scale + type: [float] + desc: + - name: fact_scale1 + type: [float] + desc: + - name: fact_scale2 + type: [float] + desc: + - name: outgoing_scales + type: [float, n-2] + desc: + - name: diagram_index + type: [int] + desc: + vectorized: False + +mlm_clustering_leptonic: + inputs: + - name: momenta + type: [float, n, 4] + desc: + - name: random + type: [float] + desc: + - name: state_machine + type: [int, single, m] + desc: + - name: external_masses + type: [float, n] + desc: + - name: bw_masses + type: [float, single, k] + desc: + - name: bw_widths + type: [float, single, k] + desc: + - name: bw_cutoffs + type: [float, single] + desc: + - name: jet_radius + type: [float, single] + desc: + outputs: + - name: ren_scale + type: [float] + desc: + - name: fact_scale1 + type: [float] + desc: + - name: fact_scale2 + type: [float] + desc: + - name: outgoing_scales + type: [float, n-2] + desc: + - name: diagram_index + type: [int] + desc: + vectorized: False + --- title: Chili diff --git a/madspace/src/compgraphs/instruction_set_mixin.inc b/madspace/src/compgraphs/instruction_set_mixin.inc index a2dbe6cd8..600850d96 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -123,53 +123,55 @@ InstructionOwner instructions[] { mi("scale_transverse_mass", 105, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), mi("scale_half_transverse_mass", 106, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), mi("scale_partonic_energy", 107, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("chili_forward", 108, true, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("chili_inverse", 109, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}}), - InstructionOwner(new MatrixElementInstruction(110, true)), - mi("collect_channel_weights", 111, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}, {DataType::dt_int, true, {"c"}, true}}, {{DataType::dt_float, false, {"c"}, false}}), - mi("interpolate_pdf", 112, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n"}, false}, {DataType::dt_float, true, {"a"}, false}, {DataType::dt_float, true, {"b"}, false}, {DataType::dt_float, true, {16, "c", "d"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("interpolate_alpha_s", 113, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {"b+1"}, false}, {DataType::dt_float, true, {4, "b"}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("matmul", 114, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"m", "n"}, false}, {DataType::dt_float, true, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("relu", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("leaky_relu", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("elu", 117, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("gelu", 118, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("sigmoid", 119, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softplus", 120, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - InstructionOwner(new RqsReshapeInstruction(121, true)), - mi("rqs_find_bin", 122, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b+1"}, false}}, {{DataType::dt_float, false, {"n", 6}, false}}), - mi("rqs_forward", 123, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), - mi("rqs_inverse", 124, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), - mi("softmax", 125, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softmax_prior", 126, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("sample_discrete", 127, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_inverse", 128, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_probs", 129, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_probs_inverse", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("discrete_histogram", 131, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, true, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}}), - mi("permute_momenta", 132, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, true, {"m", "n"}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), - mi("gather", 133, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("gather_int", 134, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), - mi("gather_vector", 135, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n", "m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("select_int", 136, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), - mi("select", 137, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("select_vector", 138, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), - mi("argsort", 139, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), - mi("quantile", 140, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}), - mi("one_hot", 141, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("madnis_abs_weight", 142, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_softclip", 143, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_variance", 144, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_single_channel_variance", 145, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_multi_channel_variance", 146, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), - InstructionOwner(new NonzeroInstruction(147, true)), - InstructionOwner(new BatchGatherInstruction(148, true)), - InstructionOwner(new BatchScatterInstruction(149, true)), - InstructionOwner(new RandomInstruction(150, true)), - InstructionOwner(new RandomIntInstruction(151, true)), - InstructionOwner(new UnweightInstruction(152, true)), - mi("vegas_forward", 153, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), - mi("vegas_inverse", 154, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), - mi("vegas_histogram", 155, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"n", "b"}, false}, {DataType::dt_int, true, {"n", "b"}, false}}), - mi("histogram", 156, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"b+2"}, false}, {DataType::dt_float, true, {"b+2"}, false}}), + mi("mlm_clustering_hadronic", 108, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"m"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"k"}, false}, {DataType::dt_float, true, {"k"}, false}, {DataType::dt_float, true, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n-2"}, false}, {DataType::dt_int, false, {}, false}}), + mi("mlm_clustering_leptonic", 109, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"m"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"k"}, false}, {DataType::dt_float, true, {"k"}, false}, {DataType::dt_float, true, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n-2"}, false}, {DataType::dt_int, false, {}, false}}), + mi("chili_forward", 110, true, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("chili_inverse", 111, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}}), + InstructionOwner(new MatrixElementInstruction(112, true)), + mi("collect_channel_weights", 113, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}, {DataType::dt_int, true, {"c"}, true}}, {{DataType::dt_float, false, {"c"}, false}}), + mi("interpolate_pdf", 114, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n"}, false}, {DataType::dt_float, true, {"a"}, false}, {DataType::dt_float, true, {"b"}, false}, {DataType::dt_float, true, {16, "c", "d"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("interpolate_alpha_s", 115, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {"b+1"}, false}, {DataType::dt_float, true, {4, "b"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("matmul", 116, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"m", "n"}, false}, {DataType::dt_float, true, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("relu", 117, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("leaky_relu", 118, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("elu", 119, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("gelu", 120, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("sigmoid", 121, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softplus", 122, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + InstructionOwner(new RqsReshapeInstruction(123, true)), + mi("rqs_find_bin", 124, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b+1"}, false}}, {{DataType::dt_float, false, {"n", 6}, false}}), + mi("rqs_forward", 125, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("rqs_inverse", 126, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("softmax", 127, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softmax_prior", 128, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("sample_discrete", 129, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_inverse", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_probs", 131, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_probs_inverse", 132, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("discrete_histogram", 133, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, true, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}}), + mi("permute_momenta", 134, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, true, {"m", "n"}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), + mi("gather", 135, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("gather_int", 136, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), + mi("gather_vector", 137, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n", "m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_int", 138, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), + mi("select", 139, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_vector", 140, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), + mi("argsort", 141, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), + mi("quantile", 142, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}), + mi("one_hot", 143, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("madnis_abs_weight", 144, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_softclip", 145, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_variance", 146, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_single_channel_variance", 147, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_multi_channel_variance", 148, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), + InstructionOwner(new NonzeroInstruction(149, true)), + InstructionOwner(new BatchGatherInstruction(150, true)), + InstructionOwner(new BatchScatterInstruction(151, true)), + InstructionOwner(new RandomInstruction(152, true)), + InstructionOwner(new RandomIntInstruction(153, true)), + InstructionOwner(new UnweightInstruction(154, true)), + mi("vegas_forward", 155, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("vegas_inverse", 156, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("vegas_histogram", 157, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"n", "b"}, false}, {DataType::dt_int, true, {"n", "b"}, false}}), + mi("histogram", 158, true, {{DataType::dt_float, false, {std::monostate{}}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"b"}, true}}, {{DataType::dt_float, true, {"b+2"}, false}, {DataType::dt_float, true, {"b+2"}, false}}), }; diff --git a/madspace/src/cpu/runtime_backward_mixin.inc b/madspace/src/cpu/runtime_backward_mixin.inc index 8a75bf498..cff5828b8 100644 --- a/madspace/src/cpu/runtime_backward_mixin.inc +++ b/madspace/src/cpu/runtime_backward_mixin.inc @@ -52,66 +52,66 @@ case 24: case 25: backward_batch_foreach, backward_kernel_square, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 116: backward_op_matmul(instr, locals, local_grads, device); break; -case 115: +case 117: backward_batch_foreach, backward_kernel_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 116: +case 118: backward_batch_foreach, backward_kernel_leaky_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 117: +case 119: backward_batch_foreach, backward_kernel_elu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 118: +case 120: backward_batch_foreach, backward_kernel_gelu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 119: +case 121: backward_batch_foreach, backward_kernel_sigmoid, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 120: +case 122: backward_batch_foreach, backward_kernel_softplus, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 121: +case 123: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 122: +case 124: backward_batch_foreach, backward_kernel_rqs_find_bin, 5, 4, 2, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 123: +case 125: backward_batch_foreach, backward_kernel_rqs_forward, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 124: +case 126: backward_batch_foreach, backward_kernel_rqs_inverse, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 125: +case 127: backward_batch_foreach, backward_kernel_softmax, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 126: +case 128: backward_batch_foreach, backward_kernel_softmax_prior, 2, 2, 1, DeviceType>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 130: +case 132: backward_batch_foreach, backward_kernel_sample_discrete_probs_inverse, 4, 2, 1, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 133: +case 135: backward_batch_foreach, backward_kernel_gather, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 137: +case 139: backward_batch_foreach, backward_kernel_select, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 142: +case 144: backward_batch_foreach, backward_kernel_madnis_abs_weight, 3, 2, 1, DeviceType>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 143: +case 145: backward_batch_foreach, backward_kernel_madnis_softclip, 5, 4, 1, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 144: +case 146: backward_batch_foreach, backward_kernel_madnis_variance, 5, 4, 1, DeviceType>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 145: +case 147: backward_batch_foreach, backward_kernel_madnis_single_channel_variance, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 146: +case 148: backward_batch_foreach, backward_kernel_madnis_multi_channel_variance, 3, 2, 1, DeviceType>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; diff --git a/madspace/src/cpu/runtime_mixin.inc b/madspace/src/cpu/runtime_mixin.inc index 64a3974be..a77ba4ce0 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -326,149 +326,155 @@ case 107: batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 108: - batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_mlm_clustering_hadronic, 8, 5, 1, DeviceType>, 8, 5>(instr, locals, device); break; case 109: - batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_mlm_clustering_leptonic, 8, 5, 1, DeviceType>, 8, 5>(instr, locals, device); break; case 110: - op_matrix_element(instr, locals, device); + batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 111: - batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 112: - batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 113: - batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 114: - op_matmul(instr, locals, device); + batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 115: - batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 116: - batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 117: - batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 118: - batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 119: - batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 120: - batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 121: - op_rqs_reshape(instr, locals, device); + batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 122: - batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 123: - batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 124: - batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); break; case 125: - batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 126: - batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 127: - batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 128: - batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 129: - batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 130: - batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 131: - op_discrete_histogram(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 132: - batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 133: - batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 134: - batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 135: - batch_foreach, kernel_gather_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 136: - batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 137: - batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 138: - batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 139: - batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 140: - op_quantile(instr, locals, device); + batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 142: - batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + op_quantile(instr, locals, device); break; case 143: - batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 144: - batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 145: - batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 146: - batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 147: - op_nonzero(instr, locals, device); + batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 148: - op_batch_gather(instr, locals, device); + batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 149: - op_batch_scatter(instr, locals, device); + op_nonzero(instr, locals, device); break; case 150: - op_random(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 151: - op_random_int(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 152: - op_unweight(instr, locals, device); + op_random(instr, locals, device); break; case 153: - batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_random_int(instr, locals, device); break; case 154: - batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_unweight(instr, locals, device); break; case 155: - op_vegas_histogram(instr, locals, device); + batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 156: + batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 157: + op_vegas_histogram(instr, locals, device); + break; +case 158: op_histogram(instr, locals, device); break; diff --git a/madspace/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 823fb5e86..d83797b9b 100644 --- a/madspace/src/driver/channel_generator.cpp +++ b/madspace/src/driver/channel_generator.cpp @@ -271,6 +271,11 @@ void ChannelEventGenerator::init_field_indices() { } else { _field_indices.partial_weight_product = -1; } + if (index_map.contains("cluster_scales")) { + _field_indices.cluster_scales = index_map.at("cluster_scales"); + } else { + _field_indices.cluster_scales = -1; + } _field_indices.random = index_map.at("random"); _field_indices.rest = _field_indices.random + 1; } @@ -644,6 +649,18 @@ void ChannelEventGenerator::write_events( } } + if (_field_indices.cluster_scales != -1) { + auto cluster_view = + unweighted_events.at(_field_indices.cluster_scales).view(); + for (std::size_t i = 0; i < w_view.size(); ++i) { + auto scales = cluster_view[i]; + for (std::size_t j = 0; j < _particle_count; ++j) { + auto particle = event_buffer.particle(i, j); + particle.cluster_scale() = j >= 2 ? scales[j - 2] : 0.; + } + } + } + _event_file.write(event_buffer); _weight_file.write(weight_buffer); _status.count_unweighted += diff --git a/madspace/src/driver/event_generator.cpp b/madspace/src/driver/event_generator.cpp index af8e3a59b..3a078d3ea 100644 --- a/madspace/src/driver/event_generator.cpp +++ b/madspace/src/driver/event_generator.cpp @@ -399,10 +399,14 @@ void EventGenerator::combine_to_lhe_npy( buffer_out.event(i).from_lhe_event(lhe_event); std::size_t j = 0; for (; j < lhe_event.particles.size(); ++j) { - buffer_out.particle(i, j).from_lhe_particle(lhe_event.particles[j]); + buffer_out.particle(i, j).from_lhe_particle( + lhe_event.particles[j], lhe_event.has_cluster_scales + ); } for (; j < lhe_completer.max_particle_count(); ++j) { - buffer_out.particle(i, j).from_lhe_particle(LHEParticle{}); + buffer_out.particle(i, j).from_lhe_particle( + LHEParticle{}, lhe_event.has_cluster_scales + ); } } event_file.write(buffer_out); @@ -606,6 +610,8 @@ void EventGenerator::read_and_combine( bool has_beam2 = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_beam2; bool has_partial = _channels.at(0)->event_layout_extra_flags() & EventRecord::f_partial_weights; + bool has_clustering = + _channels.at(0)->particle_layout_extra_flags() & ParticleRecord::f_clustering; std::random_device rand_device; std::mt19937 rand_gen(rand_device()); @@ -671,6 +677,9 @@ void EventGenerator::read_and_combine( particle_out.px() = particle_in.px(); particle_out.py() = particle_in.py(); particle_out.pz() = particle_in.pz(); + if (has_clustering) { + particle_out.cluster_scale() = particle_in.cluster_scale(); + } } for (; i < buffer.particle_count(); ++i) { auto particle_out = buffer.particle(event_index, i); @@ -711,6 +720,15 @@ void EventGenerator::fill_lhe_event( } ); } + bool has_clustering = + _channels.at(0)->particle_layout_extra_flags() & ParticleRecord::f_clustering; + lhe_event.has_cluster_scales = has_clustering; + if (has_clustering) { + for (std::size_t i = 0; i < lhe_event.particles.size(); ++i) { + lhe_event.particles.at(i).cluster_scale = + buffer.particle(event_index, i).cluster_scale(); + } + } lhe_completer.complete_event_data( lhe_event, event_in.subprocess_index(), diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 32b8f1e28..03ada99aa 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -53,6 +53,18 @@ void LHEEvent::format_to(std::string& buffer) const { particle.spin ); } + if (has_cluster_scales) { + std::format_to(insert_iter, "\n"); + } buffer += "\n"; } @@ -165,6 +177,8 @@ LHECompleter::LHECompleter( std::span(permutation.begin() + 2, permutation.end()))) { e_min.at(index) = mass; momentum_masks.at(index) = 1 << perm_index; + // TODO: is this correct or is the permutation the wrong way + // around? for (std::size_t i = 0; std::size_t color_index : colors) { prop_colors.at(colors.size() * index + i) = args.color_flows.at(matrix_flavor_index) diff --git a/madspace/src/gpu/runtime_backward_mixin.inc b/madspace/src/gpu/runtime_backward_mixin.inc index 452946c95..91346fd0f 100644 --- a/madspace/src/gpu/runtime_backward_mixin.inc +++ b/madspace/src/gpu/runtime_backward_mixin.inc @@ -52,66 +52,66 @@ case 24: case 25: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 116: backward_op_matmul(instr, locals, local_grads, device); break; -case 115: +case 117: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 116: +case 118: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 117: +case 119: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 118: +case 120: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 119: +case 121: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 120: +case 122: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 121: +case 123: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 122: +case 124: backward_batch_foreach, 5, 4, 2>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 123: +case 125: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 124: +case 126: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 125: +case 127: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 126: +case 128: backward_batch_foreach, 2, 2, 1>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 130: +case 132: backward_batch_foreach, 4, 2, 1>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 133: +case 135: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 137: +case 139: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 142: +case 144: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 143: +case 145: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 144: +case 146: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 145: +case 147: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 146: +case 148: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; diff --git a/madspace/src/gpu/runtime_mixin.inc b/madspace/src/gpu/runtime_mixin.inc index 316c746a3..a72a42749 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -326,149 +326,155 @@ case 107: batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 108: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 8, 5, 1>, 8, 5>(instr, locals, device); break; case 109: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 8, 5, 1>, 8, 5>(instr, locals, device); break; case 110: - op_matrix_element(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 111: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 112: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 113: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 114: - op_matmul(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 115: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 116: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 117: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 118: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 119: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 120: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 121: - op_rqs_reshape(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 122: - batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 123: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 124: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); break; case 125: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 126: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 127: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 128: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 129: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 130: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 131: - op_discrete_histogram(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 132: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 133: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 134: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 135: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 136: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 137: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 138: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 139: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 140: - op_quantile(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 142: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + op_quantile(instr, locals, device); break; case 143: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 144: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 145: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 146: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 147: - op_nonzero(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 148: - op_batch_gather(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 149: - op_batch_scatter(instr, locals, device); + op_nonzero(instr, locals, device); break; case 150: - op_random(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 151: - op_random_int(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 152: - op_unweight(instr, locals, device); + op_random(instr, locals, device); break; case 153: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_random_int(instr, locals, device); break; case 154: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_unweight(instr, locals, device); break; case 155: - op_vegas_histogram(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 156: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 157: + op_vegas_histogram(instr, locals, device); + break; +case 158: op_histogram(instr, locals, device); break; diff --git a/madspace/src/kernels/kernels.hpp b/madspace/src/kernels/kernels.hpp index 7c2d828ef..307397aef 100644 --- a/madspace/src/kernels/kernels.hpp +++ b/madspace/src/kernels/kernels.hpp @@ -7,6 +7,7 @@ #include "kinematics.hpp" #include "lup_det.hpp" #include "math.hpp" +#include "mlm.hpp" #include "multichannel.hpp" #include "nn.hpp" #include "observables.hpp" diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index 67fe1cb67..187165e51 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -12,7 +12,7 @@ constexpr double INV_GEV2_TO_PB = 0.38937937217186e9; template struct FourMom { KERNELSPEC FVal& operator[](int i) { return p[i]; } - KERNELSPEC FVal& operator[](int i) const { return p[i]; } + KERNELSPEC const FVal& operator[](int i) const { return p[i]; } FVal p[4]; }; @@ -192,6 +192,16 @@ KERNELSPEC FVal esquare(FourMom p) { return p[1] * p[1] + p[2] * p[2] + p[3] * p[3]; } +template +KERNELSPEC FVal ldot(FourMom p1, FourMom p2) { + return p1[0] * p2[0] - p1[1] * p2[1] - p1[2] * p2[2] - p1[3] * p2[3]; +} + +template +KERNELSPEC FVal edot(FourMom p1, FourMom p2) { + return p1[1] * p2[1] + p1[2] * p2[2] + p1[3] * p2[3]; +} + template KERNELSPEC FourMom rotate(FourMom p, FourMom q) { auto qt2 = q[1] * q[1] + q[2] * q[2]; diff --git a/madspace/src/kernels/mlm.hpp b/madspace/src/kernels/mlm.hpp new file mode 100644 index 000000000..1e30b63e1 --- /dev/null +++ b/madspace/src/kernels/mlm.hpp @@ -0,0 +1,418 @@ +#pragma once + +#include "definitions.hpp" +#include "kinematics.hpp" + +namespace madspace { +namespace kernels { + +constexpr int N_EXT_MAX = 12; +constexpr double ONE_PLUS_TINY = 1.000001; + +// mT^2 = E^2 - pz^2 (hadronic) or E^2 (lepton collider). +// based on djb_clus from Template/NLO/SubProcesses/cluster.f +template +KERNELSPEC FVal djb_clus(const FourMom& p, bool hadronic) { + auto r = hadronic ? (p[0] - p[3]) * (p[0] + p[3]) : p[0] * p[0]; + return max(r, 0.0); +} + +// kt/Durham clustering measure for two partons. +// mass1/mass2 are the tracked clustering masses kept in a separate array and updated by +// update_momenta - NOT the Lorentz-invariant masses computed from the 4-momentum. +// jet_radius is the jet-radius parameter (Fortran common /to_dj/D). +// based on dj_clus from Template/NLO/SubProcesses/cluster.f +template +KERNELSPEC FVal dj_clus( + const FourMom& p1, + const FourMom& p2, + FVal mass1, + FVal mass2, + bool hadronic, + FVal jet_radius +) { + if (!hadronic) { + // Durham e+e- measure: 2*min(E1^2, E2^2)*(1 - cos_theta) + auto p1a = sqrt(esquare(p1)); + auto p2a = sqrt(esquare(p2)); + if (p1a * p2a == 0.0) { + return 0.0; + } + auto costh = edot(p1, p2) / (p1a * p2a); + return 2.0 * min(p1[0] * p1[0], p2[0] * p2[0]) * max(1.0 - costh, 0.0); + } + + // hadronic: massless+massive pair clusters to the lighter parton's mT^2 + bool massive1 = (mass1 > 0.0); + bool massive2 = (mass2 > 0.0); + if (!massive1 && massive2) { + return djb_clus(p1, true) * ONE_PLUS_TINY; + } + if (massive1 && !massive2) { + return djb_clus(p2, true) * ONE_PLUS_TINY; + } + + // both massless or both massive: generalised kt measure in (eta, phi) + auto pt1_sq = p1[1] * p1[1] + p1[2] * p1[2]; + auto pt2_sq = p2[1] * p2[1] + p2[2] * p2[2]; + if (pt1_sq == 0.0 || pt2_sq == 0.0) { + return 0.0; + } + auto p1a = sqrt(pt1_sq + p1[3] * p1[3]); + auto p2a = sqrt(pt2_sq + p2[3] * p2[3]); + auto eta1 = 0.5 * log((p1a + p1[3]) / (p1a - p1[3])); + auto eta2 = 0.5 * log((p2a + p2[3]) / (p2a - p2[3])); + auto m_max_sq = max(mass1 * mass1, mass2 * mass2); + auto dphi_cos = (p1[1] * p2[1] + p1[2] * p2[2]) / sqrt(pt1_sq * pt2_sq); + return max( + m_max_sq + + min(pt1_sq, pt2_sq) * 2.0 * (cosh(eta1 - eta2) - dphi_cos) / + (jet_radius * jet_radius), + 0.0 + ); +} + +// Clustering scale for the pair (momentum1=pi, momentum2=pj). +// mass1/mass2 are the tracked clustering masses for momentum1/momentum2. +// +// Parameters replacing Fortran globals / derived quantities: +// is_initial : momentum2 is a beam particle (Fortran: j<=2 in cluster_one_step) +// hadronic : hadronic collider (Fortran: lpp[] from run.inc) +// jet_radius : jet-radius parameter (Fortran: common /to_dj/D) +// massive_in/out1/out2 replace the Fortran get_clustering_type cl[0:2] bit-array +// resonant : replaces the iBWlist lookup in Fortran cluster_scale +// mass_in/width_in: carried for future use; not consumed in the scale formula +// +// based on cluster_scale from Template/NLO/SubProcesses/cluster.f +template +KERNELSPEC FVal compute_scale( + const FourMom& momentum1, + const FourMom& momentum2, + const FourMom& momentum_sum, + FVal mass1, + FVal mass2, + bool resonant, + bool is_initial, + bool massive_in, + bool massive_out1, + bool massive_out2, + bool hadronic, + FVal jet_radius +) { + if (is_initial) { + // scale = mT of the final-state parton + // small penalty when it goes against the beam + auto scale = sqrt(djb_clus(momentum2, hadronic)); + if ((momentum1[3] < 0.0) != (momentum2[3] < 0.0)) { + scale = scale * ONE_PLUS_TINY; + } + return scale; + } + if (resonant) { + return sqrt(max(lsquare(momentum_sum), 0.0)); + } + if (!massive_in && massive_out1 && !massive_out2) { + return sqrt(fabs(ldot(momentum2, momentum_sum))) / 2.0; + } + if (!massive_in && !massive_out1 && massive_out2) { + return sqrt(fabs(ldot(momentum1, momentum_sum))) / 2.0; + } + if (massive_in && !massive_out1 && !massive_out2) { + return sqrt(max(lsquare(momentum_sum), 0.0)); + } + return sqrt(dj_clus(momentum1, momentum2, mass1, mass2, hadronic, jet_radius)); +} + +// Update momenta and the separately-tracked masses after one clustering step. +// Based on update_momenta from Template/NLO/SubProcesses/cluster.f +template +KERNELSPEC void update_momenta( + int n_part, + FourMom* momenta, + FVal* masses, + int& alive, + int i_remove, + int i_keep, + bool resonant +) { + alive &= ~(1 << i_remove); + + if (i_keep < 2) { // initial-state clustering + int j_other = 1 - i_keep; + for (int k = 0; k < 4; ++k) { + momenta[i_keep][k] -= momenta[i_remove][k]; + } + + masses[i_keep] = (masses[i_keep] > 0.0) != (masses[i_remove] > 0.0) + ? max(masses[i_keep], masses[i_remove]) + : 0.0; + + FourMom com_boost_vector = { + momenta[i_keep][0] + momenta[j_other][0], + -(momenta[i_keep][1] + momenta[j_other][1]), + -(momenta[i_keep][2] + momenta[j_other][2]), + -(momenta[i_keep][3] + momenta[j_other][3]), + }; + if (lsquare(com_boost_vector) > 100.0) { + // boost j_keep to COM frame to define the rotation axis, then apply to all + // alive particles + auto jkeep_cm = boost(momenta[i_keep], com_boost_vector, 1.0); + for (int j = 0; j < n_part; ++j) { + if (alive & (1 << j)) { + momenta[j] = rotate_inverse( + boost(momenta[j], com_boost_vector, 1.0), jkeep_cm + ); + } + } + } + } else { // final-state clustering + for (int k = 0; k < 4; ++k) { + momenta[i_keep][k] += momenta[i_remove][k]; + } + + if (resonant) { + masses[i_keep] = sqrt(max(lsquare(momenta[i_keep]), 0.0)); + } else { + masses[i_keep] = max(masses[i_keep], masses[i_remove]); + } + } +} + +template +KERNELSPEC void mlm_clustering( + FIn momenta, + FIn random, + IIn state_machine, + FIn external_masses, + FIn bw_masses, + FIn bw_widths, + FIn bw_cutoff, + FIn jet_radius, + FOut ren_scale, + FOut fact_scale1, + FOut fact_scale2, + FOut outgoing_scales, + IOut diagram_index, + bool hadronic +) { + // we do not support SIMD for now, so we can assume simple types + static_assert(std::is_same_v, int>); + static_assert(std::is_same_v, double>); + + int state = 0, cluster_count = 0; + int cluster_max = momenta.size() - 3; + int n_part = momenta.size(); + FourMom momenta_tmp[N_EXT_MAX]; + FVal masses_tmp[N_EXT_MAX]; + int alive = 0xFFFFFF; + int cluster_history[N_EXT_MAX - 3]; + FVal cluster_scales[N_EXT_MAX - 3]; + + for (int i = 0; i < n_part; ++i) { + for (int j = 0; j < 4; ++j) { + momenta_tmp[i][j] = momenta[i][j]; + } + masses_tmp[i] = external_masses[i]; + } + + int win_next_state = -1, win_data = 0; + bool win_resonant = false; + FVal win_scale = 1e308; + while (cluster_count < cluster_max) { + int data = state_machine[state]; + int next_state = state_machine[state + 1]; + int particle1 = data & 0xFF; + int particle2 = (data >> 8) & 0xFF; + int mass_index = (data >> 16) & 0xFF; + bool massive_in = (data >> 24) & 1; + bool massive_out1 = (data >> 25) & 1; + bool massive_out2 = (data >> 26) & 1; + bool is_last = (data >> 30) & 1; + bool is_initial = (particle1 < 2); + + FourMom momentum_sum{ + momenta_tmp[particle1][0] + momenta_tmp[particle2][0], + momenta_tmp[particle1][1] + momenta_tmp[particle2][1], + momenta_tmp[particle1][2] + momenta_tmp[particle2][2], + momenta_tmp[particle1][3] + momenta_tmp[particle2][3], + }; + + bool resonant = false; + if (mass_index != 0) { + FVal prop_m2 = lsquare(momentum_sum); + FVal mass = bw_masses[mass_index - 1]; + FVal width = bw_widths[mass_index - 1]; + FVal m_min = mass - width; + FVal m_max = mass + width; + resonant = (prop_m2 >= m_min * m_min) && (prop_m2 <= m_max * m_max); + } + + FVal scale = compute_scale( + momenta_tmp[particle1], + momenta_tmp[particle2], + momentum_sum, + masses_tmp[particle1], + masses_tmp[particle2], + resonant, + is_initial, + massive_in, + massive_out1, + massive_out2, + hadronic, + jet_radius + ); + + // The MG5 fortran code extracted the resonance structure from the integration + // channel. This is not always possible in MG7, so prefer resonant configs + // over non-resonant ones + if ((!win_resonant && resonant) || + (win_resonant == resonant && scale < win_scale)) { + win_next_state = next_state; + win_scale = scale; + win_data = data; + win_resonant = resonant; + } + if (is_last) { + int p1_win = win_data & 0xFF; + int p2_win = (win_data >> 8) & 0xFF; + update_momenta( + n_part, momenta_tmp, masses_tmp, alive, p1_win, p2_win, win_resonant + ); + state = win_next_state; + cluster_history[cluster_count] = win_data; + cluster_scales[cluster_count] = win_scale; + ++cluster_count; + win_scale = 1e308; + } else { + state += 2; + } + } + + // Renormalization scale: geometric mean of QCD clustering scales + // (non-QCD entries replaced by the maximum scale). + // Factorization scale: smallest QCD clustering scale. + FVal fac_scale = 1e308, max_scale = 0.0; + for (int i = 0; i < cluster_max; ++i) { + FVal scale = cluster_scales[i]; + bool is_qcd = (cluster_history[i] >> 27) & 1; + if (is_qcd && scale < fac_scale) { + fac_scale = scale; + } + if (scale > max_scale) { + max_scale = scale; + } + } + + for (int i = 0; i < n_part - 2; ++i) { + outgoing_scales[i] = 0.0; + } + FVal ren_scale_val = 1.0; + int is_last_cluster = 0b11111111'11111111'11111100; + for (int i = 0; i < cluster_max; ++i) { + FVal scale = cluster_scales[i]; + int data = cluster_history[i]; + int particle1 = data & 0xFF; + int particle2 = (data >> 8) & 0xFF; + bool is_qcd = (data >> 27) & 1; + bool is_jet1 = (data >> 28) & 1; + bool is_jet2 = (data >> 29) & 1; + if (is_qcd) { + if (is_jet1 && is_last_cluster & (1 << particle1)) { + outgoing_scales[particle1 - 2] = scale; + } + if (is_jet2 && is_last_cluster & (1 << particle2)) { + outgoing_scales[particle2 - 2] = scale; + } + ren_scale_val *= scale; + } else { + ren_scale_val *= max_scale; + } + is_last_cluster &= ~((1 << particle1) | (1 << particle2)); + } + ren_scale_val = pow(ren_scale_val, 1.0 / cluster_max); + if (fac_scale > ren_scale_val) { + fac_scale = ren_scale_val; + } + + ren_scale = ren_scale_val; + fact_scale1 = fac_scale; + fact_scale2 = fac_scale; + + int diag_count = state_machine[state]; + int rand_index = static_cast(FVal(random) * diag_count); + if (rand_index >= diag_count) { + rand_index = diag_count - 1; + } + diagram_index = state_machine[state + rand_index + 1]; +} + +template +KERNELSPEC void kernel_mlm_clustering_hadronic( + FIn momenta, + FIn random, + IIn state_machine, + FIn external_masses, + FIn bw_masses, + FIn bw_widths, + FIn bw_cutoff, + FIn jet_radius, + FOut ren_scale, + FOut fact_scale1, + FOut fact_scale2, + FOut outgoing_scales, + IOut diagram_index +) { + mlm_clustering( + momenta, + random, + state_machine, + external_masses, + bw_masses, + bw_widths, + bw_cutoff, + jet_radius, + ren_scale, + fact_scale1, + fact_scale2, + outgoing_scales, + diagram_index, + true + ); +} + +template +KERNELSPEC void kernel_mlm_clustering_leptonic( + FIn momenta, + FIn random, + IIn state_machine, + FIn external_masses, + FIn bw_masses, + FIn bw_widths, + FIn bw_cutoff, + FIn jet_radius, + FOut ren_scale, + FOut fact_scale1, + FOut fact_scale2, + FOut outgoing_scales, + IOut diagram_index +) { + mlm_clustering( + momenta, + random, + state_machine, + external_masses, + bw_masses, + bw_widths, + bw_cutoff, + jet_radius, + ren_scale, + fact_scale1, + fact_scale2, + outgoing_scales, + diagram_index, + false + ); +} + +} // namespace kernels +} // namespace madspace diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 2d808ee34..0ab7c670a 100644 --- a/madspace/src/phasespace/integrand.cpp +++ b/madspace/src/phasespace/integrand.cpp @@ -94,6 +94,12 @@ Integrand::Integrand( ret_types.push_back("partial_weight_product", batch_float); } } + if (energy_scale && energy_scale->is_mlm()) { + ret_types.push_back( + "cluster_scales", + batch_float_array(mapping.particle_count() - 2) + ); + } ret_types.push_back("random", batch_float_array(mapping.random_dim())); if (mapping.channel_count() > 1 && !std::holds_alternative(discrete_before)) { @@ -261,6 +267,11 @@ NamedVector Integrand::compute_channel_part_ret_types() const { } } + if (_energy_scale && _energy_scale->is_mlm()) { + ret.push_back("cluster_scales_acc", acc_float_array(particle_count - 2)); + ret.push_back("scale_diagram_index_acc", acc_int); + } + return ret; } @@ -547,6 +558,11 @@ NamedVector Integrand::build_channel_part( out.push_back("fact_scale2", scales.at(2)); } + if (_energy_scale && _energy_scale->is_mlm()) { + out.push_back("cluster_scales_acc", scales.at("outgoing_scales")); + out.push_back("scale_diagram_index_acc", scales.at("diagram_index")); + } + return out; } @@ -597,6 +613,9 @@ NamedVector Integrand::build_common_part( momenta_acc, _flavor_remap.size() > 0 ? fb.gather_int(flavor_id, _flavor_remap) : flavor_id, }; + if (_energy_scale && _energy_scale->is_mlm()) { + xs_args.push_back(args.at("scale_diagram_index_acc")); + } xs_args.push_back(x1_acc); xs_args.push_back(x2_acc); xs_args.push_back(flavor_id); @@ -777,6 +796,16 @@ NamedVector Integrand::build_common_part( ); } } + if (_energy_scale && _energy_scale->is_mlm()) { + auto outgoing_count = static_cast(_mapping.particle_count() - 2); + outputs.push_back( + "cluster_scales", + scatter_or_drop( + fb.full({0., batch_size_val, outgoing_count}), + args.at("cluster_scales_acc") + ) + ); + } outputs.push_back("random", optional_cut(args.at("r"))); if (has_permutations && !std::holds_alternative(_discrete_before)) { diff --git a/madspace/src/phasespace/mlm_clustering.cpp b/madspace/src/phasespace/mlm_clustering.cpp new file mode 100644 index 000000000..42e96f63e --- /dev/null +++ b/madspace/src/phasespace/mlm_clustering.cpp @@ -0,0 +1,272 @@ +#include "madspace/phasespace/mlm_clustering.hpp" + +#include +#include + +using namespace madspace; + +namespace { + +using StateKey = std::pair, std::vector>; +struct StateItem { + int next_state; + int particle1; + int particle2; + int mass_index; + bool massive_in; + bool massive_out1; + bool massive_out2; + bool is_qcd; + bool is_jet1; + bool is_jet2; +}; + +void find_clusterings( + const std::vector& particle_masks, + const std::vector& diagrams, + const nested_vector2& valid_diags, + nested_vector2& states, + std::map& state_map, + int prev_index +) { + int n_masks = particle_masks.size(); + for (int i = 0; i < n_masks - 1; ++i) { + for (int j = i + 1; j < n_masks; ++j) { + int mask_i = particle_masks.at(i), mask_j = particle_masks.at(j); + int mask = mask_i | mask_j; + auto& valid = valid_diags.at(mask); + if (valid.size() == 0) { + continue; + } + std::vector new_masks, new_diags; + new_masks.insert( + new_masks.end(), particle_masks.begin(), particle_masks.begin() + i + ); + new_masks.push_back(mask); + new_masks.insert( + new_masks.end(), + particle_masks.begin() + i + 1, + particle_masks.begin() + j + ); + new_masks.insert( + new_masks.end(), particle_masks.begin() + j + 1, particle_masks.end() + ); + for (int index : diagrams) { + if (std::find(valid.begin(), valid.end(), index) != valid.end()) { + new_diags.push_back(index); + } + } + + StateKey key; + if (new_masks.size() == 3) { + key = {{}, new_diags}; + } else { + key = {new_masks, new_diags}; + } + int index; + if (auto search = state_map.find(key); search != state_map.end()) { + index = search->second; + } else { + index = states.size(); + state_map[key] = index; + states.push_back({}); + } + states.at(prev_index) + .push_back({ + .next_state = index, + .particle1 = std::countr_zero(static_cast(mask_i)), + .particle2 = std::countr_zero(static_cast(mask_j)), + .mass_index = 0, + .massive_in = false, + .massive_out1 = false, + .massive_out2 = false, + .is_qcd = true, + .is_jet1 = true, + .is_jet2 = true, + }); + + if (new_masks.size() == 3) { + auto& current_states = states.at(index); + if (current_states.size() == 0) { + for (int diag_index : new_diags) { + current_states.push_back({ + .next_state = diag_index, + .particle1 = 0, + .particle2 = 0, + .mass_index = 0, + .massive_in = false, + .massive_out1 = false, + .massive_out2 = false, + .is_qcd = false, + .is_jet1 = false, + .is_jet2 = false, + }); + } + } + } else { + find_clusterings( + new_masks, new_diags, valid_diags, states, state_map, index + ); + } + } + } +} + +} // namespace + +MLMClustering::MLMClustering( + std::vector topologies, + nested_vector3 permutations, + nested_vector2 diagram_indices, + double bw_cutoff, + double jet_radius, + bool hadronic +) : + FunctionGenerator( + "MLMClustering", + {{"momenta", + batch_four_vec_array(topologies.at(0).outgoing_masses().size() + 2)}}, + {{"ren_scale", batch_float}, + {"fact_scale1", batch_float}, + {"fact_scale2", batch_float}, + {"outgoing_scales", + batch_float_array(topologies.at(0).outgoing_masses().size())}, + {"diagram_index", batch_int}} + ), + _bw_cutoff(bw_cutoff), + _jet_radius(jet_radius), + _hadronic(hadronic) { + auto& incoming_masses = topologies.at(0).incoming_masses(); + auto& outgoing_masses = topologies.at(0).outgoing_masses(); + std::size_t n_ext = outgoing_masses.size() + 2; + nested_vector2 valid_diags(1 << n_ext); + std::vector particle_masks; + std::vector all_diags; + + _external_masses.insert( + _external_masses.end(), incoming_masses.begin(), incoming_masses.end() + ); + _external_masses.insert( + _external_masses.end(), outgoing_masses.begin(), outgoing_masses.end() + ); + + // create a list of all diagram indices that are possible for a given clustering, + // where a binary encoding of the clustering is used + for (auto [topo, permutations, diag_indices] : + zip(topologies, permutations, diagram_indices)) { + for (auto [permutation, diag_index] : zip(permutations, diag_indices)) { + all_diags.push_back(diag_index); + particle_masks.assign(topo.decays().size(), 0); + for (std::size_t i = 2; i < permutation.size(); ++i) { + particle_masks.at(topo.outgoing_indices().at(permutation.at(i) - 2)) = 1 + << i; + } + + bool has_t_channel = topo.t_integration_order().size() > 0; + for (auto& decay : std::views::reverse(topo.decays())) { + if (decay.child_indices.size() == 0) { + continue; + } + if (decay.index == 0 && has_t_channel) { + continue; + } + if (decay.child_indices.size() > 2) { + throw std::logic_error("does not support 1->n decays with n > 2"); + } + int mask = + (particle_masks.at(decay.child_indices.at(0)) | + particle_masks.at(decay.child_indices.at(1))); + particle_masks.at(decay.index) = mask; + valid_diags.at(mask).push_back(diag_index); + } + + if (!has_t_channel) { + continue; + } + + // for the t-channel part, one of the initial state particles has to be + // involved in the clustering + for (int mask = 1; std::size_t index : topo.decays().at(0).child_indices) { + mask |= particle_masks.at(index); + valid_diags.at(mask).push_back(diag_index); + } + for (int mask = 2; + std::size_t index : + std::views::reverse(topo.decays().at(0).child_indices)) { + mask |= particle_masks.at(index); + valid_diags.at(mask).push_back(diag_index); + } + } + } + + std::vector masks; + masks.reserve(n_ext); + for (int i = 0; i < n_ext; ++i) { + masks.push_back(1 << i); + } + nested_vector2 states{{}}; + std::map state_map; + state_map[{{masks}, {all_diags}}] = 0; + find_clusterings(masks, all_diags, valid_diags, states, state_map, 0); + + std::vector first_indices; + first_indices.reserve(states.size()); + for (int offset = 0; auto& state : states) { + first_indices.push_back(offset); + if (state.at(0).particle1 == 0 && state.at(0).particle2 == 0) { + offset += 1 + state.size(); + } else { + offset += 2 * state.size(); + } + } + for (auto& state : states) { + if (state.at(0).particle1 == 0 && state.at(0).particle2 == 0) { + _cluster_state_machine.push_back(state.size()); + for (auto& item : state) { + _cluster_state_machine.push_back(item.next_state); + } + } else { + for (auto& item : state) { + _cluster_state_machine.push_back( + (item.particle1 << 0) + (item.particle2 << 8) + + (item.mass_index << 16) + (item.massive_in << 24) + + (item.massive_out1 << 25) + (item.massive_out2 << 26) + + (item.is_qcd << 27) + (item.is_jet1 << 28) + (item.is_jet2 << 29) + + ((&item == &state.back()) << 30) + ); + _cluster_state_machine.push_back(first_indices.at(item.next_state)); + } + } + } +} + +NamedVector MLMClustering::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + std::array mlm_out; + Value random = fb.squeeze(fb.random(fb.batch_size(args.values()), 1)); + if (_hadronic) { + mlm_out = fb.mlm_clustering_hadronic( + args.at(0), + random, + _cluster_state_machine, + _external_masses, + _bw_masses, + _bw_widths, + _bw_cutoff, + _jet_radius + ); + } else { + mlm_out = fb.mlm_clustering_leptonic( + args.at(0), + random, + _cluster_state_machine, + _external_masses, + _bw_masses, + _bw_widths, + _bw_cutoff, + _jet_radius + ); + } + return {return_types().keys(), {mlm_out.begin(), mlm_out.end()}}; +} diff --git a/madspace/src/phasespace/scale.cpp b/madspace/src/phasespace/scale.cpp index e6106314d..d39d71bcd 100644 --- a/madspace/src/phasespace/scale.cpp +++ b/madspace/src/phasespace/scale.cpp @@ -25,10 +25,17 @@ EnergyScale::EnergyScale( _fact_scale1(fact_scale1), _fact_scale2(fact_scale2) {} +EnergyScale::EnergyScale(const MLMClustering& clustering) : + FunctionGenerator("EnergyScale", clustering.arg_types(), clustering.return_types()), + _clustering(clustering) {} + NamedVector EnergyScale::build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const { auto momenta = args.at(0); + if (_clustering) { + return _clustering.value().build_function(fb, args); + } if (_ren_scale_fixed && _fact_scale_fixed) { auto batch_size = fb.batch_size({momenta}); return { diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index 73fceb63d..113a68ed8 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -121,6 +121,8 @@ void add_instructions(py::classh& fb) { fb.def("scale_transverse_mass", &FunctionBuilder::scale_transverse_mass, py::arg("momenta")); fb.def("scale_half_transverse_mass", &FunctionBuilder::scale_half_transverse_mass, py::arg("momenta")); fb.def("scale_partonic_energy", &FunctionBuilder::scale_partonic_energy, py::arg("momenta")); + fb.def("mlm_clustering_hadronic", &FunctionBuilder::mlm_clustering_hadronic, py::arg("momenta"), py::arg("random"), py::arg("state_machine"), py::arg("external_masses"), py::arg("bw_masses"), py::arg("bw_widths"), py::arg("bw_cutoffs"), py::arg("jet_radius")); + fb.def("mlm_clustering_leptonic", &FunctionBuilder::mlm_clustering_leptonic, py::arg("momenta"), py::arg("random"), py::arg("state_machine"), py::arg("external_masses"), py::arg("bw_masses"), py::arg("bw_widths"), py::arg("bw_cutoffs"), py::arg("jet_radius")); fb.def("chili_forward", &FunctionBuilder::chili_forward, py::arg("r"), py::arg("e_cm"), py::arg("m_out"), py::arg("pt_min"), py::arg("y_max")); fb.def("chili_inverse", &FunctionBuilder::chili_inverse, py::arg("p_ext"), py::arg("e_cm"), py::arg("m_out"), py::arg("pt_min"), py::arg("y_max")); fb.def("matrix_element", &FunctionBuilder::matrix_element, py::arg("args")); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 83a24af53..4266923a5 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -973,6 +973,23 @@ PYBIND11_MODULE(_madspace_py, m) { "initialize_globals", &DiscreteFlow::initialize_globals, py::arg("context") ); + py::classh(m, "MLMClustering") + .def( + py::init< + std::vector, + nested_vector3, + nested_vector2, + double, + double, + bool>(), + py::arg("topologies"), + py::arg("permutations"), + py::arg("diagram_indices"), + py::arg("bw_cutoff") = 15., + py::arg("jet_radius") = 0.4, + py::arg("hadronic") = true + ); + py::classh(m, "VegasGridOptimizer") .def( "add_data", @@ -1168,7 +1185,9 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("ren_scale"), py::arg("fact_scale1"), py::arg("fact_scale2") - ); + ) + .def(py::init(), py::arg("clustering")) + .def("is_mlm", &EnergyScale::is_mlm); py::classh(m, "CachedPdf").def(py::init<>()); py::classh(m, "CachedScale") diff --git a/madspace/tests/test_mlm.py b/madspace/tests/test_mlm.py new file mode 100644 index 000000000..8a8bc969e --- /dev/null +++ b/madspace/tests/test_mlm.py @@ -0,0 +1,60 @@ +import json +import os + +import numpy as np +import torch + +import madspace as ms + +BATCH_SIZE = 100 +CM_ENERGY = 13000.0 +rng = np.random.default_rng(5678) + +TTGG_FILE = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "processes", "ttgg.json" +) + + +def test_mlm_clustering(): + with open(TTGG_FILE) as f: + diagrams = json.load(f) + + topologies = [ + ms.Topology( + ms.Diagram( + d["incoming_masses"], + d["outgoing_masses"], + [ms.Propagator(*p) for p in d["propagators"]], + d["vertices"], + ) + ) + for d in diagrams + ] + permutations = [d["permutations"] for d in diagrams] + offset = 0 + diagram_indices = [] + for d in diagrams: + n_perms = len(d["permutations"]) + diagram_indices.append(list(range(offset, offset + n_perms))) + offset += n_perms + + clustering = ms.MLMClustering(topologies, permutations, diagram_indices) + + first = diagrams[0] + mapping = ms.PhaseSpaceMapping( + topologies[0], CM_ENERGY, permutations=first["permutations"] + ) + r = rng.random((BATCH_SIZE, mapping.random_dim())) + perm_count = len(first["permutations"]) + condition = ( + [] + if perm_count <= 1 + else [rng.integers(0, perm_count, BATCH_SIZE, dtype=np.int32)] + ) + p_ext, _, _, _ = mapping.map_forward([r], condition) + ren_scale, fact_scale1, fact_scale2, outgoing_scales, diagram_index = clustering( + p_ext + ) + from icecream import ic + + print(ren_scale, fact_scale1, fact_scale2, outgoing_scales, diagram_index)