From d3cff95f1842753a3195a388177df97495e9d78a Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 15 Jun 2026 17:47:44 +0200 Subject: [PATCH 01/12] started implementing mlm clustering --- madspace/CMakeLists.txt | 2 + .../compgraphs/function_builder_mixin.inc | 5 + .../madspace/compgraphs/opcode_mixin.inc | 97 ++++++++++--------- madspace/include/madspace/compgraphs/type.hpp | 3 + .../madspace/phasespace/mlm_clustering.hpp | 24 +++++ madspace/instruction_set.yaml | 25 +++++ .../src/compgraphs/instruction_set_mixin.inc | 97 ++++++++++--------- madspace/src/cpu/runtime_backward_mixin.inc | 42 ++++---- madspace/src/cpu/runtime_mixin.inc | 97 ++++++++++--------- madspace/src/driver/lhe_output.cpp | 2 + madspace/src/gpu/runtime_backward_mixin.inc | 42 ++++---- madspace/src/gpu/runtime_mixin.inc | 97 ++++++++++--------- madspace/src/phasespace/mlm_clustering.cpp | 83 ++++++++++++++++ madspace/src/python/instruction_set.hpp | 1 + 14 files changed, 385 insertions(+), 232 deletions(-) create mode 100644 madspace/include/madspace/phasespace/mlm_clustering.hpp create mode 100644 madspace/src/phasespace/mlm_clustering.cpp diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 1e7eb0adba..6a9de9c0b4 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -158,6 +158,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 @@ -213,6 +214,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 ) diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc index 94dbc20fbe..392d15e162 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/madspace/include/madspace/compgraphs/function_builder_mixin.inc @@ -461,6 +461,11 @@ Value scale_partonic_energy(Value momenta) { return instruction("scale_partonic_energy", {momenta})[0]; } +std::array mlm_clustering_hadronic(Value momenta, Value state_machine) { + auto output_vector = instruction("mlm_clustering_hadronic", {momenta, state_machine}); + 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 079c51b9dd..4d5bc962a4 100644 --- a/madspace/include/madspace/compgraphs/opcode_mixin.inc +++ b/madspace/include/madspace/compgraphs/opcode_mixin.inc @@ -102,51 +102,52 @@ scale_transverse_energy = 100, scale_transverse_mass = 101, scale_half_transverse_mass = 102, scale_partonic_energy = 103, -chili_forward = 104, -chili_inverse = 105, -matrix_element = 106, -collect_channel_weights = 107, -interpolate_pdf = 108, -interpolate_alpha_s = 109, -matmul = 110, -relu = 111, -leaky_relu = 112, -elu = 113, -gelu = 114, -sigmoid = 115, -softplus = 116, -rqs_reshape = 117, -rqs_find_bin = 118, -rqs_forward = 119, -rqs_inverse = 120, -softmax = 121, -softmax_prior = 122, -sample_discrete = 123, -sample_discrete_inverse = 124, -sample_discrete_probs = 125, -sample_discrete_probs_inverse = 126, -discrete_histogram = 127, -permute_momenta = 128, -gather = 129, -gather_int = 130, -select_int = 131, -select = 132, -select_vector = 133, -argsort = 134, -quantile = 135, -one_hot = 136, -madnis_abs_weight = 137, -madnis_softclip = 138, -madnis_variance = 139, -madnis_single_channel_variance = 140, -madnis_multi_channel_variance = 141, -nonzero = 142, -batch_gather = 143, -batch_scatter = 144, -random = 145, -random_int = 146, -unweight = 147, -vegas_forward = 148, -vegas_inverse = 149, -vegas_histogram = 150, -histogram = 151 +mlm_clustering_hadronic = 104, +chili_forward = 105, +chili_inverse = 106, +matrix_element = 107, +collect_channel_weights = 108, +interpolate_pdf = 109, +interpolate_alpha_s = 110, +matmul = 111, +relu = 112, +leaky_relu = 113, +elu = 114, +gelu = 115, +sigmoid = 116, +softplus = 117, +rqs_reshape = 118, +rqs_find_bin = 119, +rqs_forward = 120, +rqs_inverse = 121, +softmax = 122, +softmax_prior = 123, +sample_discrete = 124, +sample_discrete_inverse = 125, +sample_discrete_probs = 126, +sample_discrete_probs_inverse = 127, +discrete_histogram = 128, +permute_momenta = 129, +gather = 130, +gather_int = 131, +select_int = 132, +select = 133, +select_vector = 134, +argsort = 135, +quantile = 136, +one_hot = 137, +madnis_abs_weight = 138, +madnis_softclip = 139, +madnis_variance = 140, +madnis_single_channel_variance = 141, +madnis_multi_channel_variance = 142, +nonzero = 143, +batch_gather = 144, +batch_scatter = 145, +random = 146, +random_int = 147, +unweight = 148, +vegas_forward = 149, +vegas_inverse = 150, +vegas_histogram = 151, +histogram = 152 diff --git a/madspace/include/madspace/compgraphs/type.hpp b/madspace/include/madspace/compgraphs/type.hpp index f0d8ada0bd..787e20b219 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/phasespace/mlm_clustering.hpp b/madspace/include/madspace/phasespace/mlm_clustering.hpp new file mode 100644 index 0000000000..a21377a387 --- /dev/null +++ b/madspace/include/madspace/phasespace/mlm_clustering.hpp @@ -0,0 +1,24 @@ +#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 + ); + +private: + NamedVector build_function_impl( + FunctionBuilder& fb, const NamedVector& args + ) const override; + + std::vector _cluster_state_machine; +}; + +} // namespace madspace diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index ceeedf7051..de664a7b06 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -1947,6 +1947,31 @@ scale_partonic_energy: type: [float] desc: +mlm_clustering_hadronic: + inputs: + - name: momenta + type: [float, n, 4] + desc: + - name: state_machine + type: [int, m] + desc: + outputs: + - name: ren_scale + type: [float] + desc: + - name: fact_scale1 + type: [float] + desc: + - name: fact_scale2 + type: [float] + desc: + - name: cluster_history + type: [int, n-3] + desc: + - name: cluster_scales + type: [float, n-3] + desc: + --- title: Chili diff --git a/madspace/src/compgraphs/instruction_set_mixin.inc b/madspace/src/compgraphs/instruction_set_mixin.inc index 955d72bd25..0c1edccae6 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -119,52 +119,53 @@ InstructionOwner instructions[] { mi("scale_transverse_mass", 101, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), mi("scale_half_transverse_mass", 102, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), mi("scale_partonic_energy", 103, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("chili_forward", 104, 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", 105, 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(106, true)), - mi("collect_channel_weights", 107, 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", 108, 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", 109, 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", 110, 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", 111, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("leaky_relu", 112, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("elu", 113, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("gelu", 114, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("sigmoid", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softplus", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - InstructionOwner(new RqsReshapeInstruction(117, true)), - mi("rqs_find_bin", 118, 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", 119, 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", 120, 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", 121, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softmax_prior", 122, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("sample_discrete", 123, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_inverse", 124, 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", 125, 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", 126, 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", 127, 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", 128, 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", 129, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("gather_int", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), - mi("select_int", 131, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), - mi("select", 132, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("select_vector", 133, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), - mi("argsort", 134, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), - mi("quantile", 135, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}), - mi("one_hot", 136, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("madnis_abs_weight", 137, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_softclip", 138, 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", 139, 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", 140, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_multi_channel_variance", 141, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), - InstructionOwner(new NonzeroInstruction(142, true)), - InstructionOwner(new BatchGatherInstruction(143, true)), - InstructionOwner(new BatchScatterInstruction(144, true)), - InstructionOwner(new RandomInstruction(145, true)), - InstructionOwner(new RandomIntInstruction(146, true)), - InstructionOwner(new UnweightInstruction(147, true)), - mi("vegas_forward", 148, 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", 149, 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", 150, 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", 151, 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", 104, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n-3"}, false}, {DataType::dt_float, false, {"n-3"}, false}}), + mi("chili_forward", 105, 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", 106, 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(107, true)), + mi("collect_channel_weights", 108, 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", 109, 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", 110, 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", 111, 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", 112, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("leaky_relu", 113, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("elu", 114, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("gelu", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("sigmoid", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softplus", 117, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + InstructionOwner(new RqsReshapeInstruction(118, true)), + mi("rqs_find_bin", 119, 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", 120, 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", 121, 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", 122, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softmax_prior", 123, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("sample_discrete", 124, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_inverse", 125, 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", 126, 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", 127, 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", 128, 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", 129, 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", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("gather_int", 131, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), + mi("select_int", 132, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), + mi("select", 133, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_vector", 134, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), + mi("argsort", 135, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), + mi("quantile", 136, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}), + mi("one_hot", 137, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("madnis_abs_weight", 138, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_softclip", 139, 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", 140, 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", 141, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_multi_channel_variance", 142, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), + InstructionOwner(new NonzeroInstruction(143, true)), + InstructionOwner(new BatchGatherInstruction(144, true)), + InstructionOwner(new BatchScatterInstruction(145, true)), + InstructionOwner(new RandomInstruction(146, true)), + InstructionOwner(new RandomIntInstruction(147, true)), + InstructionOwner(new UnweightInstruction(148, true)), + mi("vegas_forward", 149, 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", 150, 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", 151, 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", 152, 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 5dd193731f..01f0b0e0a9 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 110: +case 111: backward_op_matmul(instr, locals, local_grads, device); break; -case 111: +case 112: backward_batch_foreach, backward_kernel_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 112: +case 113: backward_batch_foreach, backward_kernel_leaky_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 113: +case 114: backward_batch_foreach, backward_kernel_elu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 115: backward_batch_foreach, backward_kernel_gelu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 115: +case 116: backward_batch_foreach, backward_kernel_sigmoid, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 116: +case 117: backward_batch_foreach, backward_kernel_softplus, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 117: +case 118: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 118: +case 119: 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 119: +case 120: backward_batch_foreach, backward_kernel_rqs_forward, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 120: +case 121: backward_batch_foreach, backward_kernel_rqs_inverse, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 121: +case 122: backward_batch_foreach, backward_kernel_softmax, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 122: +case 123: backward_batch_foreach, backward_kernel_softmax_prior, 2, 2, 1, DeviceType>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 126: +case 127: 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 129: +case 130: backward_batch_foreach, backward_kernel_gather, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 132: +case 133: backward_batch_foreach, backward_kernel_select, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 137: +case 138: 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 138: +case 139: 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 139: +case 140: 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 140: +case 141: 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 141: +case 142: 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 3e3cc2400b..3e067702c5 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -314,146 +314,149 @@ case 103: batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 104: - batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_mlm_clustering_hadronic, 2, 5, 1, DeviceType>, 2, 5>(instr, locals, device); break; case 105: - batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 106: - op_matrix_element(instr, locals, device); + batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 107: - batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 108: - batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 109: - batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 110: - op_matmul(instr, locals, device); + batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 111: - batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 112: - batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 113: - batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 114: - batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 115: - batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 116: - batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 117: - op_rqs_reshape(instr, locals, device); + batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 118: - batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 119: - batch_foreach, kernel_rqs_forward, 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 120: - batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 121: - batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 122: - batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 123: - batch_foreach, kernel_sample_discrete, 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 124: - batch_foreach, kernel_sample_discrete_inverse, 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 125: - batch_foreach, kernel_sample_discrete_probs, 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 126: - batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 127: - op_discrete_histogram(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 128: - batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 129: - batch_foreach, kernel_gather, 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 130: - batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 131: - 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 132: - batch_foreach, kernel_select, 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 133: - batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 134: - batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 135: - op_quantile(instr, locals, device); + batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 136: - batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + op_quantile(instr, locals, device); break; case 137: - batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 138: - batch_foreach, kernel_madnis_softclip, 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 139: - batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 140: - batch_foreach, kernel_madnis_single_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 141: - batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 142: - op_nonzero(instr, locals, device); + batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 143: - op_batch_gather(instr, locals, device); + op_nonzero(instr, locals, device); break; case 144: - op_batch_scatter(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 145: - op_random(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 146: - op_random_int(instr, locals, device); + op_random(instr, locals, device); break; case 147: - op_unweight(instr, locals, device); + op_random_int(instr, locals, device); break; case 148: - batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_unweight(instr, locals, device); break; case 149: - batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 150: - op_vegas_histogram(instr, locals, device); + batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 151: + op_vegas_histogram(instr, locals, device); + break; +case 152: op_histogram(instr, locals, device); break; diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index 7624d0ac38..622c5e26d5 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -165,6 +165,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 59453fdabe..be422c24d5 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 110: +case 111: backward_op_matmul(instr, locals, local_grads, device); break; -case 111: +case 112: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 112: +case 113: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 113: +case 114: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 115: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 115: +case 116: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 116: +case 117: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 117: +case 118: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 118: +case 119: backward_batch_foreach, 5, 4, 2>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 119: +case 120: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 120: +case 121: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 121: +case 122: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 122: +case 123: backward_batch_foreach, 2, 2, 1>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 126: +case 127: backward_batch_foreach, 4, 2, 1>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 129: +case 130: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 132: +case 133: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 137: +case 138: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 138: +case 139: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 139: +case 140: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 140: +case 141: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 141: +case 142: 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 72271c9492..6e97b57ec7 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -314,146 +314,149 @@ case 103: batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 104: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 2, 5, 1>, 2, 5>(instr, locals, device); break; case 105: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 106: - op_matrix_element(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 107: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 108: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 109: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 110: - op_matmul(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 111: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 112: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 113: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 114: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 115: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 116: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 117: - op_rqs_reshape(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 118: - batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 119: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); break; case 120: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 121: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 122: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 123: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 124: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 125: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 126: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 127: - op_discrete_histogram(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 128: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 129: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 130: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 131: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 132: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 133: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 134: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 135: - op_quantile(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 136: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + op_quantile(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, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 139: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 140: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 141: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 142: - op_nonzero(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 143: - op_batch_gather(instr, locals, device); + op_nonzero(instr, locals, device); break; case 144: - op_batch_scatter(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 145: - op_random(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 146: - op_random_int(instr, locals, device); + op_random(instr, locals, device); break; case 147: - op_unweight(instr, locals, device); + op_random_int(instr, locals, device); break; case 148: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_unweight(instr, locals, device); break; case 149: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 150: - op_vegas_histogram(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 151: + op_vegas_histogram(instr, locals, device); + break; +case 152: op_histogram(instr, locals, device); break; diff --git a/madspace/src/phasespace/mlm_clustering.cpp b/madspace/src/phasespace/mlm_clustering.cpp new file mode 100644 index 0000000000..e4fe0cde27 --- /dev/null +++ b/madspace/src/phasespace/mlm_clustering.cpp @@ -0,0 +1,83 @@ +#include "madspace/phasespace/mlm_clustering.hpp" + +using namespace madspace; + +MLMClustering::MLMClustering( + std::vector topologies, + nested_vector3 permutations, + nested_vector2 diagram_indices +) : + 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}, + {"cluster_history", + batch_int_array(topologies.at(0).outgoing_masses().size() - 1)}, + {"cluster_scales", + batch_float_array(topologies.at(0).outgoing_masses().size() - 1)}} + ) { + std::size_t n_ext = topologies.at(0).outgoing_masses().size() + 2; + nested_vector2 valid_diags(1 << n_ext); + std::vector particle_masks; + std::vector all_diags; + + // 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))) = 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"); + } + std::size_t 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 (std::size_t 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 (std::size_t 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); + } + } + } +} + +NamedVector MLMClustering::build_function_impl( + FunctionBuilder& fb, const NamedVector& args +) const { + auto mlm_out = fb.mlm_clustering_hadronic(args.at(0), _cluster_state_machine); + return {return_types().keys(), {mlm_out.begin(), mlm_out.end()}}; +} diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index 1b2bc8e842..2a7623c193 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -117,6 +117,7 @@ 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("state_machine")); 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")); From 2d0b02ac89551e1b71186818c9240bd388457562 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Wed, 17 Jun 2026 11:04:49 +0200 Subject: [PATCH 02/12] progress on mlm --- .../compgraphs/function_builder_mixin.inc | 4 +- madspace/instruction_set.yaml | 20 ++- .../src/compgraphs/instruction_set_mixin.inc | 2 +- madspace/src/cpu/runtime_mixin.inc | 2 +- madspace/src/gpu/runtime_mixin.inc | 2 +- madspace/src/kernels/mlm.hpp | 167 ++++++++++++++++++ madspace/src/phasespace/mlm_clustering.cpp | 156 +++++++++++++++- madspace/src/python/instruction_set.hpp | 2 +- 8 files changed, 337 insertions(+), 18 deletions(-) create mode 100644 madspace/src/kernels/mlm.hpp diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc index 392d15e162..6d3f99b9a8 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/madspace/include/madspace/compgraphs/function_builder_mixin.inc @@ -461,8 +461,8 @@ Value scale_partonic_energy(Value momenta) { return instruction("scale_partonic_energy", {momenta})[0]; } -std::array mlm_clustering_hadronic(Value momenta, Value state_machine) { - auto output_vector = instruction("mlm_clustering_hadronic", {momenta, state_machine}); +std::array mlm_clustering_hadronic(Value momenta, Value random, Value state_machine, Value masses, Value widths) { + auto output_vector = instruction("mlm_clustering_hadronic", {momenta, random, state_machine, masses, widths}); return {output_vector[0], output_vector[1], output_vector[2], output_vector[3], output_vector[4]}; } diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index de664a7b06..5283a2221f 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -1952,8 +1952,17 @@ mlm_clustering_hadronic: - name: momenta type: [float, n, 4] desc: + - name: random + type: [float] + desc: - name: state_machine - type: [int, m] + type: [int, single, m] + desc: + - name: masses + type: [double, single, k] + desc: + - name: widths + type: [double, single, k] desc: outputs: - name: ren_scale @@ -1965,12 +1974,13 @@ mlm_clustering_hadronic: - name: fact_scale2 type: [float] desc: - - name: cluster_history - type: [int, n-3] + - name: outgoing_scales + type: [float, n-2] desc: - - name: cluster_scales - type: [float, n-3] + - 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 0c1edccae6..e71e59dc25 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -119,7 +119,7 @@ InstructionOwner instructions[] { mi("scale_transverse_mass", 101, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), mi("scale_half_transverse_mass", 102, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), mi("scale_partonic_energy", 103, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("mlm_clustering_hadronic", 104, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n-3"}, false}, {DataType::dt_float, false, {"n-3"}, false}}), + mi("mlm_clustering_hadronic", 104, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"m"}, false}, {DataType::dt_double, true, {"k"}, false}, {DataType::dt_double, true, {"k"}, 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", 105, 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", 106, 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(107, true)), diff --git a/madspace/src/cpu/runtime_mixin.inc b/madspace/src/cpu/runtime_mixin.inc index 3e067702c5..9ff7dee352 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -314,7 +314,7 @@ case 103: batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 104: - batch_foreach, kernel_mlm_clustering_hadronic, 2, 5, 1, DeviceType>, 2, 5>(instr, locals, device); + batch_foreach, kernel_mlm_clustering_hadronic, 5, 5, 1, DeviceType>, 5, 5>(instr, locals, device); break; case 105: batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); diff --git a/madspace/src/gpu/runtime_mixin.inc b/madspace/src/gpu/runtime_mixin.inc index 6e97b57ec7..dbc3488e69 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -314,7 +314,7 @@ case 103: batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 104: - batch_foreach, 2, 5, 1>, 2, 5>(instr, locals, device); + batch_foreach, 5, 5, 1>, 5, 5>(instr, locals, device); break; case 105: batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); diff --git a/madspace/src/kernels/mlm.hpp b/madspace/src/kernels/mlm.hpp new file mode 100644 index 0000000000..f80453a117 --- /dev/null +++ b/madspace/src/kernels/mlm.hpp @@ -0,0 +1,167 @@ +#pragma once + +#include "definitions.hpp" + +namespace madspace { +namespace kernels { + +constexpr int N_EXT_MAX = 12; + +KERNELSPEC void compute_scale( + double* momentum1, + double* momentum2, + bool resonant, + bool massive_in, + bool massive_out1, + bool massive_out2 +) { + /* + initial state: + cluster_scale=sqrt(djb_clus(pi)) + ! prefer clustering when outgoing is in the direction of incoming + if(sign(1d0,pi(3)).ne.sign(1d0,pj(3))) + $ cluster_scale=cluster_scale*one_plus_tiny + + final state: + resonant: + cluster_scale=sqrt(max(sumdot(pi,pj,1d0),0d0)) + + 2 : 1 0 1 + dj_clus = DJB_clus(p1)*(1d0+1d-6) + + 3 : 1 1 0 + dj_clus = DJB_clus(p2)*(1d0+1d-6) + + 1 : 0 0 0 + 6 : 0 1 1 + 7 : 1 1 1 + pt1 = p1(1)**2+p1(2)**2 + pt2 = p2(1)**2+p2(2)**2 + if (pt1.eq.0d0 .or. pt2.eq.0d0) then + dj_clus=0d0 + return + endif + p1a = dsqrt(pt1+p1(3)**2) + p2a = dsqrt(pt2+p2(3)**2) + eta1 = 0.5d0*log((p1a+p1(3))/(p1a-p1(3))) + eta2 = 0.5d0*log((p2a+p2(3))/(p2a-p2(3))) + dj_clus = max(m1,m2)**2+min(pt1,pt2)*2d0*(cosh(eta1-eta2)- + & (p1(1)*p2(1)+p1(2)*p2(2))/dsqrt(pt1*pt2))/D**2 + if (dj_clus.lt.0d0 .or. dj_clus.ne.dj_clus) + & dj_clus=0d0 ! prevent numerical inaccuracies + + 4 : 0 1 0 + cluster_scale=sqrt(abs(dot(pj,(pi+pj))))/2d0 + + 5 : 0 0 1 + cluster_scale=sqrt(abs(dot(pi,(pi+pj))))/2d0 + + 8 : 1 0 0 + resonant: + cluster_scale=sqrt(max(sumdot(pi,pj,1d0),0d0)) + */ +} + +template +KERNELSPEC void mlm_clustering_hadronic( + FIn momenta, + FIn random, + IIn state_machine, + FIn masses, + FIn widths, + FOut ren_scale, + FOut fact_scale1, + FOut fact_scale2, + FOut outgoing_scales, + IOut diagram_index +) { + static_assert(std::is_same_v, int>); + int state = 0, cluster_count = 0; + int cluster_max = momenta.size() - 3; + double momenta_tmp[N_EXT_MAX][4]; + int cluster_history[N_EXT_MAX - 3]; + int cluster_scales[N_EXT_MAX - 3]; + + for (int i = 0; i < momenta.size(); ++i) { + for (int j = 0; j < 4; ++j) { + momenta_tmp[i][j] = momenta[i][j]; + } + } + + // Iterate over the possible clustering states. This loop iterates over the possible + // clusterings at a given points as well as the full sequence of clusterings. + // Compared to two nested loops, this means that no thread has to wait on the GPU + int win_next_state = -1, win_data = 0; + double 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 >> 28) & 1; + + double scale = compute_scale_hadronic( + momenta_tmp[particle1], + momenta_tmp[particle2], + massive_in, + massive_out1, + massive_out2 + ); + + if (scale < win_scale) { + win_next_state = next_state; + win_scale = scale; + win_data = data; + } + if (is_last) { + state = win_next_state; + cluster_history[cluster_count] = win_data; + cluster_scales[cluster_count] = win_scale; + ++cluster_count; + } else { + ++state; + } + } + + // now determine the scales: + // renormalization scale: geometric mean of clustering scales, replacing those + // with is_qcd==false with the maximal scale + // factorization scale: minimal scale with is_qcd==true + double fac_scale = 1e308, max_scale = 0; + for (int i = 0; i < cluster_max; ++i) { + double 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; + } + } + + double ren_scale = 1.; + for (int i = 0; i < cluster_max; ++i) { + double scale = cluster_scales[i]; + bool is_qcd = (cluster_history[i] >> 27) & 1; + ren_scale *= is_qcd ? scale : max_scale; + } + ren_scale = pow(ren_scale, 1. / cluster_max); + if (fac_scale > ren_scale) { + fac_scale = ren_scale; + } + + // if multiple diagrams exist, pick one randomly + int diag_count = state_machine[state]; + int rand_index = static_cast(static_cast(random) * diag_count); + if (rand_index >= diag_count) { + rand_index = diag_count - 1; + } + diagram_index = state_machine[state + rand_index + 1]; +} + +} // namespace kernels +} // namespace madspace diff --git a/madspace/src/phasespace/mlm_clustering.cpp b/madspace/src/phasespace/mlm_clustering.cpp index e4fe0cde27..ec5d2b78d6 100644 --- a/madspace/src/phasespace/mlm_clustering.cpp +++ b/madspace/src/phasespace/mlm_clustering.cpp @@ -1,7 +1,111 @@ #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; +}; + +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 = false, + }); + + if (new_masks.size() == 3) { + auto& current_states = states.at(index); + 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, + }); + } + } else { + find_clusterings( + new_masks, new_diags, valid_diags, states, state_map, index + ); + } + } + } +} + +} // namespace + MLMClustering::MLMClustering( std::vector topologies, nested_vector3 permutations, @@ -20,9 +124,9 @@ MLMClustering::MLMClustering( batch_float_array(topologies.at(0).outgoing_masses().size() - 1)}} ) { std::size_t n_ext = topologies.at(0).outgoing_masses().size() + 2; - nested_vector2 valid_diags(1 << n_ext); - std::vector particle_masks; - std::vector all_diags; + nested_vector2 valid_diags(1 << n_ext); + std::vector particle_masks; + std::vector all_diags; // create a list of all diagram indices that are possible for a given clustering, // where a binary encoding of the clustering is used @@ -47,7 +151,7 @@ MLMClustering::MLMClustering( if (decay.child_indices.size() > 2) { throw std::logic_error("does not support 1->n decays with n > 2"); } - std::size_t mask = + int mask = (particle_masks.at(decay.child_indices.at(0)) | particle_masks.at(decay.child_indices.at(1))); particle_masks.at(decay.index) = mask; @@ -60,12 +164,11 @@ MLMClustering::MLMClustering( // for the t-channel part, one of the initial state particles has to be // involved in the clustering - for (std::size_t mask = 1; - std::size_t index : topo.decays().at(0).child_indices) { + 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 (std::size_t mask = 2; + for (int mask = 2; std::size_t index : std::views::reverse(topo.decays().at(0).child_indices)) { mask |= particle_masks.at(index); @@ -73,6 +176,45 @@ MLMClustering::MLMClustering( } } } + + 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 == &state.back()) << 28) + ); + _cluster_state_machine.push_back(first_indices.at(item.next_state)); + } + } + } } NamedVector MLMClustering::build_function_impl( diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index 2a7623c193..b1b52dc5e0 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -117,7 +117,7 @@ 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("state_machine")); + fb.def("mlm_clustering_hadronic", &FunctionBuilder::mlm_clustering_hadronic, py::arg("momenta"), py::arg("random"), py::arg("state_machine"), py::arg("masses"), py::arg("widths")); 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")); From 04f43eb99d8ead9c7530629bf5d02c7ff7655f22 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 25 Jun 2026 13:42:38 +0200 Subject: [PATCH 03/12] initial slop version of the mlm clustering scale --- madspace/src/kernels/mlm.hpp | 372 ++++++++++++++++++++++++++++------- 1 file changed, 306 insertions(+), 66 deletions(-) diff --git a/madspace/src/kernels/mlm.hpp b/madspace/src/kernels/mlm.hpp index f80453a117..62a7cec212 100644 --- a/madspace/src/kernels/mlm.hpp +++ b/madspace/src/kernels/mlm.hpp @@ -1,65 +1,279 @@ #pragma once #include "definitions.hpp" +#include "kinematics.hpp" + +#include namespace madspace { namespace kernels { constexpr int N_EXT_MAX = 12; -KERNELSPEC void compute_scale( - double* momentum1, - double* momentum2, +// --- Low-level kinematic helpers (mirror boost and rotate_inverse from +// kinematics.hpp +// but operating on plain double[4] so they can be used with the local scratch +// arrays) --- + +// Lorentz boost of k by the frame defined by p_boost (sign=+1 forward, -1 inverse). +// Mirrors boost from kinematics.hpp. +KERNELSPEC void +boost4(const double k[4], const double p_boost[4], double sign, double out[4]) { + double p2 = p_boost[0] * p_boost[0] - p_boost[1] * p_boost[1] - + p_boost[2] * p_boost[2] - p_boost[3] * p_boost[3]; + double rsq = sqrt(p2 > EPS2 ? p2 : EPS2); + double k_dot_p = k[1] * p_boost[1] + k[2] * p_boost[2] + k[3] * p_boost[3]; + double e = (k[0] * p_boost[0] + sign * k_dot_p) / rsq; + double c1 = sign * (k[0] + e) / (rsq + p_boost[0]); + out[0] = e; + out[1] = k[1] + c1 * p_boost[1]; + out[2] = k[2] + c1 * p_boost[2]; + out[3] = k[3] + c1 * p_boost[3]; +} + +// Rotate p into the canonical frame where q's spatial direction is the z-axis. +// Mirrors rotate_inverse from kinematics.hpp. +KERNELSPEC void rotate_inv4(const double p[4], const double q[4], double out[4]) { + double qt2 = q[1] * q[1] + q[2] * q[2]; + double qq2 = qt2 + q[3] * q[3]; + double qt = sqrt(qt2 > EPS2 ? qt2 : EPS2); + double qq = sqrt(qq2 > EPS2 ? qq2 : EPS2); + out[0] = p[0]; + if (qt2 < EPS2) { + double s = (q[3] < 0.0) ? -1.0 : 1.0; + out[1] = s * p[1]; + out[2] = p[2]; + out[3] = s * p[3]; + return; + } + out[1] = q[1] * q[3] / (qq * qt) * p[1] + q[2] * q[3] / (qq * qt) * p[2] - + p[3] * qt / qq; + out[2] = -q[2] / qt * p[1] + q[1] / qt * p[2]; + out[3] = q[1] / qq * p[1] + q[2] / qq * p[2] + q[3] / qq * p[3]; +} + +// --- Mass tracking --- + +// Update momenta and separately-tracked masses after one clustering step. +// +// i_remove : index of the final-state particle absorbed into the combination (always >= +// 2) j_keep : index of the particle kept with the combined momentum (0/1 = initial +// state, else final) is_bw : Breit-Wigner resonance — combined mass = invariant +// mass; else max of daughters alive : alive[i] is true while particle i has not yet +// been removed n_part : total number of particles (bound for the alive[] loop) +// +// Mass update rules from Fortran update_momenta: +// Final-state non-BW : mass[j_keep] = max(mass[j_keep], mass[i_remove]) +// Final-state BW : mass[j_keep] = sqrt((p_i + p_j)^2) +// Initial-state (XOR) : mass[j_keep] = max(mass[j_keep], mass[i_remove]) (exactly +// one massive) Initial-state (same): mass[j_keep] = 0 (both massive or both massless) +// +// Momentum update rules: +// Final-state : momenta[j_keep] += momenta[i_remove] +// Initial-state: momenta[j_keep] -= momenta[i_remove]; then if CM s > 100 GeV^2, +// boost all alive particles to the new CM frame and rotate so that +// the new j_keep momentum is along z. +KERNELSPEC void update_momenta( + int n_part, + double momenta[N_EXT_MAX][4], + double masses[N_EXT_MAX], + bool alive[N_EXT_MAX], + int i_remove, + int j_keep, + bool is_bw +) { + alive[i_remove] = false; + + if (j_keep < 2) { + // initial-state clustering + int j_other = 1 - j_keep; // the other beam particle (0-indexed) + for (int k = 0; k < 4; ++k) { + momenta[j_keep][k] -= momenta[i_remove][k]; + } + + // mass: max if exactly one daughter is massive, else 0 + bool m0 = (masses[j_keep] > 0.0); + bool m1 = (masses[i_remove] > 0.0); + masses[j_keep] = (m0 != m1) ? fmax(masses[j_keep], masses[i_remove]) : 0.0; + + // CM boost vector: (E_tot, -px_tot, -py_tot, -pz_tot) of the two beam particles + double pcmsp[4]; + pcmsp[0] = momenta[j_keep][0] + momenta[j_other][0]; + pcmsp[1] = -(momenta[j_keep][1] + momenta[j_other][1]); + pcmsp[2] = -(momenta[j_keep][2] + momenta[j_other][2]); + pcmsp[3] = -(momenta[j_keep][3] + momenta[j_other][3]); + + double inv_sq = pcmsp[0] * pcmsp[0] - pcmsp[1] * pcmsp[1] - + pcmsp[2] * pcmsp[2] - pcmsp[3] * pcmsp[3]; + if (inv_sq > 100.0) { + // boost j_keep to CM frame to define the rotation axis + double jkeep_cm[4]; + boost4(momenta[j_keep], pcmsp, 1.0, jkeep_cm); + + // boost all alive particles and rotate so that j_keep is along z + for (int j = 0; j < n_part; ++j) { + if (!alive[j]) { + continue; + } + double p_cm[4], p_rot[4]; + boost4(momenta[j], pcmsp, 1.0, p_cm); + rotate_inv4(p_cm, jkeep_cm, p_rot); + for (int k = 0; k < 4; ++k) { + momenta[j][k] = p_rot[k]; + } + } + } + // if inv_sq <= 100: no boost needed, momenta already correct after subtraction + + } else { + // final-state clustering: sum momenta, update mass + for (int k = 0; k < 4; ++k) { + momenta[j_keep][k] += momenta[i_remove][k]; + } + + if (is_bw) { + const double* p = momenta[j_keep]; + double m2 = p[0] * p[0] - p[1] * p[1] - p[2] * p[2] - p[3] * p[3]; + masses[j_keep] = sqrt(m2 > 0.0 ? m2 : 0.0); + } else { + masses[j_keep] = fmax(masses[j_keep], masses[i_remove]); + } + } +} + +// --- Clustering scale helpers --- + +KERNELSPEC double minkowski_dot(const double* p1, const double* p2) { + return p1[0] * p2[0] - p1[1] * p2[1] - p1[2] * p2[2] - p1[3] * p2[3]; +} + +// mT^2 = E^2 - pz^2 for one parton (hadronic), or E^2 (lepton collider). +// Replaces Fortran DJB_clus from cluster.f:2201. +KERNELSPEC double djb_clus(const double* p, bool hadronic) { + double r = hadronic ? (p[0] - p[3]) * (p[0] + p[3]) : p[0] * p[0]; + return r < 0.0 ? 0.0 : r; +} + +// kt/Durham clustering measure for two partons. +// mass1/mass2 are the tracked clustering masses from masses[] (set by update_momenta), +// NOT the Lorentz-invariant masses computed from the 4-momentum. +// D is the jet-radius parameter from Fortran common /to_dj/D. +// Replaces Fortran dj_clus from cluster.f:2144. +KERNELSPEC double dj_clus( + const double* p1, + const double* p2, + double mass1, + double mass2, + bool hadronic, + double D +) { + constexpr double one_plus_tiny = 1.0 + 1e-6; + if (!hadronic) { + // Durham e+e- measure: 2*min(E1^2,E2^2)*(1-cos_theta) + double p1a = sqrt(p1[1] * p1[1] + p1[2] * p1[2] + p1[3] * p1[3]); + double p2a = sqrt(p2[1] * p2[1] + p2[2] * p2[2] + p2[3] * p2[3]); + if (p1a * p2a == 0.0) { + return 0.0; + } + double costh = (p1[1] * p2[1] + p1[2] * p2[2] + p1[3] * p2[3]) / (p1a * p2a); + return 2.0 * fmin(p1[0] * p1[0], p2[0] * p2[0]) * fmax(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) + double pt1_sq = p1[1] * p1[1] + p1[2] * p1[2]; + double pt2_sq = p2[1] * p2[1] + p2[2] * p2[2]; + if (pt1_sq == 0.0 || pt2_sq == 0.0) { + return 0.0; + } + double p1a = sqrt(pt1_sq + p1[3] * p1[3]); + double p2a = sqrt(pt2_sq + p2[3] * p2[3]); + double eta1 = 0.5 * log((p1a + p1[3]) / (p1a - p1[3])); + double eta2 = 0.5 * log((p2a + p2[3]) / (p2a - p2[3])); + double m_max_sq = fmax(mass1 * mass1, mass2 * mass2); + double dphi_cos = (p1[1] * p2[1] + p1[2] * p2[2]) / sqrt(pt1_sq * pt2_sq); + double r = m_max_sq + + fmin(pt1_sq, pt2_sq) * 2.0 * (cosh(eta1 - eta2) - dphi_cos) / (D * D); + return r < 0.0 ? 0.0 : r; +} + +// Clustering scale for the pair (momentum1=pi, momentum2=pj). +// mass1/mass2 are the tracked clustering masses for momentum1/momentum2. +// +// Parameters replacing Fortran globals: +// is_initial : momentum2 is a beam particle (Fortran: j<=2 in cluster_one_step) +// hadronic : hadronic collider (Fortran: lpp[] from run.inc) +// D : jet-radius parameter (Fortran: common /to_dj/D) +// +// massive_in/out1/out2 replace the Fortran get_clustering_type cl[0:2] bit-array: +// massive_in = cl[0] has bit 2 or 4 (intermediate/mother) +// massive_out1 = cl[1] has bit 2 or 4 (momentum1 = final-state particle pi) +// massive_out2 = cl[2] has bit 2 or 4 (momentum2 = pj) +// +// resonant replaces the iBWlist lookup in Fortran cluster_scale. +// mass_in/width_in are carried for future use; not consumed in the scale formula. +// +// Replaces Fortran cluster_scale from cluster.f:1159. +KERNELSPEC double compute_scale( + const double* momentum1, + const double* momentum2, + double mass1, + double mass2, + double mass_in, + double width_in, bool resonant, + bool is_initial, bool massive_in, bool massive_out1, - bool massive_out2 + bool massive_out2, + bool hadronic, + double D ) { - /* - initial state: - cluster_scale=sqrt(djb_clus(pi)) - ! prefer clustering when outgoing is in the direction of incoming - if(sign(1d0,pi(3)).ne.sign(1d0,pj(3))) - $ cluster_scale=cluster_scale*one_plus_tiny - - final state: - resonant: - cluster_scale=sqrt(max(sumdot(pi,pj,1d0),0d0)) - - 2 : 1 0 1 - dj_clus = DJB_clus(p1)*(1d0+1d-6) - - 3 : 1 1 0 - dj_clus = DJB_clus(p2)*(1d0+1d-6) - - 1 : 0 0 0 - 6 : 0 1 1 - 7 : 1 1 1 - pt1 = p1(1)**2+p1(2)**2 - pt2 = p2(1)**2+p2(2)**2 - if (pt1.eq.0d0 .or. pt2.eq.0d0) then - dj_clus=0d0 - return - endif - p1a = dsqrt(pt1+p1(3)**2) - p2a = dsqrt(pt2+p2(3)**2) - eta1 = 0.5d0*log((p1a+p1(3))/(p1a-p1(3))) - eta2 = 0.5d0*log((p2a+p2(3))/(p2a-p2(3))) - dj_clus = max(m1,m2)**2+min(pt1,pt2)*2d0*(cosh(eta1-eta2)- - & (p1(1)*p2(1)+p1(2)*p2(2))/dsqrt(pt1*pt2))/D**2 - if (dj_clus.lt.0d0 .or. dj_clus.ne.dj_clus) - & dj_clus=0d0 ! prevent numerical inaccuracies - - 4 : 0 1 0 - cluster_scale=sqrt(abs(dot(pj,(pi+pj))))/2d0 - - 5 : 0 0 1 - cluster_scale=sqrt(abs(dot(pi,(pi+pj))))/2d0 - - 8 : 1 0 0 - resonant: - cluster_scale=sqrt(max(sumdot(pi,pj,1d0),0d0)) - */ + constexpr double one_plus_tiny = 1.000001; + + if (is_initial) { + // scale = mT of the final-state parton; small penalty when it goes against the + // beam + double scale = sqrt(djb_clus(momentum1, hadronic)); + if ((momentum1[3] < 0.0) != (momentum2[3] < 0.0)) { + scale *= one_plus_tiny; + } + return scale; + } + + double sum[4] = { + momentum1[0] + momentum2[0], + momentum1[1] + momentum2[1], + momentum1[2] + momentum2[2], + momentum1[3] + momentum2[3] + }; + + if (resonant) { + return sqrt(fmax(minkowski_dot(sum, sum), 0.0)); + } + + // Map massive_in/out1/out2 booleans to Fortran get_clustering_type itypes: + // type 4: massless mother -> massive out1 + massless out2 + if (!massive_in && massive_out1 && !massive_out2) { + return sqrt(fabs(minkowski_dot(momentum2, sum))) / 2.0; + } + // type 5: massless mother -> massless out1 + massive out2 + if (!massive_in && !massive_out1 && massive_out2) { + return sqrt(fabs(minkowski_dot(momentum1, sum))) / 2.0; + } + // type 8: massive mother -> massless out1 + massless out2 + if (massive_in && !massive_out1 && !massive_out2) { + return sqrt(fmax(minkowski_dot(sum, sum), 0.0)); + } + // types 1,2,3,6,7: all-massless, massive-emitting-massless, all-massive + return sqrt(dj_clus(momentum1, momentum2, mass1, mass2, hadronic, D)); } template @@ -69,6 +283,7 @@ KERNELSPEC void mlm_clustering_hadronic( IIn state_machine, FIn masses, FIn widths, + FIn D, FOut ren_scale, FOut fact_scale1, FOut fact_scale2, @@ -76,21 +291,24 @@ KERNELSPEC void mlm_clustering_hadronic( IOut diagram_index ) { 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(); double momenta_tmp[N_EXT_MAX][4]; + double masses_tmp[N_EXT_MAX]; + bool alive[N_EXT_MAX]; int cluster_history[N_EXT_MAX - 3]; - int cluster_scales[N_EXT_MAX - 3]; + double cluster_scales[N_EXT_MAX - 3]; - for (int i = 0; i < momenta.size(); ++i) { + 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] = masses[i]; + alive[i] = true; } - // Iterate over the possible clustering states. This loop iterates over the possible - // clusterings at a given points as well as the full sequence of clusterings. - // Compared to two nested loops, this means that no thread has to wait on the GPU int win_next_state = -1, win_data = 0; double win_scale = 1e308; while (cluster_count < cluster_max) { @@ -103,13 +321,24 @@ KERNELSPEC void mlm_clustering_hadronic( bool massive_out1 = (data >> 25) & 1; bool massive_out2 = (data >> 26) & 1; bool is_last = (data >> 28) & 1; + bool is_initial = (particle2 < 2); + bool resonant = + (mass_index != 0) && (static_cast(widths[mass_index]) > 0.0); - double scale = compute_scale_hadronic( + double scale = compute_scale( momenta_tmp[particle1], momenta_tmp[particle2], + masses_tmp[particle1], + masses_tmp[particle2], + static_cast(masses[mass_index]), + static_cast(widths[mass_index]), + resonant, + is_initial, massive_in, massive_out1, - massive_out2 + massive_out2, + true, // hadronic + static_cast(D) ); if (scale < win_scale) { @@ -118,20 +347,27 @@ KERNELSPEC void mlm_clustering_hadronic( win_data = data; } if (is_last) { + int p1_win = win_data & 0xFF; + int p2_win = (win_data >> 8) & 0xFF; + bool win_resonant = ((win_data >> 16) & 0xFF) != 0 && + (static_cast(widths[(win_data >> 16) & 0xFF]) > 0.0); + 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; } } - // now determine the scales: - // renormalization scale: geometric mean of clustering scales, replacing those - // with is_qcd==false with the maximal scale - // factorization scale: minimal scale with is_qcd==true - double fac_scale = 1e308, max_scale = 0; + // Renormalization scale: geometric mean of QCD clustering scales + // (non-QCD entries replaced by the max scale). + // Factorization scale: smallest QCD clustering scale. + double fac_scale = 1e308, max_scale = 0.0; for (int i = 0; i < cluster_max; ++i) { double scale = cluster_scales[i]; bool is_qcd = (cluster_history[i] >> 27) & 1; @@ -143,17 +379,21 @@ KERNELSPEC void mlm_clustering_hadronic( } } - double ren_scale = 1.; + double ren_scale_val = 1.0; for (int i = 0; i < cluster_max; ++i) { double scale = cluster_scales[i]; bool is_qcd = (cluster_history[i] >> 27) & 1; - ren_scale *= is_qcd ? scale : max_scale; + ren_scale_val *= is_qcd ? scale : max_scale; } - ren_scale = pow(ren_scale, 1. / cluster_max); - if (fac_scale > ren_scale) { - fac_scale = ren_scale; + 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; + // if multiple diagrams exist, pick one randomly int diag_count = state_machine[state]; int rand_index = static_cast(static_cast(random) * diag_count); From 85c81d49d210176982033d5dd4432562e34ed663 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 25 Jun 2026 14:22:50 +0200 Subject: [PATCH 04/12] begin desloppification --- madspace/src/kernels/kinematics.hpp | 10 + madspace/src/kernels/mlm.hpp | 452 ++++++++++++++-------------- 2 files changed, 230 insertions(+), 232 deletions(-) diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index cacfdcac40..26705daf64 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -139,6 +139,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 index 62a7cec212..f6ea2b3360 100644 --- a/madspace/src/kernels/mlm.hpp +++ b/madspace/src/kernels/mlm.hpp @@ -3,252 +3,112 @@ #include "definitions.hpp" #include "kinematics.hpp" -#include - namespace madspace { namespace kernels { constexpr int N_EXT_MAX = 12; -// --- Low-level kinematic helpers (mirror boost and rotate_inverse from -// kinematics.hpp -// but operating on plain double[4] so they can be used with the local scratch -// arrays) --- - -// Lorentz boost of k by the frame defined by p_boost (sign=+1 forward, -1 inverse). -// Mirrors boost from kinematics.hpp. -KERNELSPEC void -boost4(const double k[4], const double p_boost[4], double sign, double out[4]) { - double p2 = p_boost[0] * p_boost[0] - p_boost[1] * p_boost[1] - - p_boost[2] * p_boost[2] - p_boost[3] * p_boost[3]; - double rsq = sqrt(p2 > EPS2 ? p2 : EPS2); - double k_dot_p = k[1] * p_boost[1] + k[2] * p_boost[2] + k[3] * p_boost[3]; - double e = (k[0] * p_boost[0] + sign * k_dot_p) / rsq; - double c1 = sign * (k[0] + e) / (rsq + p_boost[0]); - out[0] = e; - out[1] = k[1] + c1 * p_boost[1]; - out[2] = k[2] + c1 * p_boost[2]; - out[3] = k[3] + c1 * p_boost[3]; -} - -// Rotate p into the canonical frame where q's spatial direction is the z-axis. -// Mirrors rotate_inverse from kinematics.hpp. -KERNELSPEC void rotate_inv4(const double p[4], const double q[4], double out[4]) { - double qt2 = q[1] * q[1] + q[2] * q[2]; - double qq2 = qt2 + q[3] * q[3]; - double qt = sqrt(qt2 > EPS2 ? qt2 : EPS2); - double qq = sqrt(qq2 > EPS2 ? qq2 : EPS2); - out[0] = p[0]; - if (qt2 < EPS2) { - double s = (q[3] < 0.0) ? -1.0 : 1.0; - out[1] = s * p[1]; - out[2] = p[2]; - out[3] = s * p[3]; - return; - } - out[1] = q[1] * q[3] / (qq * qt) * p[1] + q[2] * q[3] / (qq * qt) * p[2] - - p[3] * qt / qq; - out[2] = -q[2] / qt * p[1] + q[1] / qt * p[2]; - out[3] = q[1] / qq * p[1] + q[2] / qq * p[2] + q[3] / qq * p[3]; -} - -// --- Mass tracking --- - -// Update momenta and separately-tracked masses after one clustering step. -// -// i_remove : index of the final-state particle absorbed into the combination (always >= -// 2) j_keep : index of the particle kept with the combined momentum (0/1 = initial -// state, else final) is_bw : Breit-Wigner resonance — combined mass = invariant -// mass; else max of daughters alive : alive[i] is true while particle i has not yet -// been removed n_part : total number of particles (bound for the alive[] loop) -// -// Mass update rules from Fortran update_momenta: -// Final-state non-BW : mass[j_keep] = max(mass[j_keep], mass[i_remove]) -// Final-state BW : mass[j_keep] = sqrt((p_i + p_j)^2) -// Initial-state (XOR) : mass[j_keep] = max(mass[j_keep], mass[i_remove]) (exactly -// one massive) Initial-state (same): mass[j_keep] = 0 (both massive or both massless) -// -// Momentum update rules: -// Final-state : momenta[j_keep] += momenta[i_remove] -// Initial-state: momenta[j_keep] -= momenta[i_remove]; then if CM s > 100 GeV^2, -// boost all alive particles to the new CM frame and rotate so that -// the new j_keep momentum is along z. -KERNELSPEC void update_momenta( - int n_part, - double momenta[N_EXT_MAX][4], - double masses[N_EXT_MAX], - bool alive[N_EXT_MAX], - int i_remove, - int j_keep, - bool is_bw -) { - alive[i_remove] = false; - - if (j_keep < 2) { - // initial-state clustering - int j_other = 1 - j_keep; // the other beam particle (0-indexed) - for (int k = 0; k < 4; ++k) { - momenta[j_keep][k] -= momenta[i_remove][k]; - } - - // mass: max if exactly one daughter is massive, else 0 - bool m0 = (masses[j_keep] > 0.0); - bool m1 = (masses[i_remove] > 0.0); - masses[j_keep] = (m0 != m1) ? fmax(masses[j_keep], masses[i_remove]) : 0.0; - - // CM boost vector: (E_tot, -px_tot, -py_tot, -pz_tot) of the two beam particles - double pcmsp[4]; - pcmsp[0] = momenta[j_keep][0] + momenta[j_other][0]; - pcmsp[1] = -(momenta[j_keep][1] + momenta[j_other][1]); - pcmsp[2] = -(momenta[j_keep][2] + momenta[j_other][2]); - pcmsp[3] = -(momenta[j_keep][3] + momenta[j_other][3]); - - double inv_sq = pcmsp[0] * pcmsp[0] - pcmsp[1] * pcmsp[1] - - pcmsp[2] * pcmsp[2] - pcmsp[3] * pcmsp[3]; - if (inv_sq > 100.0) { - // boost j_keep to CM frame to define the rotation axis - double jkeep_cm[4]; - boost4(momenta[j_keep], pcmsp, 1.0, jkeep_cm); - - // boost all alive particles and rotate so that j_keep is along z - for (int j = 0; j < n_part; ++j) { - if (!alive[j]) { - continue; - } - double p_cm[4], p_rot[4]; - boost4(momenta[j], pcmsp, 1.0, p_cm); - rotate_inv4(p_cm, jkeep_cm, p_rot); - for (int k = 0; k < 4; ++k) { - momenta[j][k] = p_rot[k]; - } - } - } - // if inv_sq <= 100: no boost needed, momenta already correct after subtraction - - } else { - // final-state clustering: sum momenta, update mass - for (int k = 0; k < 4; ++k) { - momenta[j_keep][k] += momenta[i_remove][k]; - } - - if (is_bw) { - const double* p = momenta[j_keep]; - double m2 = p[0] * p[0] - p[1] * p[1] - p[2] * p[2] - p[3] * p[3]; - masses[j_keep] = sqrt(m2 > 0.0 ? m2 : 0.0); - } else { - masses[j_keep] = fmax(masses[j_keep], masses[i_remove]); - } - } -} - -// --- Clustering scale helpers --- - -KERNELSPEC double minkowski_dot(const double* p1, const double* p2) { - return p1[0] * p2[0] - p1[1] * p2[1] - p1[2] * p2[2] - p1[3] * p2[3]; -} - -// mT^2 = E^2 - pz^2 for one parton (hadronic), or E^2 (lepton collider). -// Replaces Fortran DJB_clus from cluster.f:2201. -KERNELSPEC double djb_clus(const double* p, bool hadronic) { - double r = hadronic ? (p[0] - p[3]) * (p[0] + p[3]) : p[0] * p[0]; - return r < 0.0 ? 0.0 : r; +// 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(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 from masses[] (set by update_momenta), -// NOT the Lorentz-invariant masses computed from the 4-momentum. -// D is the jet-radius parameter from Fortran common /to_dj/D. -// Replaces Fortran dj_clus from cluster.f:2144. -KERNELSPEC double dj_clus( - const double* p1, - const double* p2, - double mass1, - double mass2, +// 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( + FourMom p1, + FourMom p2, + FVal mass1, + FVal mass2, bool hadronic, - double D + FVal jet_radius ) { constexpr double one_plus_tiny = 1.0 + 1e-6; if (!hadronic) { - // Durham e+e- measure: 2*min(E1^2,E2^2)*(1-cos_theta) - double p1a = sqrt(p1[1] * p1[1] + p1[2] * p1[2] + p1[3] * p1[3]); - double p2a = sqrt(p2[1] * p2[1] + p2[2] * p2[2] + p2[3] * p2[3]); + // 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; } - double costh = (p1[1] * p2[1] + p1[2] * p2[2] + p1[3] * p2[3]) / (p1a * p2a); - return 2.0 * fmin(p1[0] * p1[0], p2[0] * p2[0]) * fmax(1.0 - costh, 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; + return djb_clus(p1, true) * one_plus_tiny; } if (massive1 && !massive2) { - return djb_clus(p2, true) * one_plus_tiny; + return djb_clus(p2, true) * one_plus_tiny; } // both massless or both massive: generalised kt measure in (eta, phi) - double pt1_sq = p1[1] * p1[1] + p1[2] * p1[2]; - double pt2_sq = p2[1] * p2[1] + p2[2] * p2[2]; + 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; } - double p1a = sqrt(pt1_sq + p1[3] * p1[3]); - double p2a = sqrt(pt2_sq + p2[3] * p2[3]); - double eta1 = 0.5 * log((p1a + p1[3]) / (p1a - p1[3])); - double eta2 = 0.5 * log((p2a + p2[3]) / (p2a - p2[3])); - double m_max_sq = fmax(mass1 * mass1, mass2 * mass2); - double dphi_cos = (p1[1] * p2[1] + p1[2] * p2[2]) / sqrt(pt1_sq * pt2_sq); - double r = m_max_sq + - fmin(pt1_sq, pt2_sq) * 2.0 * (cosh(eta1 - eta2) - dphi_cos) / (D * D); - return r < 0.0 ? 0.0 : r; + 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: -// is_initial : momentum2 is a beam particle (Fortran: j<=2 in cluster_one_step) -// hadronic : hadronic collider (Fortran: lpp[] from run.inc) -// D : jet-radius parameter (Fortran: common /to_dj/D) -// -// massive_in/out1/out2 replace the Fortran get_clustering_type cl[0:2] bit-array: -// massive_in = cl[0] has bit 2 or 4 (intermediate/mother) -// massive_out1 = cl[1] has bit 2 or 4 (momentum1 = final-state particle pi) -// massive_out2 = cl[2] has bit 2 or 4 (momentum2 = pj) +// 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 // -// resonant replaces the iBWlist lookup in Fortran cluster_scale. -// mass_in/width_in are carried for future use; not consumed in the scale formula. -// -// Replaces Fortran cluster_scale from cluster.f:1159. -KERNELSPEC double compute_scale( - const double* momentum1, - const double* momentum2, - double mass1, - double mass2, - double mass_in, - double width_in, +// based on cluster_scale from Template/NLO/SubProcesses/cluster.f +template +KERNELSPEC FVal compute_scale( + FourMom momentum1, + FourMom momentum2, + FVal mass1, + FVal mass2, bool resonant, bool is_initial, bool massive_in, bool massive_out1, bool massive_out2, bool hadronic, - double D + FVal jet_radius ) { constexpr double one_plus_tiny = 1.000001; if (is_initial) { - // scale = mT of the final-state parton; small penalty when it goes against the - // beam - double scale = sqrt(djb_clus(momentum1, hadronic)); + // scale = mT of the final-state parton + // small penalty when it goes against the beam + auto scale = sqrt(djb_clus(momentum1, hadronic)); if ((momentum1[3] < 0.0) != (momentum2[3] < 0.0)) { - scale *= one_plus_tiny; + scale = scale * one_plus_tiny; } return scale; } - double sum[4] = { + FourMom sum = { momentum1[0] + momentum2[0], momentum1[1] + momentum2[1], momentum1[2] + momentum2[2], @@ -256,50 +116,120 @@ KERNELSPEC double compute_scale( }; if (resonant) { - return sqrt(fmax(minkowski_dot(sum, sum), 0.0)); + return sqrt(max(lsquare(sum), 0.0)); } - // Map massive_in/out1/out2 booleans to Fortran get_clustering_type itypes: - // type 4: massless mother -> massive out1 + massless out2 + // Map massive_in/out1/out2 to Fortran get_clustering_type itypes: if (!massive_in && massive_out1 && !massive_out2) { - return sqrt(fabs(minkowski_dot(momentum2, sum))) / 2.0; + return sqrt(fabs(ldot(momentum2, sum))) / 2.0; } - // type 5: massless mother -> massless out1 + massive out2 if (!massive_in && !massive_out1 && massive_out2) { - return sqrt(fabs(minkowski_dot(momentum1, sum))) / 2.0; + return sqrt(fabs(ldot(momentum1, sum))) / 2.0; } - // type 8: massive mother -> massless out1 + massless out2 if (massive_in && !massive_out1 && !massive_out2) { - return sqrt(fmax(minkowski_dot(sum, sum), 0.0)); + return sqrt(max(lsquare(sum), 0.0)); } - // types 1,2,3,6,7: all-massless, massive-emitting-massless, all-massive - return sqrt(dj_clus(momentum1, momentum2, mass1, mass2, hadronic, D)); + return sqrt(dj_clus(momentum1, momentum2, mass1, mass2, hadronic, jet_radius)); } +// Update momenta and the separately-tracked masses after one clustering step. +// Uses boost and rotate_inverse from kinematics.hpp. +// +// i_remove : final-state particle absorbed into the combination (always >= 2, +// 0-indexed) j_keep : particle kept with the combined momentum (0/1 = initial state, +// else final) is_bw : BW resonance — combined mass = invariant mass of sum; else max +// of daughters +// +// Mass update rules mirror Fortran update_momenta (cluster.f:871): +// Final non-BW : mass[j_keep] = max(mass[j_keep], mass[i_remove]) +// Final BW : mass[j_keep] = sqrt((pi + pj)^2) +// Initial, one massive: mass[j_keep] = max(mass[j_keep], mass[i_remove]) +// Initial, same : mass[j_keep] = 0 template -KERNELSPEC void mlm_clustering_hadronic( +KERNELSPEC void update_momenta( + int n_part, + FourMom momenta[N_EXT_MAX], + FVal masses[N_EXT_MAX], + bool alive[N_EXT_MAX], + int i_remove, + int j_keep, + bool is_bw +) { + alive[i_remove] = false; + + if (j_keep < 2) { + // initial-state clustering + int j_other = 1 - j_keep; + for (int k = 0; k < 4; ++k) { + momenta[j_keep][k] -= momenta[i_remove][k]; + } + + // mass: take max if exactly one daughter is massive, else 0 + bool m0 = (masses[j_keep] > 0.0); + bool m1 = (masses[i_remove] > 0.0); + masses[j_keep] = + (m0 != m1) ? max(masses[j_keep], masses[i_remove]) : FVal(0.0); + + // CM boost vector: (E_tot, -px_tot, -py_tot, -pz_tot) of the two beam particles + FourMom pcmsp = { + momenta[j_keep][0] + momenta[j_other][0], + -(momenta[j_keep][1] + momenta[j_other][1]), + -(momenta[j_keep][2] + momenta[j_other][2]), + -(momenta[j_keep][3] + momenta[j_other][3]) + }; + + if (lsquare(pcmsp) > 100.0) { + // boost j_keep to CM to define the rotation axis, then apply to all alive + // particles + auto jkeep_cm = boost(momenta[j_keep], pcmsp, 1.0); + for (int j = 0; j < n_part; ++j) { + if (!alive[j]) { + continue; + } + momenta[j] = + rotate_inverse(boost(momenta[j], pcmsp, 1.0), jkeep_cm); + } + } + + } else { + // final-state clustering + for (int k = 0; k < 4; ++k) { + momenta[j_keep][k] += momenta[i_remove][k]; + } + + if (is_bw) { + masses[j_keep] = sqrt(max(lsquare(momenta[j_keep]), 0.0)); + } else { + masses[j_keep] = max(masses[j_keep], masses[i_remove]); + } + } +} + +template +KERNELSPEC void mlm_clustering( FIn momenta, FIn random, IIn state_machine, FIn masses, FIn widths, - FIn D, + FIn jet_radius, FOut ren_scale, FOut fact_scale1, FOut fact_scale2, FOut outgoing_scales, - IOut diagram_index + IOut diagram_index, + bool hadronic ) { 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(); - double momenta_tmp[N_EXT_MAX][4]; - double masses_tmp[N_EXT_MAX]; + FourMom momenta_tmp[N_EXT_MAX]; + FVal masses_tmp[N_EXT_MAX]; bool alive[N_EXT_MAX]; int cluster_history[N_EXT_MAX - 3]; - double cluster_scales[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) { @@ -310,7 +240,7 @@ KERNELSPEC void mlm_clustering_hadronic( } int win_next_state = -1, win_data = 0; - double win_scale = 1e308; + FVal win_scale = 1e308; while (cluster_count < cluster_max) { int data = state_machine[state]; int next_state = state_machine[state + 1]; @@ -322,23 +252,22 @@ KERNELSPEC void mlm_clustering_hadronic( bool massive_out2 = (data >> 26) & 1; bool is_last = (data >> 28) & 1; bool is_initial = (particle2 < 2); - bool resonant = - (mass_index != 0) && (static_cast(widths[mass_index]) > 0.0); + bool resonant = (mass_index != 0) && (widths[mass_index] > 0.0); - double scale = compute_scale( + FVal scale = compute_scale( momenta_tmp[particle1], momenta_tmp[particle2], masses_tmp[particle1], masses_tmp[particle2], - static_cast(masses[mass_index]), - static_cast(widths[mass_index]), + masses[mass_index], + widths[mass_index], resonant, is_initial, massive_in, massive_out1, massive_out2, - true, // hadronic - static_cast(D) + true, + jet_radius ); if (scale < win_scale) { @@ -349,9 +278,9 @@ KERNELSPEC void mlm_clustering_hadronic( if (is_last) { int p1_win = win_data & 0xFF; int p2_win = (win_data >> 8) & 0xFF; - bool win_resonant = ((win_data >> 16) & 0xFF) != 0 && - (static_cast(widths[(win_data >> 16) & 0xFF]) > 0.0); - update_momenta( + int mi_win = (win_data >> 16) & 0xFF; + bool win_resonant = (mi_win != 0) && (widths[mi_win] > 0.0); + update_momenta( n_part, momenta_tmp, masses_tmp, alive, p1_win, p2_win, win_resonant ); state = win_next_state; @@ -365,11 +294,11 @@ KERNELSPEC void mlm_clustering_hadronic( } // Renormalization scale: geometric mean of QCD clustering scales - // (non-QCD entries replaced by the max scale). + // (non-QCD entries replaced by the maximum scale). // Factorization scale: smallest QCD clustering scale. - double fac_scale = 1e308, max_scale = 0.0; + FVal fac_scale = 1e308, max_scale = 0.0; for (int i = 0; i < cluster_max; ++i) { - double scale = cluster_scales[i]; + FVal scale = cluster_scales[i]; bool is_qcd = (cluster_history[i] >> 27) & 1; if (is_qcd && scale < fac_scale) { fac_scale = scale; @@ -379,9 +308,9 @@ KERNELSPEC void mlm_clustering_hadronic( } } - double ren_scale_val = 1.0; + FVal ren_scale_val = 1.0; for (int i = 0; i < cluster_max; ++i) { - double scale = cluster_scales[i]; + FVal scale = cluster_scales[i]; bool is_qcd = (cluster_history[i] >> 27) & 1; ren_scale_val *= is_qcd ? scale : max_scale; } @@ -394,14 +323,73 @@ KERNELSPEC void mlm_clustering_hadronic( fact_scale1 = fac_scale; fact_scale2 = fac_scale; - // if multiple diagrams exist, pick one randomly int diag_count = state_machine[state]; - int rand_index = static_cast(static_cast(random) * diag_count); + 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 mlm_clustering_hadronic( + FIn momenta, + FIn random, + IIn state_machine, + FIn masses, + FIn widths, + FIn jet_radius, + FOut ren_scale, + FOut fact_scale1, + FOut fact_scale2, + FOut outgoing_scales, + IOut diagram_index +) { + mlm_clustering( + momenta, + random, + state_machine, + masses, + widths, + jet_radius, + ren_scale, + fact_scale1, + fact_scale2, + outgoing_scales, + diagram_index, + true + ); +} + +template +KERNELSPEC void mlm_clustering_leptonic( + FIn momenta, + FIn random, + IIn state_machine, + FIn masses, + FIn widths, + FIn jet_radius, + FOut ren_scale, + FOut fact_scale1, + FOut fact_scale2, + FOut outgoing_scales, + IOut diagram_index +) { + mlm_clustering( + momenta, + random, + state_machine, + masses, + widths, + jet_radius, + ren_scale, + fact_scale1, + fact_scale2, + outgoing_scales, + diagram_index, + false + ); +} + } // namespace kernels } // namespace madspace From 306950c5ffff3b4947e7f4e0c58ed93325303d22 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Thu, 25 Jun 2026 20:28:31 +0200 Subject: [PATCH 05/12] more progress on mlm --- madspace/CMakeLists.txt | 1 + .../compgraphs/function_builder_mixin.inc | 9 +- .../madspace/compgraphs/opcode_mixin.inc | 97 ++++++++-------- madspace/instruction_set.yaml | 61 +++++++++- .../src/compgraphs/instruction_set_mixin.inc | 99 ++++++++-------- madspace/src/cpu/runtime_backward_mixin.inc | 42 +++---- madspace/src/cpu/runtime_mixin.inc | 99 ++++++++-------- madspace/src/gpu/runtime_backward_mixin.inc | 42 +++---- madspace/src/gpu/runtime_mixin.inc | 99 ++++++++-------- madspace/src/kernels/kernels.hpp | 1 + madspace/src/kernels/kinematics.hpp | 2 +- madspace/src/kernels/mlm.hpp | 108 +++++++++++------- madspace/src/python/instruction_set.hpp | 3 +- 13 files changed, 379 insertions(+), 284 deletions(-) diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index b54850aa0a..df2e82c484 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -302,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 ea10552d04..ad31f0fb0c 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/madspace/include/madspace/compgraphs/function_builder_mixin.inc @@ -461,8 +461,13 @@ 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 masses, Value widths) { - auto output_vector = instruction("mlm_clustering_hadronic", {momenta, random, state_machine, masses, widths}); +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]}; } diff --git a/madspace/include/madspace/compgraphs/opcode_mixin.inc b/madspace/include/madspace/compgraphs/opcode_mixin.inc index 4d5bc962a4..ba363e13a1 100644 --- a/madspace/include/madspace/compgraphs/opcode_mixin.inc +++ b/madspace/include/madspace/compgraphs/opcode_mixin.inc @@ -103,51 +103,52 @@ scale_transverse_mass = 101, scale_half_transverse_mass = 102, scale_partonic_energy = 103, mlm_clustering_hadronic = 104, -chili_forward = 105, -chili_inverse = 106, -matrix_element = 107, -collect_channel_weights = 108, -interpolate_pdf = 109, -interpolate_alpha_s = 110, -matmul = 111, -relu = 112, -leaky_relu = 113, -elu = 114, -gelu = 115, -sigmoid = 116, -softplus = 117, -rqs_reshape = 118, -rqs_find_bin = 119, -rqs_forward = 120, -rqs_inverse = 121, -softmax = 122, -softmax_prior = 123, -sample_discrete = 124, -sample_discrete_inverse = 125, -sample_discrete_probs = 126, -sample_discrete_probs_inverse = 127, -discrete_histogram = 128, -permute_momenta = 129, -gather = 130, -gather_int = 131, -select_int = 132, -select = 133, -select_vector = 134, -argsort = 135, -quantile = 136, -one_hot = 137, -madnis_abs_weight = 138, -madnis_softclip = 139, -madnis_variance = 140, -madnis_single_channel_variance = 141, -madnis_multi_channel_variance = 142, -nonzero = 143, -batch_gather = 144, -batch_scatter = 145, -random = 146, -random_int = 147, -unweight = 148, -vegas_forward = 149, -vegas_inverse = 150, -vegas_histogram = 151, -histogram = 152 +mlm_clustering_leptonic = 105, +chili_forward = 106, +chili_inverse = 107, +matrix_element = 108, +collect_channel_weights = 109, +interpolate_pdf = 110, +interpolate_alpha_s = 111, +matmul = 112, +relu = 113, +leaky_relu = 114, +elu = 115, +gelu = 116, +sigmoid = 117, +softplus = 118, +rqs_reshape = 119, +rqs_find_bin = 120, +rqs_forward = 121, +rqs_inverse = 122, +softmax = 123, +softmax_prior = 124, +sample_discrete = 125, +sample_discrete_inverse = 126, +sample_discrete_probs = 127, +sample_discrete_probs_inverse = 128, +discrete_histogram = 129, +permute_momenta = 130, +gather = 131, +gather_int = 132, +select_int = 133, +select = 134, +select_vector = 135, +argsort = 136, +quantile = 137, +one_hot = 138, +madnis_abs_weight = 139, +madnis_softclip = 140, +madnis_variance = 141, +madnis_single_channel_variance = 142, +madnis_multi_channel_variance = 143, +nonzero = 144, +batch_gather = 145, +batch_scatter = 146, +random = 147, +random_int = 148, +unweight = 149, +vegas_forward = 150, +vegas_inverse = 151, +vegas_histogram = 152, +histogram = 153 diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 0a4d9413c9..9ff993a92a 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -1958,11 +1958,64 @@ mlm_clustering_hadronic: - name: state_machine type: [int, single, m] desc: - - name: masses - type: [double, single, k] + - name: external_masses + type: [float, n] desc: - - name: widths - type: [double, single, k] + - 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 diff --git a/madspace/src/compgraphs/instruction_set_mixin.inc b/madspace/src/compgraphs/instruction_set_mixin.inc index e71e59dc25..be0daf3231 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -119,53 +119,54 @@ InstructionOwner instructions[] { mi("scale_transverse_mass", 101, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), mi("scale_half_transverse_mass", 102, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), mi("scale_partonic_energy", 103, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("mlm_clustering_hadronic", 104, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"m"}, false}, {DataType::dt_double, true, {"k"}, false}, {DataType::dt_double, true, {"k"}, 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", 105, 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", 106, 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(107, true)), - mi("collect_channel_weights", 108, 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", 109, 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", 110, 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", 111, 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", 112, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("leaky_relu", 113, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("elu", 114, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("gelu", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("sigmoid", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softplus", 117, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - InstructionOwner(new RqsReshapeInstruction(118, true)), - mi("rqs_find_bin", 119, 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", 120, 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", 121, 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", 122, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softmax_prior", 123, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("sample_discrete", 124, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_inverse", 125, 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", 126, 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", 127, 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", 128, 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", 129, 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", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("gather_int", 131, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), - mi("select_int", 132, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), - mi("select", 133, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("select_vector", 134, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), - mi("argsort", 135, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), - mi("quantile", 136, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}), - mi("one_hot", 137, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("madnis_abs_weight", 138, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_softclip", 139, 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", 140, 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", 141, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_multi_channel_variance", 142, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), - InstructionOwner(new NonzeroInstruction(143, true)), - InstructionOwner(new BatchGatherInstruction(144, true)), - InstructionOwner(new BatchScatterInstruction(145, true)), - InstructionOwner(new RandomInstruction(146, true)), - InstructionOwner(new RandomIntInstruction(147, true)), - InstructionOwner(new UnweightInstruction(148, true)), - mi("vegas_forward", 149, 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", 150, 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", 151, 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", 152, 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", 104, 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", 105, 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", 106, 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", 107, 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(108, true)), + mi("collect_channel_weights", 109, 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", 110, 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", 111, 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", 112, 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", 113, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("leaky_relu", 114, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("elu", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("gelu", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("sigmoid", 117, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softplus", 118, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + InstructionOwner(new RqsReshapeInstruction(119, true)), + mi("rqs_find_bin", 120, 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", 121, 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", 122, 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", 123, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softmax_prior", 124, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("sample_discrete", 125, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_inverse", 126, 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", 127, 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", 128, 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", 129, 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", 130, 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", 131, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("gather_int", 132, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), + mi("select_int", 133, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), + mi("select", 134, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_vector", 135, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), + mi("argsort", 136, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), + mi("quantile", 137, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}), + mi("one_hot", 138, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("madnis_abs_weight", 139, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_softclip", 140, 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", 141, 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", 142, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_multi_channel_variance", 143, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), + InstructionOwner(new NonzeroInstruction(144, true)), + InstructionOwner(new BatchGatherInstruction(145, true)), + InstructionOwner(new BatchScatterInstruction(146, true)), + InstructionOwner(new RandomInstruction(147, true)), + InstructionOwner(new RandomIntInstruction(148, true)), + InstructionOwner(new UnweightInstruction(149, true)), + mi("vegas_forward", 150, 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", 151, 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", 152, 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", 153, 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 01f0b0e0a9..1a92043e6d 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 111: +case 112: backward_op_matmul(instr, locals, local_grads, device); break; -case 112: +case 113: backward_batch_foreach, backward_kernel_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 113: +case 114: backward_batch_foreach, backward_kernel_leaky_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 115: backward_batch_foreach, backward_kernel_elu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 115: +case 116: backward_batch_foreach, backward_kernel_gelu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 116: +case 117: backward_batch_foreach, backward_kernel_sigmoid, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 117: +case 118: backward_batch_foreach, backward_kernel_softplus, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 118: +case 119: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 119: +case 120: 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 120: +case 121: backward_batch_foreach, backward_kernel_rqs_forward, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 121: +case 122: backward_batch_foreach, backward_kernel_rqs_inverse, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 122: +case 123: backward_batch_foreach, backward_kernel_softmax, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 123: +case 124: backward_batch_foreach, backward_kernel_softmax_prior, 2, 2, 1, DeviceType>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 127: +case 128: 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 130: +case 131: backward_batch_foreach, backward_kernel_gather, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 133: +case 134: backward_batch_foreach, backward_kernel_select, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 138: +case 139: 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 139: +case 140: 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 140: +case 141: 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 141: +case 142: 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 142: +case 143: 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 9ff7dee352..68e688f743 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -314,149 +314,152 @@ case 103: batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 104: - batch_foreach, kernel_mlm_clustering_hadronic, 5, 5, 1, DeviceType>, 5, 5>(instr, locals, device); + batch_foreach, kernel_mlm_clustering_hadronic, 8, 5, 1, DeviceType>, 8, 5>(instr, locals, device); break; case 105: - batch_foreach, kernel_chili_forward, 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 106: - batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 107: - op_matrix_element(instr, locals, device); + batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 108: - batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 109: - batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 110: - batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 111: - op_matmul(instr, locals, device); + batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 112: - batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 113: - batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 114: - batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 115: - batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 116: - batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 117: - batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 118: - op_rqs_reshape(instr, locals, device); + batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 119: - batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 120: - batch_foreach, kernel_rqs_forward, 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 121: - batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 122: - batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 123: - batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 124: - batch_foreach, kernel_sample_discrete, 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 125: - batch_foreach, kernel_sample_discrete_inverse, 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 126: - batch_foreach, kernel_sample_discrete_probs, 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 127: - batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 128: - op_discrete_histogram(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 129: - batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 130: - batch_foreach, kernel_gather, 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 131: - batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 132: - 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 133: - batch_foreach, kernel_select, 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 134: - batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 135: - batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 136: - op_quantile(instr, locals, device); + batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 137: - batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + op_quantile(instr, locals, device); break; case 138: - batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 139: - batch_foreach, kernel_madnis_softclip, 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 140: - batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 141: - batch_foreach, kernel_madnis_single_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 142: - batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 143: - op_nonzero(instr, locals, device); + batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 144: - op_batch_gather(instr, locals, device); + op_nonzero(instr, locals, device); break; case 145: - op_batch_scatter(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 146: - op_random(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 147: - op_random_int(instr, locals, device); + op_random(instr, locals, device); break; case 148: - op_unweight(instr, locals, device); + op_random_int(instr, locals, device); break; case 149: - batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_unweight(instr, locals, device); break; case 150: - batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 151: - op_vegas_histogram(instr, locals, device); + batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 152: + op_vegas_histogram(instr, locals, device); + break; +case 153: op_histogram(instr, locals, device); break; diff --git a/madspace/src/gpu/runtime_backward_mixin.inc b/madspace/src/gpu/runtime_backward_mixin.inc index be422c24d5..57bfd1fc73 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 111: +case 112: backward_op_matmul(instr, locals, local_grads, device); break; -case 112: +case 113: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 113: +case 114: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 115: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 115: +case 116: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 116: +case 117: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 117: +case 118: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 118: +case 119: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 119: +case 120: backward_batch_foreach, 5, 4, 2>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 120: +case 121: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 121: +case 122: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 122: +case 123: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 123: +case 124: backward_batch_foreach, 2, 2, 1>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 127: +case 128: backward_batch_foreach, 4, 2, 1>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 130: +case 131: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 133: +case 134: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 138: +case 139: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 139: +case 140: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 140: +case 141: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 141: +case 142: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 142: +case 143: 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 dbc3488e69..f975f0adc8 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -314,149 +314,152 @@ case 103: batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 104: - batch_foreach, 5, 5, 1>, 5, 5>(instr, locals, device); + batch_foreach, 8, 5, 1>, 8, 5>(instr, locals, device); break; case 105: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 8, 5, 1>, 8, 5>(instr, locals, device); break; case 106: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 107: - op_matrix_element(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 108: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 109: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 110: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 111: - op_matmul(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 112: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 113: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 114: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 115: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 116: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(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: - op_rqs_reshape(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 119: - batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 120: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); break; case 121: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 122: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 123: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 124: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 125: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 126: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 127: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 128: - op_discrete_histogram(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 129: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 130: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 131: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 132: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 133: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 134: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 135: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 136: - op_quantile(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 137: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + op_quantile(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, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 140: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 141: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 142: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 143: - op_nonzero(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 144: - op_batch_gather(instr, locals, device); + op_nonzero(instr, locals, device); break; case 145: - op_batch_scatter(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 146: - op_random(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 147: - op_random_int(instr, locals, device); + op_random(instr, locals, device); break; case 148: - op_unweight(instr, locals, device); + op_random_int(instr, locals, device); break; case 149: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_unweight(instr, locals, device); break; case 150: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 151: - op_vegas_histogram(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 152: + op_vegas_histogram(instr, locals, device); + break; +case 153: op_histogram(instr, locals, device); break; diff --git a/madspace/src/kernels/kernels.hpp b/madspace/src/kernels/kernels.hpp index 7c2d828efc..307397aef0 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 26705daf64..04dd0fd322 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]; }; diff --git a/madspace/src/kernels/mlm.hpp b/madspace/src/kernels/mlm.hpp index f6ea2b3360..a68979f816 100644 --- a/madspace/src/kernels/mlm.hpp +++ b/madspace/src/kernels/mlm.hpp @@ -18,7 +18,7 @@ KERNELSPEC FVal djb_clus(FourMom p, bool hadronic) { // 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. +// 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 @@ -84,8 +84,8 @@ KERNELSPEC FVal dj_clus( // based on cluster_scale from Template/NLO/SubProcesses/cluster.f template KERNELSPEC FVal compute_scale( - FourMom momentum1, - FourMom momentum2, + const FourMom& momentum1, + const FourMom& momentum2, FVal mass1, FVal mass2, bool resonant, @@ -148,14 +148,14 @@ KERNELSPEC FVal compute_scale( template KERNELSPEC void update_momenta( int n_part, - FourMom momenta[N_EXT_MAX], - FVal masses[N_EXT_MAX], - bool alive[N_EXT_MAX], + FourMom* momenta, + FVal* masses, + int& alive, int i_remove, int j_keep, - bool is_bw + bool resonant ) { - alive[i_remove] = false; + alive &= ~(1 << i_remove); if (j_keep < 2) { // initial-state clustering @@ -165,10 +165,9 @@ KERNELSPEC void update_momenta( } // mass: take max if exactly one daughter is massive, else 0 - bool m0 = (masses[j_keep] > 0.0); - bool m1 = (masses[i_remove] > 0.0); - masses[j_keep] = - (m0 != m1) ? max(masses[j_keep], masses[i_remove]) : FVal(0.0); + masses[j_keep] = (masses[j_keep] > 0.0) != (masses[i_remove] > 0.0) + ? max(masses[j_keep], masses[i_remove]) + : FVal(0.0); // CM boost vector: (E_tot, -px_tot, -py_tot, -pz_tot) of the two beam particles FourMom pcmsp = { @@ -183,11 +182,10 @@ KERNELSPEC void update_momenta( // particles auto jkeep_cm = boost(momenta[j_keep], pcmsp, 1.0); for (int j = 0; j < n_part; ++j) { - if (!alive[j]) { - continue; + if (alive & (1 << j)) { + momenta[j] = + rotate_inverse(boost(momenta[j], pcmsp, 1.0), jkeep_cm); } - momenta[j] = - rotate_inverse(boost(momenta[j], pcmsp, 1.0), jkeep_cm); } } @@ -197,7 +195,7 @@ KERNELSPEC void update_momenta( momenta[j_keep][k] += momenta[i_remove][k]; } - if (is_bw) { + if (resonant) { masses[j_keep] = sqrt(max(lsquare(momenta[j_keep]), 0.0)); } else { masses[j_keep] = max(masses[j_keep], masses[i_remove]); @@ -210,8 +208,10 @@ KERNELSPEC void mlm_clustering( FIn momenta, FIn random, IIn state_machine, - FIn masses, - FIn widths, + FIn external_masses, + FIn bw_masses, + FIn bw_widths, + FIn bw_cutoff, FIn jet_radius, FOut ren_scale, FOut fact_scale1, @@ -220,14 +220,16 @@ KERNELSPEC void mlm_clustering( 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]; - bool alive[N_EXT_MAX]; + int alive = 0xFFFFFF; int cluster_history[N_EXT_MAX - 3]; FVal cluster_scales[N_EXT_MAX - 3]; @@ -235,11 +237,11 @@ KERNELSPEC void mlm_clustering( for (int j = 0; j < 4; ++j) { momenta_tmp[i][j] = momenta[i][j]; } - masses_tmp[i] = masses[i]; - alive[i] = true; + 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]; @@ -252,34 +254,50 @@ KERNELSPEC void mlm_clustering( bool massive_out2 = (data >> 26) & 1; bool is_last = (data >> 28) & 1; bool is_initial = (particle2 < 2); - bool resonant = (mass_index != 0) && (widths[mass_index] > 0.0); + + bool resonant = false; + if (mass_index != 0) { + FourMom mom_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], + }; + FVal prop_m2 = lsquare(mom_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], masses_tmp[particle1], masses_tmp[particle2], - masses[mass_index], - widths[mass_index], resonant, is_initial, massive_in, massive_out1, massive_out2, - true, + hadronic, jet_radius ); - if (scale < win_scale) { + // 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; - int mi_win = (win_data >> 16) & 0xFF; - bool win_resonant = (mi_win != 0) && (widths[mi_win] > 0.0); update_momenta( n_part, momenta_tmp, masses_tmp, alive, p1_win, p2_win, win_resonant ); @@ -332,12 +350,14 @@ KERNELSPEC void mlm_clustering( } template -KERNELSPEC void mlm_clustering_hadronic( +KERNELSPEC void kernel_mlm_clustering_hadronic( FIn momenta, FIn random, IIn state_machine, - FIn masses, - FIn widths, + FIn external_masses, + FIn bw_masses, + FIn bw_widths, + FIn bw_cutoff, FIn jet_radius, FOut ren_scale, FOut fact_scale1, @@ -345,12 +365,14 @@ KERNELSPEC void mlm_clustering_hadronic( FOut outgoing_scales, IOut diagram_index ) { - mlm_clustering( + mlm_clustering( momenta, random, state_machine, - masses, - widths, + external_masses, + bw_masses, + bw_widths, + bw_cutoff, jet_radius, ren_scale, fact_scale1, @@ -362,12 +384,14 @@ KERNELSPEC void mlm_clustering_hadronic( } template -KERNELSPEC void mlm_clustering_leptonic( +KERNELSPEC void kernel_mlm_clustering_leptonic( FIn momenta, FIn random, IIn state_machine, - FIn masses, - FIn widths, + FIn external_masses, + FIn bw_masses, + FIn bw_widths, + FIn bw_cutoff, FIn jet_radius, FOut ren_scale, FOut fact_scale1, @@ -375,12 +399,14 @@ KERNELSPEC void mlm_clustering_leptonic( FOut outgoing_scales, IOut diagram_index ) { - mlm_clustering( + mlm_clustering( momenta, random, state_machine, - masses, - widths, + external_masses, + bw_masses, + bw_widths, + bw_cutoff, jet_radius, ren_scale, fact_scale1, diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index fc952bdcdc..9dc330b23f 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -117,7 +117,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("masses"), py::arg("widths")); + 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")); From 350bde867fafa901817de9bce4dc99cb4af2d8fc Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Sun, 28 Jun 2026 21:58:12 +0200 Subject: [PATCH 06/12] some more cleaning up in mlm --- madspace/src/kernels/mlm.hpp | 109 ++++++++++++++--------------------- 1 file changed, 44 insertions(+), 65 deletions(-) diff --git a/madspace/src/kernels/mlm.hpp b/madspace/src/kernels/mlm.hpp index a68979f816..5f0bb4f262 100644 --- a/madspace/src/kernels/mlm.hpp +++ b/madspace/src/kernels/mlm.hpp @@ -7,6 +7,7 @@ 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 @@ -30,7 +31,6 @@ KERNELSPEC FVal dj_clus( bool hadronic, FVal jet_radius ) { - constexpr double one_plus_tiny = 1.0 + 1e-6; if (!hadronic) { // Durham e+e- measure: 2*min(E1^2, E2^2)*(1 - cos_theta) auto p1a = sqrt(esquare(p1)); @@ -41,15 +41,17 @@ KERNELSPEC FVal dj_clus( 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; + return djb_clus(p1, true) * ONE_PLUS_TINY; } if (massive1 && !massive2) { - return djb_clus(p2, true) * one_plus_tiny; + 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]; @@ -86,6 +88,7 @@ template KERNELSPEC FVal compute_scale( const FourMom& momentum1, const FourMom& momentum2, + const FourMom& momentum_sum, FVal mass1, FVal mass2, bool resonant, @@ -96,55 +99,33 @@ KERNELSPEC FVal compute_scale( bool hadronic, FVal jet_radius ) { - constexpr double one_plus_tiny = 1.000001; if (is_initial) { // scale = mT of the final-state parton // small penalty when it goes against the beam auto scale = sqrt(djb_clus(momentum1, hadronic)); if ((momentum1[3] < 0.0) != (momentum2[3] < 0.0)) { - scale = scale * one_plus_tiny; + scale = scale * ONE_PLUS_TINY; } return scale; } - - FourMom sum = { - momentum1[0] + momentum2[0], - momentum1[1] + momentum2[1], - momentum1[2] + momentum2[2], - momentum1[3] + momentum2[3] - }; - if (resonant) { - return sqrt(max(lsquare(sum), 0.0)); + return sqrt(max(lsquare(momentum_sum), 0.0)); } - - // Map massive_in/out1/out2 to Fortran get_clustering_type itypes: if (!massive_in && massive_out1 && !massive_out2) { - return sqrt(fabs(ldot(momentum2, sum))) / 2.0; + return sqrt(fabs(ldot(momentum2, momentum_sum))) / 2.0; } if (!massive_in && !massive_out1 && massive_out2) { - return sqrt(fabs(ldot(momentum1, sum))) / 2.0; + return sqrt(fabs(ldot(momentum1, momentum_sum))) / 2.0; } if (massive_in && !massive_out1 && !massive_out2) { - return sqrt(max(lsquare(sum), 0.0)); + 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. -// Uses boost and rotate_inverse from kinematics.hpp. -// -// i_remove : final-state particle absorbed into the combination (always >= 2, -// 0-indexed) j_keep : particle kept with the combined momentum (0/1 = initial state, -// else final) is_bw : BW resonance — combined mass = invariant mass of sum; else max -// of daughters -// -// Mass update rules mirror Fortran update_momenta (cluster.f:871): -// Final non-BW : mass[j_keep] = max(mass[j_keep], mass[i_remove]) -// Final BW : mass[j_keep] = sqrt((pi + pj)^2) -// Initial, one massive: mass[j_keep] = max(mass[j_keep], mass[i_remove]) -// Initial, same : mass[j_keep] = 0 +// Based on update_momenta from Template/NLO/SubProcesses/cluster.f template KERNELSPEC void update_momenta( int n_part, @@ -152,53 +133,49 @@ KERNELSPEC void update_momenta( FVal* masses, int& alive, int i_remove, - int j_keep, + int i_keep, bool resonant ) { alive &= ~(1 << i_remove); - if (j_keep < 2) { - // initial-state clustering - int j_other = 1 - j_keep; + if (i_keep < 2) { // initial-state clustering + int j_other = 1 - i_keep; for (int k = 0; k < 4; ++k) { - momenta[j_keep][k] -= momenta[i_remove][k]; + momenta[i_keep][k] -= momenta[i_remove][k]; } // mass: take max if exactly one daughter is massive, else 0 - masses[j_keep] = (masses[j_keep] > 0.0) != (masses[i_remove] > 0.0) - ? max(masses[j_keep], masses[i_remove]) - : FVal(0.0); + masses[i_keep] = (masses[i_keep] > 0.0) != (masses[i_remove] > 0.0) + ? max(masses[i_keep], masses[i_remove]) + : 0.0; - // CM boost vector: (E_tot, -px_tot, -py_tot, -pz_tot) of the two beam particles - FourMom pcmsp = { - momenta[j_keep][0] + momenta[j_other][0], - -(momenta[j_keep][1] + momenta[j_other][1]), - -(momenta[j_keep][2] + momenta[j_other][2]), - -(momenta[j_keep][3] + momenta[j_other][3]) + 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(pcmsp) > 100.0) { - // boost j_keep to CM to define the rotation axis, then apply to all alive - // particles - auto jkeep_cm = boost(momenta[j_keep], pcmsp, 1.0); + 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], pcmsp, 1.0), jkeep_cm); + momenta[j] = rotate_inverse( + boost(momenta[j], com_boost_vector, 1.0), jkeep_cm + ); } } } - - } else { - // final-state clustering + } else { // final-state clustering for (int k = 0; k < 4; ++k) { - momenta[j_keep][k] += momenta[i_remove][k]; + momenta[i_keep][k] += momenta[i_remove][k]; } if (resonant) { - masses[j_keep] = sqrt(max(lsquare(momenta[j_keep]), 0.0)); + masses[i_keep] = sqrt(max(lsquare(momenta[i_keep]), 0.0)); } else { - masses[j_keep] = max(masses[j_keep], masses[i_remove]); + masses[i_keep] = max(masses[i_keep], masses[i_remove]); } } } @@ -255,15 +232,16 @@ KERNELSPEC void mlm_clustering( bool is_last = (data >> 28) & 1; bool is_initial = (particle2 < 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) { - FourMom mom_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], - }; - FVal prop_m2 = lsquare(mom_sum); + 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; @@ -274,6 +252,7 @@ KERNELSPEC void mlm_clustering( FVal scale = compute_scale( momenta_tmp[particle1], momenta_tmp[particle2], + momentum_sum, masses_tmp[particle1], masses_tmp[particle2], resonant, From 3661af9f8e5e668c1b2bcb3ddbdfa13c5b834dfb Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 29 Jun 2026 17:03:29 +0200 Subject: [PATCH 07/12] output clustering scale for external particles --- madspace/src/kernels/mlm.hpp | 27 ++++++++++++++++++---- madspace/src/phasespace/mlm_clustering.cpp | 5 +++- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/madspace/src/kernels/mlm.hpp b/madspace/src/kernels/mlm.hpp index 5f0bb4f262..290efa04cb 100644 --- a/madspace/src/kernels/mlm.hpp +++ b/madspace/src/kernels/mlm.hpp @@ -144,7 +144,6 @@ KERNELSPEC void update_momenta( momenta[i_keep][k] -= momenta[i_remove][k]; } - // mass: take max if exactly one daughter is massive, else 0 masses[i_keep] = (masses[i_keep] > 0.0) != (masses[i_remove] > 0.0) ? max(masses[i_keep], masses[i_remove]) : 0.0; @@ -229,7 +228,7 @@ KERNELSPEC void mlm_clustering( bool massive_in = (data >> 24) & 1; bool massive_out1 = (data >> 25) & 1; bool massive_out2 = (data >> 26) & 1; - bool is_last = (data >> 28) & 1; + bool is_last = (data >> 30) & 1; bool is_initial = (particle2 < 2); FourMom momentum_sum{ @@ -305,11 +304,31 @@ KERNELSPEC void mlm_clustering( } } + 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]; - bool is_qcd = (cluster_history[i] >> 27) & 1; - ren_scale_val *= is_qcd ? scale : max_scale; + 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) { diff --git a/madspace/src/phasespace/mlm_clustering.cpp b/madspace/src/phasespace/mlm_clustering.cpp index ec5d2b78d6..9c7753e094 100644 --- a/madspace/src/phasespace/mlm_clustering.cpp +++ b/madspace/src/phasespace/mlm_clustering.cpp @@ -17,6 +17,8 @@ struct StateItem { bool massive_out1; bool massive_out2; bool is_qcd; + bool is_jet1; + bool is_jet2; }; void find_clusterings( @@ -209,7 +211,8 @@ MLMClustering::MLMClustering( (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 == &state.back()) << 28) + (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)); } From b075f6d475e1bfec27155dccbcc71cb83537f45b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 3 Jul 2026 15:23:08 +0200 Subject: [PATCH 08/12] mlm compiling and running, no reasonable output yet --- madspace/include/madspace/phasespace.hpp | 1 + .../madspace/phasespace/mlm_clustering.hpp | 11 +++- madspace/src/phasespace/mlm_clustering.cpp | 64 +++++++++++++++---- madspace/src/python/madspace.cpp | 17 +++++ madspace/tests/test_mlm.py | 60 +++++++++++++++++ 5 files changed, 141 insertions(+), 12 deletions(-) create mode 100644 madspace/tests/test_mlm.py diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index ca6cb3cf8d..ef927428a9 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 index a21377a387..aa7ce71632 100644 --- a/madspace/include/madspace/phasespace/mlm_clustering.hpp +++ b/madspace/include/madspace/phasespace/mlm_clustering.hpp @@ -10,7 +10,10 @@ class MLMClustering : public FunctionGenerator { MLMClustering( std::vector topologies, nested_vector3 permutations, - nested_vector2 diagram_indices + nested_vector2 diagram_indices, + double bw_cutoff = 15, + double jet_radius = 0.4, + bool hadronic = true ); private: @@ -19,6 +22,12 @@ class MLMClustering : public FunctionGenerator { ) 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/src/phasespace/mlm_clustering.cpp b/madspace/src/phasespace/mlm_clustering.cpp index 9c7753e094..60a4f2fede 100644 --- a/madspace/src/phasespace/mlm_clustering.cpp +++ b/madspace/src/phasespace/mlm_clustering.cpp @@ -64,7 +64,7 @@ void find_clusterings( key = {new_masks, new_diags}; } int index; - if (auto search = state_map.find(key); search == state_map.end()) { + if (auto search = state_map.find(key); search != state_map.end()) { index = search->second; } else { index = states.size(); @@ -80,7 +80,9 @@ void find_clusterings( .massive_in = false, .massive_out1 = false, .massive_out2 = false, - .is_qcd = false, + .is_qcd = true, + .is_jet1 = true, + .is_jet2 = true, }); if (new_masks.size() == 3) { @@ -95,6 +97,8 @@ void find_clusterings( .massive_out1 = false, .massive_out2 = false, .is_qcd = false, + .is_jet1 = false, + .is_jet2 = false, }); } } else { @@ -111,7 +115,10 @@ void find_clusterings( MLMClustering::MLMClustering( std::vector topologies, nested_vector3 permutations, - nested_vector2 diagram_indices + nested_vector2 diagram_indices, + double bw_cutoff, + double jet_radius, + bool hadronic ) : FunctionGenerator( "MLMClustering", @@ -120,16 +127,27 @@ MLMClustering::MLMClustering( {{"ren_scale", batch_float}, {"fact_scale1", batch_float}, {"fact_scale2", batch_float}, - {"cluster_history", - batch_int_array(topologies.at(0).outgoing_masses().size() - 1)}, - {"cluster_scales", - batch_float_array(topologies.at(0).outgoing_masses().size() - 1)}} - ) { - std::size_t n_ext = topologies.at(0).outgoing_masses().size() + 2; + {"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] : @@ -138,7 +156,7 @@ MLMClustering::MLMClustering( 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))) = 1 + particle_masks.at(topo.outgoing_indices().at(permutation.at(i) - 2)) = 1 << i; } @@ -223,6 +241,30 @@ MLMClustering::MLMClustering( NamedVector MLMClustering::build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const { - auto mlm_out = fb.mlm_clustering_hadronic(args.at(0), _cluster_state_machine); + 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/python/madspace.cpp b/madspace/src/python/madspace.cpp index 3d3313c8d6..181a51520b 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -952,6 +952,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", diff --git a/madspace/tests/test_mlm.py b/madspace/tests/test_mlm.py new file mode 100644 index 0000000000..8a8bc969e1 --- /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) From bed30160262699966d15ffa056737d9fc9233ee6 Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Fri, 3 Jul 2026 16:25:34 +0200 Subject: [PATCH 09/12] mlm scale computation bugfixes --- madspace/src/kernels/mlm.hpp | 13 +++++----- madspace/src/phasespace/mlm_clustering.cpp | 28 ++++++++++++---------- 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/madspace/src/kernels/mlm.hpp b/madspace/src/kernels/mlm.hpp index 290efa04cb..1e30b63e1b 100644 --- a/madspace/src/kernels/mlm.hpp +++ b/madspace/src/kernels/mlm.hpp @@ -12,7 +12,7 @@ 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(FourMom p, bool hadronic) { +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); } @@ -24,8 +24,8 @@ KERNELSPEC FVal djb_clus(FourMom p, bool hadronic) { // based on dj_clus from Template/NLO/SubProcesses/cluster.f template KERNELSPEC FVal dj_clus( - FourMom p1, - FourMom p2, + const FourMom& p1, + const FourMom& p2, FVal mass1, FVal mass2, bool hadronic, @@ -99,11 +99,10 @@ KERNELSPEC FVal compute_scale( 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(momentum1, hadronic)); + auto scale = sqrt(djb_clus(momentum2, hadronic)); if ((momentum1[3] < 0.0) != (momentum2[3] < 0.0)) { scale = scale * ONE_PLUS_TINY; } @@ -229,7 +228,7 @@ KERNELSPEC void mlm_clustering( bool massive_out1 = (data >> 25) & 1; bool massive_out2 = (data >> 26) & 1; bool is_last = (data >> 30) & 1; - bool is_initial = (particle2 < 2); + bool is_initial = (particle1 < 2); FourMom momentum_sum{ momenta_tmp[particle1][0] + momenta_tmp[particle2][0], @@ -285,7 +284,7 @@ KERNELSPEC void mlm_clustering( ++cluster_count; win_scale = 1e308; } else { - ++state; + state += 2; } } diff --git a/madspace/src/phasespace/mlm_clustering.cpp b/madspace/src/phasespace/mlm_clustering.cpp index 60a4f2fede..42e96f63e0 100644 --- a/madspace/src/phasespace/mlm_clustering.cpp +++ b/madspace/src/phasespace/mlm_clustering.cpp @@ -87,19 +87,21 @@ void find_clusterings( if (new_masks.size() == 3) { auto& current_states = states.at(index); - 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, - }); + 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( From f5060f1934010016f31e35f0321e3fc73bc679ca Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 6 Jul 2026 15:06:59 +0200 Subject: [PATCH 10/12] add mlm as option for dynamical scale --- .../iolibs/template_files/mg7/madevent.py | 39 +++++++++++++------ .../iolibs/template_files/mg7/run_card.toml | 3 +- .../include/madspace/phasespace/scale.hpp | 3 ++ madspace/src/phasespace/scale.cpp | 7 ++++ madspace/src/python/madspace.cpp | 3 +- 5 files changed, 41 insertions(+), 14 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 57c6f73c91..33039f36d2 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 diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 83af645394..8a94ca1e57 100644 --- a/madgraph/iolibs/template_files/mg7/run_card.toml +++ b/madgraph/iolibs/template_files/mg7/run_card.toml @@ -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/include/madspace/phasespace/scale.hpp b/madspace/include/madspace/phasespace/scale.hpp index b53235821b..d55bc8efaa 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,7 @@ class EnergyScale : public FunctionGenerator { double fact_scale1, double fact_scale2 ); + EnergyScale(const MLMClustering& clustering); private: NamedVector build_function_impl( @@ -48,6 +50,7 @@ class EnergyScale : public FunctionGenerator { double _ren_scale; double _fact_scale1; double _fact_scale2; + std::optional _clustering; }; } // namespace madspace diff --git a/madspace/src/phasespace/scale.cpp b/madspace/src/phasespace/scale.cpp index e6106314d5..d39d71bcd7 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/madspace.cpp b/madspace/src/python/madspace.cpp index 64ec9baf05..07ddde36b9 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1185,7 +1185,8 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("ren_scale"), py::arg("fact_scale1"), py::arg("fact_scale2") - ); + ) + .def(py::init(), py::arg("clustering")); py::classh(m, "CachedPdf").def(py::init<>()); py::classh(m, "CachedScale") From bde606ddb8f3f6e4ed2f7c5923d9b78eefcdc50d Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 6 Jul 2026 15:51:49 +0200 Subject: [PATCH 11/12] cluster scales in lhe output --- madgraph/iolibs/template_files/mg7/run_card.toml | 2 +- madspace/include/madspace/driver/lhe_output.hpp | 2 ++ madspace/include/madspace/phasespace/scale.hpp | 2 ++ madspace/src/driver/lhe_output.cpp | 12 ++++++++++++ 4 files changed, 17 insertions(+), 1 deletion(-) diff --git a/madgraph/iolibs/template_files/mg7/run_card.toml b/madgraph/iolibs/template_files/mg7/run_card.toml index 8a94ca1e57..47418315b8 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 diff --git a/madspace/include/madspace/driver/lhe_output.hpp b/madspace/include/madspace/driver/lhe_output.hpp index 5451d6f257..f29c612f70 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/scale.hpp b/madspace/include/madspace/phasespace/scale.hpp index d55bc8efaa..f95cfbed38 100644 --- a/madspace/include/madspace/phasespace/scale.hpp +++ b/madspace/include/madspace/phasespace/scale.hpp @@ -39,6 +39,8 @@ class EnergyScale : public FunctionGenerator { ); EnergyScale(const MLMClustering& clustering); + bool is_mlm() const { return _clustering.has_value(); } + private: NamedVector build_function_impl( FunctionBuilder& fb, const NamedVector& args diff --git a/madspace/src/driver/lhe_output.cpp b/madspace/src/driver/lhe_output.cpp index edea68b065..33ee1752a3 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, ""); + } buffer += "\n"; } From a0a006985fd8ecf4efe3af297602a3a4266e418b Mon Sep 17 00:00:00 2001 From: Theo Heimel Date: Mon, 6 Jul 2026 17:07:55 +0200 Subject: [PATCH 12/12] pass mlm output to event file and matrix element --- .../iolibs/template_files/mg7/madevent.py | 13 +++++++-- .../madspace/driver/channel_generator.hpp | 1 + madspace/include/madspace/driver/io.hpp | 5 +++- madspace/src/driver/channel_generator.cpp | 17 +++++++++++ madspace/src/driver/event_generator.cpp | 22 ++++++++++++-- madspace/src/driver/lhe_output.cpp | 16 +++++----- madspace/src/phasespace/integrand.cpp | 29 +++++++++++++++++++ madspace/src/python/madspace.cpp | 3 +- 8 files changed, 92 insertions(+), 14 deletions(-) diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 33039f36d2..e52957daad 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -1306,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, ) @@ -1317,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/madspace/include/madspace/driver/channel_generator.hpp b/madspace/include/madspace/driver/channel_generator.hpp index 957037d158..dae774f2dd 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 783de70da5..1ab7f211d7 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/src/driver/channel_generator.cpp b/madspace/src/driver/channel_generator.cpp index 823fb5e860..d83797b9b0 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 af8e3a59ba..3a078d3ea0 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 33ee1752a3..03ada99aac 100644 --- a/madspace/src/driver/lhe_output.cpp +++ b/madspace/src/driver/lhe_output.cpp @@ -55,15 +55,15 @@ void LHEEvent::format_to(std::string& buffer) const { } if (has_cluster_scales) { std::format_to(insert_iter, ""); + std::format_to(insert_iter, ">\n"); } buffer += "\n"; } diff --git a/madspace/src/phasespace/integrand.cpp b/madspace/src/phasespace/integrand.cpp index 2d808ee34a..0ab7c670a6 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/python/madspace.cpp b/madspace/src/python/madspace.cpp index 07ddde36b9..4266923a58 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -1186,7 +1186,8 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("fact_scale1"), py::arg("fact_scale2") ) - .def(py::init(), py::arg("clustering")); + .def(py::init(), py::arg("clustering")) + .def("is_mlm", &EnergyScale::is_mlm); py::classh(m, "CachedPdf").def(py::init<>()); py::classh(m, "CachedScale")