diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc index 592b031dc..5b60df2cd 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/madspace/include/madspace/compgraphs/function_builder_mixin.inc @@ -300,23 +300,33 @@ std::array t_inv_value_and_min_max(Value pa, Value pb, Value p1, Value return {output_vector[0], output_vector[1], output_vector[2]}; } -std::array t1_inv_min_max_doublet(Value pa, Value pb, Value m1, Value mir_min) { - auto output_vector = instruction("t1_inv_min_max_doublet", {pa, pb, m1, mir_min}); +std::array t_inv_min_max_cut(Value pa, Value pb, Value m1, Value m2, Value etmin_1, Value etmin_2) { + auto output_vector = instruction("t_inv_min_max_cut", {pa, pb, m1, m2, etmin_1, etmin_2}); return {output_vector[0], output_vector[1]}; } -std::array t1_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value m1, Value mir_min) { - auto output_vector = instruction("t1_inv_value_and_min_max_doublet", {pa, pb, p1, m1, mir_min}); +std::array t_inv_value_and_min_max_cut(Value pa, Value pb, Value p1, Value p2, Value etmin_1, Value etmin_2) { + auto output_vector = instruction("t_inv_value_and_min_max_cut", {pa, pb, p1, p2, etmin_1, etmin_2}); return {output_vector[0], output_vector[1], output_vector[2]}; } -std::array t2_inv_min_max_doublet(Value pa, Value pb, Value m1, Value mir_min, Value t1_abs) { - auto output_vector = instruction("t2_inv_min_max_doublet", {pa, pb, m1, mir_min, t1_abs}); +std::array t1_inv_min_max_doublet(Value pa, Value pb, Value m1, Value mir_min, Value etmin_i, Value etmin_ir) { + auto output_vector = instruction("t1_inv_min_max_doublet", {pa, pb, m1, mir_min, etmin_i, etmin_ir}); return {output_vector[0], output_vector[1]}; } -std::array t2_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value m1, Value mir_min, Value t1_abs) { - auto output_vector = instruction("t2_inv_value_and_min_max_doublet", {pa, pb, p1, m1, mir_min, t1_abs}); +std::array t1_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value m1, Value mir_min, Value etmin_i, Value etmin_ir) { + auto output_vector = instruction("t1_inv_value_and_min_max_doublet", {pa, pb, p1, m1, mir_min, etmin_i, etmin_ir}); + return {output_vector[0], output_vector[1], output_vector[2]}; +} + +std::array t2_inv_min_max_doublet(Value pa, Value pb, Value m1, Value mir_min, Value t1_abs, Value etmin_i, Value etmin_ir) { + auto output_vector = instruction("t2_inv_min_max_doublet", {pa, pb, m1, mir_min, t1_abs, etmin_i, etmin_ir}); + return {output_vector[0], output_vector[1]}; +} + +std::array t2_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value m1, Value mir_min, Value t1_abs, Value etmin_i, Value etmin_ir) { + auto output_vector = instruction("t2_inv_value_and_min_max_doublet", {pa, pb, p1, m1, mir_min, t1_abs, etmin_i, etmin_ir}); return {output_vector[0], output_vector[1], output_vector[2]}; } @@ -330,6 +340,16 @@ std::array s23_value_and_min_max(Value pa, Value pb, Value p3, Value t return {output_vector[0], output_vector[1], output_vector[2]}; } +std::array s23_min_max_cut(Value pa, Value pb, Value p3, Value t1_abs, Value m1, Value m2, Value etmin_1, Value etmin_2, Value drcut, Value s23_min_cut) { + auto output_vector = instruction("s23_min_max_cut", {pa, pb, p3, t1_abs, m1, m2, etmin_1, etmin_2, drcut, s23_min_cut}); + return {output_vector[0], output_vector[1]}; +} + +std::array s23_value_and_min_max_cut(Value pa, Value pb, Value p3, Value t1_abs, Value p1, Value p2, Value etmin_1, Value etmin_2, Value drcut, Value s23_min_cut) { + auto output_vector = instruction("s23_value_and_min_max_cut", {pa, pb, p3, t1_abs, p1, p2, etmin_1, etmin_2, drcut, s23_min_cut}); + return {output_vector[0], output_vector[1], output_vector[2]}; +} + Value invariants_from_momenta(Value p_ext, Value factors) { return instruction("invariants_from_momenta", {p_ext, factors})[0]; } diff --git a/madspace/include/madspace/compgraphs/opcode_mixin.inc b/madspace/include/madspace/compgraphs/opcode_mixin.inc index 079c51b9d..4898175ad 100644 --- a/madspace/include/madspace/compgraphs/opcode_mixin.inc +++ b/madspace/include/madspace/compgraphs/opcode_mixin.inc @@ -67,86 +67,90 @@ three_body_decay = 65, three_body_decay_inverse = 66, t_inv_min_max = 67, t_inv_value_and_min_max = 68, -t1_inv_min_max_doublet = 69, -t1_inv_value_and_min_max_doublet = 70, -t2_inv_min_max_doublet = 71, -t2_inv_value_and_min_max_doublet = 72, -s23_min_max = 73, -s23_value_and_min_max = 74, -invariants_from_momenta = 75, -sde2_channel_weights = 76, -subchannel_weights = 77, -apply_subchannel_weights = 78, -pt_eta_phi_x = 79, -mirror_momenta = 80, -momenta_to_x1x2 = 81, -uniform_invariant = 82, -uniform_invariant_inverse = 83, -breit_wigner_invariant = 84, -breit_wigner_invariant_inverse = 85, -stable_invariant = 86, -stable_invariant_inverse = 87, -stable_invariant_nu = 88, -stable_invariant_nu_inverse = 89, -fast_rambo_massless = 90, -fast_rambo_massless_inverse = 91, -fast_rambo_massless_com = 92, -fast_rambo_massive = 93, -fast_rambo_massive_inverse = 94, -fast_rambo_massive_com = 95, -cut_unphysical = 96, -cut_one = 97, -cut_all = 98, -cut_any = 99, -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 +t_inv_min_max_cut = 69, +t_inv_value_and_min_max_cut = 70, +t1_inv_min_max_doublet = 71, +t1_inv_value_and_min_max_doublet = 72, +t2_inv_min_max_doublet = 73, +t2_inv_value_and_min_max_doublet = 74, +s23_min_max = 75, +s23_value_and_min_max = 76, +s23_min_max_cut = 77, +s23_value_and_min_max_cut = 78, +invariants_from_momenta = 79, +sde2_channel_weights = 80, +subchannel_weights = 81, +apply_subchannel_weights = 82, +pt_eta_phi_x = 83, +mirror_momenta = 84, +momenta_to_x1x2 = 85, +uniform_invariant = 86, +uniform_invariant_inverse = 87, +breit_wigner_invariant = 88, +breit_wigner_invariant_inverse = 89, +stable_invariant = 90, +stable_invariant_inverse = 91, +stable_invariant_nu = 92, +stable_invariant_nu_inverse = 93, +fast_rambo_massless = 94, +fast_rambo_massless_inverse = 95, +fast_rambo_massless_com = 96, +fast_rambo_massive = 97, +fast_rambo_massive_inverse = 98, +fast_rambo_massive_com = 99, +cut_unphysical = 100, +cut_one = 101, +cut_all = 102, +cut_any = 103, +scale_transverse_energy = 104, +scale_transverse_mass = 105, +scale_half_transverse_mass = 106, +scale_partonic_energy = 107, +chili_forward = 108, +chili_inverse = 109, +matrix_element = 110, +collect_channel_weights = 111, +interpolate_pdf = 112, +interpolate_alpha_s = 113, +matmul = 114, +relu = 115, +leaky_relu = 116, +elu = 117, +gelu = 118, +sigmoid = 119, +softplus = 120, +rqs_reshape = 121, +rqs_find_bin = 122, +rqs_forward = 123, +rqs_inverse = 124, +softmax = 125, +softmax_prior = 126, +sample_discrete = 127, +sample_discrete_inverse = 128, +sample_discrete_probs = 129, +sample_discrete_probs_inverse = 130, +discrete_histogram = 131, +permute_momenta = 132, +gather = 133, +gather_int = 134, +select_int = 135, +select = 136, +select_vector = 137, +argsort = 138, +quantile = 139, +one_hot = 140, +madnis_abs_weight = 141, +madnis_softclip = 142, +madnis_variance = 143, +madnis_single_channel_variance = 144, +madnis_multi_channel_variance = 145, +nonzero = 146, +batch_gather = 147, +batch_scatter = 148, +random = 149, +random_int = 150, +unweight = 151, +vegas_forward = 152, +vegas_inverse = 153, +vegas_histogram = 154, +histogram = 155 diff --git a/madspace/include/madspace/phasespace/base.hpp b/madspace/include/madspace/phasespace/base.hpp index 618efc621..f5ca88c6c 100644 --- a/madspace/include/madspace/phasespace/base.hpp +++ b/madspace/include/madspace/phasespace/base.hpp @@ -50,6 +50,7 @@ class Mapping { const NamedVector& output_types() const { return _output_types; } const NamedVector& condition_types() const { return _condition_types; } const std::string& name() const { return _name; } + virtual std::size_t discrete_dim() const { return 0; } protected: // TODO: make parameters const ref diff --git a/madspace/include/madspace/phasespace/color_ordered_mapping.hpp b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp index c6d339743..2c05dc9cf 100644 --- a/madspace/include/madspace/phasespace/color_ordered_mapping.hpp +++ b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp @@ -14,13 +14,31 @@ class ColorOrderedMapping : public Mapping { public: // color_order: 0-indexed permutation of {0, ..., n-1} (n = n_out + 2). // Particles 0 and 1 are the two incoming beams. + // + // Optional cuts (all indexed by 0-based outgoing-particle index, i.e. the + // same indexing as the masses m_out and the entries of the color order + // minus 2): + // * pt_min[i] : minimum transverse momentum of outgoing particle i. + // * m_inv_min[i][j] : minimum invariant mass of the pair (i, j). + // * dr_min[i][j] : minimum delta-R separation of the pair (i, j). + // Passing any non-empty cut container enables cut-aware sampling. The cuts + // are translated into invariant-space bounds exactly as in the Fortran + // reference (phase_space_gen23): per-subset invariant-mass floors + // (invm_min) on the sampled s-channel masses, the adjacent-pair floor on + // the 2->3 s23 invariant, and a |t| floor (pt^2) on each peeled particle. + // Empty containers (the default) reproduce the previous cut-free behaviour + // exactly. ColorOrderedMapping( const std::vector& color_order, double t_invariant_power = 0.8, - double s_invariant_power = 0.8 + double s_invariant_power = 0.8, + const std::vector& pt_min = {}, + const std::vector>& m_inv_min = {}, + const std::vector>& dr_min = {} ); std::size_t random_dim() const { return _random_dim; } + std::size_t discrete_dim() const override { return _discrete_dim; } private: Result build_forward_impl( @@ -34,6 +52,13 @@ class ColorOrderedMapping : public Mapping { const NamedVector& conditions ) const override; + // pt^2 of outgoing particle i (0 if no pt cut on it). + double pt2(std::size_t i) const; + // Cut-derived invariant-mass^2 floor (gen23 invm_min, without the mass^2 + // term, which is applied separately) for a subset of outgoing particles. + // Returns 0 when cuts are disabled or the subset has fewer than 2 members. + double cut_floor(const std::vector& subset) const; + // 0-indexed outgoing-particle indices (values in {0,...,n_out-1}). // _set1 contains the outgoing particles attached to beam 0's side, // _set2 those attached to beam 1's side, in peel order. @@ -41,6 +66,10 @@ class ColorOrderedMapping : public Mapping { std::vector _set2; std::size_t _n_out; std::size_t _random_dim; + // Number of discrete two-solution choices (one per 2->3 peel). These are + // supplied/recovered as a separate batch_int channel, not through the + // continuous random_dim() block (opt-in r_disc). + std::size_t _discrete_dim; // True iff exactly one of (set1, set2) has size 1 (and the other >= 2). // In that case the central block is DoubleT instead of 2->2. bool _use_double_t; @@ -50,6 +79,12 @@ class ColorOrderedMapping : public Mapping { // produced as a single t-channel chain seeded directly off the beams. bool _use_single_chain; + // Cut configuration (empty => all bounds resolve to 0 = no cut). + std::vector _pt_min; + std::vector> _m_inv_min; + std::vector> _dr_min; + bool _has_cut; + Invariant _uniform_invariant; TwoToTwoParticleScattering _com_scattering; TwoToTwoParticleScattering _lab_scattering; diff --git a/madspace/include/madspace/phasespace/cuts.hpp b/madspace/include/madspace/phasespace/cuts.hpp index e0de627d7..7d8347fc7 100644 --- a/madspace/include/madspace/phasespace/cuts.hpp +++ b/madspace/include/madspace/phasespace/cuts.hpp @@ -4,6 +4,8 @@ #include "madspace/phasespace/base.hpp" #include "madspace/phasespace/observable.hpp" +#include +#include #include namespace madspace { @@ -23,12 +25,20 @@ class Cuts : public FunctionGenerator { double sqrt_s_min() const; std::vector eta_max() const; std::vector pt_min() const; + std::vector> m_inv_min() const; + std::vector> dr_min() const; private: NamedVector build_function_impl( FunctionBuilder& fb, const NamedVector& args ) const override; + std::vector> pairwise_min( + Observable::ObservableOption obs, + const std::function< + std::vector>(const Observable&)>& pairs + ) const; + std::vector _cut_data; }; diff --git a/madspace/include/madspace/phasespace/observable.hpp b/madspace/include/madspace/phasespace/observable.hpp index 909f255a6..093db8d9f 100644 --- a/madspace/include/madspace/phasespace/observable.hpp +++ b/madspace/include/madspace/phasespace/observable.hpp @@ -48,6 +48,8 @@ class Observable : public FunctionGenerator { const std::string& name = "" ); ObservableOption observable() const { return _observable; } + const nested_vector2& indices() const { return _indices; } + bool sum_momenta() const { return _sum_momenta; } std::vector simple_observable_indices() const { if (_sum_momenta || _sum_observable || _indices.size() != 1) { return {}; diff --git a/madspace/include/madspace/phasespace/phasespace.hpp b/madspace/include/madspace/phasespace/phasespace.hpp index 617f39fa6..b1827565d 100644 --- a/madspace/include/madspace/phasespace/phasespace.hpp +++ b/madspace/include/madspace/phasespace/phasespace.hpp @@ -2,6 +2,7 @@ #include "madspace/phasespace/base.hpp" #include "madspace/phasespace/chili.hpp" +#include "madspace/phasespace/color_ordered_mapping.hpp" #include "madspace/phasespace/cuts.hpp" #include "madspace/phasespace/invariants.hpp" #include "madspace/phasespace/luminosity.hpp" @@ -14,7 +15,7 @@ namespace madspace { class PhaseSpaceMapping : public Mapping { public: - enum TChannelMode { propagator, rambo, chili }; + enum TChannelMode { propagator, rambo, chili, color_ordered }; PhaseSpaceMapping( const Topology& topology, @@ -23,7 +24,8 @@ class PhaseSpaceMapping : public Mapping { double invariant_power = 0.8, TChannelMode t_channel_mode = propagator, const std::optional& cuts = std::nullopt, - const std::vector>& permutations = {} + const std::vector>& permutations = {}, + const std::optional>& color_order = std::nullopt ); PhaseSpaceMapping( @@ -32,12 +34,14 @@ class PhaseSpaceMapping : public Mapping { bool leptonic = false, double invariant_power = 0.8, TChannelMode mode = rambo, - const std::optional& cuts = std::nullopt + const std::optional& cuts = std::nullopt, + const std::optional>& color_order = std::nullopt ); std::size_t random_dim() const { return 3 * _topology.outgoing_masses().size() - (_leptonic ? 4 : 2); } + std::size_t discrete_dim() const override { return _n_discrete; } std::size_t particle_count() const { return _topology.outgoing_masses().size() + 2; } @@ -61,8 +65,14 @@ class PhaseSpaceMapping : public Mapping { double _sqrt_s_lab; bool _leptonic; bool _map_luminosity; + std::size_t _n_discrete; std::vector _s_invariants; - std::variant + std::variant< + TPropagatorMapping, + FastRamboMapping, + ChiliMapping, + ColorOrderedMapping, + std::monostate> _t_mapping; std::vector> _s_decays; nested_vector2 _permutations; diff --git a/madspace/include/madspace/phasespace/t_propagator_mapping.hpp b/madspace/include/madspace/phasespace/t_propagator_mapping.hpp index 162de6183..81a97423d 100644 --- a/madspace/include/madspace/phasespace/t_propagator_mapping.hpp +++ b/madspace/include/madspace/phasespace/t_propagator_mapping.hpp @@ -12,7 +12,9 @@ namespace madspace { class TPropagatorMapping : public Mapping { public: TPropagatorMapping( - const std::vector& integration_order, double invariant_power = 0.8 + const std::vector& integration_order, + double invariant_power = 0.8, + const std::vector& pt_min = {} ); std::size_t random_dim() const { return 3 * _integration_order.size() - 1; } @@ -28,8 +30,13 @@ class TPropagatorMapping : public Mapping { const NamedVector& conditions ) const override; + // pt^2 of the outgoing particle at position i (0 if no pt cut). + double pt2(std::size_t i) const; + std::vector _integration_order; std::vector _sample_sides; + std::vector _pt_min; + bool _has_cut; Invariant _uniform_invariant; TwoToTwoParticleScattering _com_scattering; TwoToTwoParticleScattering _lab_scattering; diff --git a/madspace/include/madspace/phasespace/three_particle.hpp b/madspace/include/madspace/phasespace/three_particle.hpp index 64e1485c9..e7d3a7700 100644 --- a/madspace/include/madspace/phasespace/three_particle.hpp +++ b/madspace/include/madspace/phasespace/three_particle.hpp @@ -33,9 +33,13 @@ class TwoToThreeParticleScattering : public Mapping { double t_width = 0, double s_invariant_power = 0, double s_mass = 0, - double s_width = 0 + double s_width = 0, + bool has_cut = false ); + // one discrete input: which of the 2 two-body solutions to take + std::size_t discrete_dim() const override { return 1; } + private: Result build_forward_impl( FunctionBuilder& fb, @@ -50,6 +54,7 @@ class TwoToThreeParticleScattering : public Mapping { Invariant _t_invariant; Invariant _s_invariant; + bool _has_cut; }; } // namespace madspace diff --git a/madspace/include/madspace/phasespace/two_particle.hpp b/madspace/include/madspace/phasespace/two_particle.hpp index 5faf94994..238c6b135 100644 --- a/madspace/include/madspace/phasespace/two_particle.hpp +++ b/madspace/include/madspace/phasespace/two_particle.hpp @@ -28,7 +28,11 @@ class TwoBodyDecay : public Mapping { class TwoToTwoParticleScattering : public Mapping { public: TwoToTwoParticleScattering( - bool com, double invariant_power = 0, double mass = 0, double width = 0 + bool com, + double invariant_power = 0, + double mass = 0, + double width = 0, + bool has_cut = false ); private: @@ -45,6 +49,7 @@ class TwoToTwoParticleScattering : public Mapping { bool _com; Invariant _invariant; + bool _has_cut; }; class DoubleT : public Mapping { @@ -55,7 +60,8 @@ class DoubleT : public Mapping { double t1_width = 0, double t2_invariant_power = 0, double t2_mass = 0, - double t2_width = 0 + double t2_width = 0, + bool has_cut = false ); private: @@ -72,6 +78,7 @@ class DoubleT : public Mapping { Invariant _t1_invariant; Invariant _t2_invariant; + bool _has_cut; }; } // namespace madspace diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 571fd9dd1..8e10d14c4 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -1242,6 +1242,65 @@ t_inv_value_and_min_max: type: [float] desc: +t_inv_min_max_cut: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: m1 + type: [float] + desc: + - name: m2 + type: [float] + desc: + - name: etmin_1 + type: [float] + desc: + - name: etmin_2 + type: [float] + desc: + outputs: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + +t_inv_value_and_min_max_cut: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: p1 + type: [float, 4] + desc: + - name: p2 + type: [float, 4] + desc: + - name: etmin_1 + type: [float] + desc: + - name: etmin_2 + type: [float] + desc: + outputs: + - name: t_abs + type: [float] + desc: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + t1_inv_min_max_doublet: inputs: - name: pa @@ -1256,6 +1315,12 @@ t1_inv_min_max_doublet: - name: mir_min type: [float] desc: + - name: etmin_i + type: [float] + desc: + - name: etmin_ir + type: [float] + desc: outputs: - name: t_min type: [float] @@ -1281,6 +1346,12 @@ t1_inv_value_and_min_max_doublet: - name: mir_min type: [float] desc: + - name: etmin_i + type: [float] + desc: + - name: etmin_ir + type: [float] + desc: outputs: - name: t_abs type: [float] @@ -1309,6 +1380,12 @@ t2_inv_min_max_doublet: - name: t1_abs type: [float] desc: + - name: etmin_i + type: [float] + desc: + - name: etmin_ir + type: [float] + desc: outputs: - name: t_min type: [float] @@ -1337,6 +1414,12 @@ t2_inv_value_and_min_max_doublet: - name: t1_abs type: [float] desc: + - name: etmin_i + type: [float] + desc: + - name: etmin_ir + type: [float] + desc: outputs: - name: t_abs type: [float] @@ -1407,6 +1490,89 @@ s23_value_and_min_max: type: [float] desc: +s23_min_max_cut: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: p3 + type: [float, 4] + desc: + - name: t1_abs + type: [float] + desc: + - name: m1 + type: [float] + desc: + - name: m2 + type: [float] + desc: + - name: etmin_1 + type: [float] + desc: + - name: etmin_2 + type: [float] + desc: + - name: drcut + type: [float] + desc: + - name: s23_min_cut + type: [float] + desc: + outputs: + - name: s23_min + type: [float] + desc: + - name: s23_max + type: [float] + desc: + +s23_value_and_min_max_cut: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: p3 + type: [float, 4] + desc: + - name: t1_abs + type: [float] + desc: + - name: p1 + type: [float, 4] + desc: + - name: p2 + type: [float, 4] + desc: + - name: etmin_1 + type: [float] + desc: + - name: etmin_2 + type: [float] + desc: + - name: drcut + type: [float] + desc: + - name: s23_min_cut + type: [float] + desc: + outputs: + - name: s_23 + type: [float] + desc: + - name: s23_min + type: [float] + desc: + - name: s23_max + type: [float] + desc: + invariants_from_momenta: inputs: - name: p_ext diff --git a/madspace/src/compgraphs/instruction_set_mixin.inc b/madspace/src/compgraphs/instruction_set_mixin.inc index 30b23bab6..b8fa7b59e 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -84,87 +84,91 @@ InstructionOwner instructions[] { mi("three_body_decay_inverse", 66, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), mi("t_inv_min_max", 67, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), mi("t_inv_value_and_min_max", 68, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("t1_inv_min_max_doublet", 69, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("t1_inv_value_and_min_max_doublet", 70, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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("t2_inv_min_max_doublet", 71, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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("t2_inv_value_and_min_max_doublet", 72, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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}, {DataType::dt_float, false, {}, false}}), - mi("s23_min_max", 73, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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("s23_value_and_min_max", 74, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("invariants_from_momenta", 75, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {"m", "n"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("sde2_channel_weights", 76, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}}, {{DataType::dt_float, false, {"c"}, false}}), - mi("subchannel_weights", 77, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, true, {"g"}, false}}, {{DataType::dt_float, false, {"c"}, false}}), - mi("apply_subchannel_weights", 78, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"d"}, false}, {DataType::dt_int, false, {"d"}, false}}, {{DataType::dt_float, false, {"d"}, false}}), - mi("pt_eta_phi_x", 79, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n+2"}, false}}), - mi("mirror_momenta", 80, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), - mi("momenta_to_x1x2", 81, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("uniform_invariant", 82, 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("uniform_invariant_inverse", 83, 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("breit_wigner_invariant", 84, 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("breit_wigner_invariant_inverse", 85, 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant", 86, 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}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant_inverse", 87, 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}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant_nu", 88, 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant_nu_inverse", 89, 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massless", 90, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massless_inverse", 91, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massless_com", 92, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massive", 93, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massive_inverse", 94, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massive_com", 95, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), - mi("cut_unphysical", 96, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("cut_one", 97, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("cut_all", 98, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("cut_any", 99, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("scale_transverse_energy", 100, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - 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, false, {}, 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("t_inv_min_max_cut", 69, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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}, {DataType::dt_float, false, {}, false}}), + mi("t_inv_value_and_min_max_cut", 70, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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("t1_inv_min_max_doublet", 71, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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}, {DataType::dt_float, false, {}, false}}), + mi("t1_inv_value_and_min_max_doublet", 72, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("t2_inv_min_max_doublet", 73, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("t2_inv_value_and_min_max_doublet", 74, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("s23_min_max", 75, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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("s23_value_and_min_max", 76, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("s23_min_max_cut", 77, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("s23_value_and_min_max_cut", 78, true, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("invariants_from_momenta", 79, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {"m", "n"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("sde2_channel_weights", 80, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}}, {{DataType::dt_float, false, {"c"}, false}}), + mi("subchannel_weights", 81, true, {{DataType::dt_float, false, {"m"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_float, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, false, {"c", "n"}, false}, {DataType::dt_int, true, {"g"}, false}}, {{DataType::dt_float, false, {"c"}, false}}), + mi("apply_subchannel_weights", 82, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"d"}, false}, {DataType::dt_int, false, {"d"}, false}}, {{DataType::dt_float, false, {"d"}, false}}), + mi("pt_eta_phi_x", 83, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n+2"}, false}}), + mi("mirror_momenta", 84, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), + mi("momenta_to_x1x2", 85, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("uniform_invariant", 86, 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("uniform_invariant_inverse", 87, 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("breit_wigner_invariant", 88, 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("breit_wigner_invariant_inverse", 89, 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant", 90, 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}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant_inverse", 91, 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}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant_nu", 92, 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("stable_invariant_nu_inverse", 93, 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massless", 94, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massless_inverse", 95, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massless_com", 96, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massive", 97, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {4}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massive_inverse", 98, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("fast_rambo_massive_com", 99, true, {{DataType::dt_float, false, {"3n-4"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("cut_unphysical", 100, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("cut_one", 101, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("cut_all", 102, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("cut_any", 103, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_transverse_energy", 104, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_transverse_mass", 105, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_half_transverse_mass", 106, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("scale_partonic_energy", 107, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("chili_forward", 108, true, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}}), + mi("chili_inverse", 109, true, {{DataType::dt_float, false, {"n+2", 4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"3n-2"}, false}, {DataType::dt_float, false, {}, false}}), + InstructionOwner(new MatrixElementInstruction(110, true)), + mi("collect_channel_weights", 111, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}, {DataType::dt_int, true, {"c"}, true}}, {{DataType::dt_float, false, {"c"}, false}}), + mi("interpolate_pdf", 112, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {"n"}, false}, {DataType::dt_float, true, {"a"}, false}, {DataType::dt_float, true, {"b"}, false}, {DataType::dt_float, true, {16, "c", "d"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("interpolate_alpha_s", 113, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {"b+1"}, false}, {DataType::dt_float, true, {4, "b"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("matmul", 114, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"m", "n"}, false}, {DataType::dt_float, true, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("relu", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("leaky_relu", 116, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("elu", 117, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("gelu", 118, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("sigmoid", 119, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softplus", 120, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + InstructionOwner(new RqsReshapeInstruction(121, true)), + mi("rqs_find_bin", 122, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b"}, false}, {DataType::dt_float, false, {"n", "b+1"}, false}}, {{DataType::dt_float, false, {"n", 6}, false}}), + mi("rqs_forward", 123, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("rqs_inverse", 124, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n", 6}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("softmax", 125, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), + mi("softmax_prior", 126, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("sample_discrete", 127, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_inverse", 128, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_probs", 129, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_probs_inverse", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("discrete_histogram", 131, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, true, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}}), + mi("permute_momenta", 132, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, true, {"m", "n"}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), + mi("gather", 133, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("gather_int", 134, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), + mi("select_int", 135, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), + mi("select", 136, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_vector", 137, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), + mi("argsort", 138, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), + mi("quantile", 139, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, true, {}, false}}, {{DataType::dt_float, true, {}, false}}), + mi("one_hot", 140, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("madnis_abs_weight", 141, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_softclip", 142, 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", 143, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_single_channel_variance", 144, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_multi_channel_variance", 145, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), + InstructionOwner(new NonzeroInstruction(146, true)), + InstructionOwner(new BatchGatherInstruction(147, true)), + InstructionOwner(new BatchScatterInstruction(148, true)), + InstructionOwner(new RandomInstruction(149, true)), + InstructionOwner(new RandomIntInstruction(150, true)), + InstructionOwner(new UnweightInstruction(151, true)), + mi("vegas_forward", 152, 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", 153, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, true, {"n", "b"}, false}}, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}), + mi("vegas_histogram", 154, 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", 155, 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 5dd193731..2b5349849 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 114: backward_op_matmul(instr, locals, local_grads, device); break; -case 111: +case 115: backward_batch_foreach, backward_kernel_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 112: +case 116: backward_batch_foreach, backward_kernel_leaky_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 113: +case 117: backward_batch_foreach, backward_kernel_elu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 118: backward_batch_foreach, backward_kernel_gelu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 115: +case 119: backward_batch_foreach, backward_kernel_sigmoid, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 116: +case 120: backward_batch_foreach, backward_kernel_softplus, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 117: +case 121: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 118: +case 122: 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 123: 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 124: 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 125: backward_batch_foreach, backward_kernel_softmax, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 122: +case 126: 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 130: 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 133: backward_batch_foreach, backward_kernel_gather, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 132: +case 136: backward_batch_foreach, backward_kernel_select, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 137: +case 141: 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 142: 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 143: 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 144: 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 145: 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 3e3cc2400..a454016df 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -209,251 +209,263 @@ case 68: batch_foreach, kernel_t_inv_value_and_min_max, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); break; case 69: - batch_foreach, kernel_t1_inv_min_max_doublet, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_t_inv_min_max_cut, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); break; case 70: - batch_foreach, kernel_t1_inv_value_and_min_max_doublet, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device); + batch_foreach, kernel_t_inv_value_and_min_max_cut, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 71: - batch_foreach, kernel_t2_inv_min_max_doublet, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_t1_inv_min_max_doublet, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); break; case 72: - batch_foreach, kernel_t2_inv_value_and_min_max_doublet, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_t1_inv_value_and_min_max_doublet, 7, 3, 1, DeviceType>, 7, 3>(instr, locals, device); break; case 73: - batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); + batch_foreach, kernel_t2_inv_min_max_doublet, 7, 2, 1, DeviceType>, 7, 2>(instr, locals, device); break; case 74: - batch_foreach, kernel_s23_value_and_min_max, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_t2_inv_value_and_min_max_doublet, 8, 3, 1, DeviceType>, 8, 3>(instr, locals, device); break; case 75: - batch_foreach, kernel_invariants_from_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); break; case 76: - batch_foreach, kernel_sde2_channel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_s23_value_and_min_max, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 77: - batch_foreach, kernel_subchannel_weights, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_s23_min_max_cut, 10, 2, 1, DeviceType>, 10, 2>(instr, locals, device); break; case 78: - batch_foreach, kernel_apply_subchannel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_s23_value_and_min_max_cut, 10, 3, 1, DeviceType>, 10, 3>(instr, locals, device); break; case 79: - batch_foreach, kernel_pt_eta_phi_x, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_invariants_from_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 80: - batch_foreach, kernel_mirror_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sde2_channel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 81: - batch_foreach, kernel_momenta_to_x1x2, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_subchannel_weights, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 82: - batch_foreach, kernel_uniform_invariant, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_apply_subchannel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 83: - batch_foreach, kernel_uniform_invariant_inverse, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_pt_eta_phi_x, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 84: - batch_foreach, kernel_breit_wigner_invariant, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_mirror_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 85: - batch_foreach, kernel_breit_wigner_invariant_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_momenta_to_x1x2, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 86: - batch_foreach, kernel_stable_invariant, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_uniform_invariant, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 87: - batch_foreach, kernel_stable_invariant_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_uniform_invariant_inverse, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 88: - batch_foreach, kernel_stable_invariant_nu, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_breit_wigner_invariant, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 89: - batch_foreach, kernel_stable_invariant_nu_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_breit_wigner_invariant_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 90: - batch_foreach, kernel_fast_rambo_massless, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 91: - batch_foreach, kernel_fast_rambo_massless_inverse, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device); + batch_foreach, kernel_stable_invariant_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 92: - batch_foreach, kernel_fast_rambo_massless_com, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant_nu, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 93: - batch_foreach, kernel_fast_rambo_massive, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant_nu_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 94: - batch_foreach, kernel_fast_rambo_massive_inverse, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 95: - batch_foreach, kernel_fast_rambo_massive_com, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless_inverse, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device); break; case 96: - batch_foreach, kernel_cut_unphysical, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless_com, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 97: - batch_foreach, kernel_cut_one, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 98: - batch_foreach, kernel_cut_all, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive_inverse, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); break; case 99: - batch_foreach, kernel_cut_any, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive_com, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 100: - batch_foreach, kernel_scale_transverse_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_unphysical, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 101: - batch_foreach, kernel_scale_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_one, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 102: - batch_foreach, kernel_scale_half_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_all, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 103: - batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_cut_any, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 104: - batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_scale_transverse_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 105: - batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_scale_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 106: - op_matrix_element(instr, locals, device); + batch_foreach, kernel_scale_half_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 107: - batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_scale_partonic_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 108: - batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 109: - batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 110: - op_matmul(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 111: - batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 112: - batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 113: - batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_alpha_s, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 114: - batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 115: - batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_relu, 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_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 117: - op_rqs_reshape(instr, locals, device); + batch_foreach, kernel_elu, 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); + batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 119: - batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 120: - batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 121: - batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 122: - batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); break; case 123: - batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); break; case 124: - batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_inverse, 2, 2, 2, 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_softmax, 1, 1, DeviceType>, 1, 1>(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_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 127: - op_discrete_histogram(instr, locals, device); + batch_foreach, kernel_sample_discrete, 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); + batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 129: - batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 130: - batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_probs_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 131: - batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 132: - batch_foreach, kernel_select, 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 133: - batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 134: - batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 135: - op_quantile(instr, locals, device); + batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 136: - batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 137: - batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_select_vector, 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_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 139: - batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + op_quantile(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_one_hot, 2, 1, 1, DeviceType>, 2, 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_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 142: - op_nonzero(instr, locals, device); + batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 143: - op_batch_gather(instr, locals, device); + batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 144: - op_batch_scatter(instr, locals, device); + batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 145: - op_random(instr, locals, device); + batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 146: - op_random_int(instr, locals, device); + op_nonzero(instr, locals, device); break; case 147: - op_unweight(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 148: - batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 149: - batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_random(instr, locals, device); break; case 150: - op_vegas_histogram(instr, locals, device); + op_random_int(instr, locals, device); break; case 151: + op_unweight(instr, locals, device); + break; +case 152: + batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 153: + batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 154: + op_vegas_histogram(instr, locals, device); + break; +case 155: 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 59453fdab..0980fa79d 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 114: backward_op_matmul(instr, locals, local_grads, device); break; -case 111: +case 115: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 112: +case 116: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 113: +case 117: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 114: +case 118: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 115: +case 119: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 116: +case 120: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 117: +case 121: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 118: +case 122: backward_batch_foreach, 5, 4, 2>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 119: +case 123: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 120: +case 124: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 121: +case 125: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 122: +case 126: backward_batch_foreach, 2, 2, 1>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 126: +case 130: backward_batch_foreach, 4, 2, 1>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 129: +case 133: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 132: +case 136: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 137: +case 141: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 138: +case 142: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 139: +case 143: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 140: +case 144: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 141: +case 145: 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 72271c949..e1e6cb286 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -209,251 +209,263 @@ case 68: batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); break; case 69: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); break; case 70: - batch_foreach, 5, 3, 1>, 5, 3>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 71: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); break; case 72: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 7, 3, 1>, 7, 3>(instr, locals, device); break; case 73: - batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); + batch_foreach, 7, 2, 1>, 7, 2>(instr, locals, device); break; case 74: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 8, 3, 1>, 8, 3>(instr, locals, device); break; case 75: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); break; case 76: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 77: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 10, 2, 1>, 10, 2>(instr, locals, device); break; case 78: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 10, 3, 1>, 10, 3>(instr, locals, device); break; case 79: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 80: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 81: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 82: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 83: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 84: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 85: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 86: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 87: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 88: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 89: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 90: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 91: - batch_foreach, 2, 3, 1>, 2, 3>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 92: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 93: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 94: - batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 95: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 2, 3, 1>, 2, 3>(instr, locals, device); break; case 96: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 97: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 98: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); break; case 99: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 100: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 101: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 102: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 103: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 104: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 105: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 106: - op_matrix_element(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 107: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 108: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 109: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 110: - op_matmul(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 111: - batch_foreach, 1, 1>, 1, 1>(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); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 113: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 114: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matmul(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); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 119: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 120: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 121: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 122: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); break; case 123: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 124: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 125: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 126: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(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); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 129: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 130: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 131: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 132: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 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, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 136: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 137: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 138: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 139: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + op_quantile(instr, locals, device); break; case 140: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 142: - op_nonzero(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 143: - op_batch_gather(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 144: - op_batch_scatter(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 145: - op_random(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 146: - op_random_int(instr, locals, device); + op_nonzero(instr, locals, device); break; case 147: - op_unweight(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 148: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_batch_scatter(instr, locals, device); break; case 149: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_random(instr, locals, device); break; case 150: - op_vegas_histogram(instr, locals, device); + op_random_int(instr, locals, device); break; case 151: + op_unweight(instr, locals, device); + break; +case 152: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 153: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 154: + op_vegas_histogram(instr, locals, device); + break; +case 155: op_histogram(instr, locals, device); break; diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index cacfdcac4..67fe1cb67 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -129,6 +129,59 @@ t2_inv_min_max_doublet(FVal s, FVal m1_2, FVal mir_min_2, FVal t1_ab return {t2_min, t2_max}; } +// ETmin floor on the double-t central single particle (mass^2 m1_2) and its recoil +// system, ported from AmpliCol double_t. AmpliCol writes signed-t (lines 605-608); +// we work in |t| = -t, so its tmin=max/tmax=min map to our |t|_min/|t|_max with both +// edges negated. +template +KERNELSPEC Pair, FVal> doublet_t1_etmin_clamp( + FVal t_min, + FVal t_max, + FVal s, + FVal m1_2, + FVal etmin_i, + FVal etmin_ir +) { + auto sqrt_s = sqrt(max(s, EPS)); + auto ei2 = etmin_i * etmin_i; + auto eir2 = etmin_ir * etmin_ir; + auto lam = s * s + ei2 * ei2 + eir2 * eir2 - 2.0 * s * ei2 - 2.0 * ei2 * eir2 - + 2.0 * eir2 * s; + auto pzmax = sqrt(max(lam, 0.)) / (2.0 * sqrt_s); + auto Eimax = sqrt_s - sqrt(max(eir2 + pzmax * pzmax, 0.)); + auto cut_min = sqrt_s * (Eimax - pzmax) - m1_2; + auto cut_max = sqrt_s * (Eimax + pzmax) - m1_2; + auto active = (etmin_i > EPS) | (etmin_ir > EPS); + auto lo = where(active, max(t_min, cut_min), t_min); + auto hi = where(active, min(t_max, cut_max), t_max); + auto ok = hi > lo; + return {where(ok, lo, t_min), where(ok, hi, t_max)}; +} + +// ETmin floor on |t2| given the sampled |t1| (= t1_abs). AmpliCol double_t lines +// 625-626 (signed t). +template +KERNELSPEC Pair, FVal> doublet_t2_etmin_clamp( + FVal t_min, + FVal t_max, + FVal s, + FVal m1_2, + FVal t1_abs, + FVal etmin_i, + FVal etmin_ir +) { + auto ei2 = etmin_i * etmin_i; + auto eir2 = etmin_ir * etmin_ir; + auto denom2 = s - t1_abs - m1_2; + auto cut_min = s * ei2 / max(m1_2 + t1_abs, EPS) - m1_2; + auto cut_max = s * (1.0 - eir2 / max(denom2, EPS)) - m1_2; + auto active = (etmin_i > EPS) | (etmin_ir > EPS); + auto lo = where(active, max(t_min, cut_min), t_min); + auto hi = where(active & (denom2 > EPS), min(t_max, cut_max), t_max); + auto ok = hi > lo; + return {where(ok, lo, t_min), where(ok, hi, t_max)}; +} + template KERNELSPEC FVal lsquare(FourMom p) { return p[0] * p[0] - p[1] * p[1] - p[2] * p[2] - p[3] * p[3]; @@ -487,6 +540,131 @@ KERNELSPEC void kernel_t_inv_value_and_min_max( t_abs = -t_temp; } +// Cut-derived |t| interval from the AmpliCol/gen23 ETmin construction +// (phase_space_gen23.f03, gen23_one_step, lines 698-727). +// Convention chosen to match the kernels' m1/m2 mass labels: +// particle 1 == recoil system "ir" (mass^2 m1_2, ETmin floor etmin_1) +// particle 2 == single peeled jet "i" (mass^2 m2_2, ETmin floor etmin_2) +// t is measured w.r.t. beam a: t = (pa - p1)^2 = (pb - p2)^2, so the angle->|t| +// offset carries the recoil mass^2 (m1). +template +KERNELSPEC Pair, FVal> t_cut_bounds_etmin( + FourMom p_tot, + FourMom pa, + FVal m1_2, + FVal m2_2, + FVal etmin_1, + FVal etmin_2 +) { + auto pt_tot2 = p_tot[1] * p_tot[1] + p_tot[2] * p_tot[2]; + auto s = lsquare(p_tot); + auto piir0 = sqrt(max(s, 0.) + pt_tot2); + auto ppa0 = (p_tot[0] * pa[0] - p_tot[3] * pa[3]) / piir0; + // peeled jet (2): pt floor^2 = etmin_2^2 - m2^2; recoil (1) floor corrected + // for the system's own pt. + auto pt2_floor = etmin_2 * etmin_2 - m2_2; + auto d = sqrt(pt_tot2) - sqrt(max(pt2_floor, 0.)); + auto et_corr = sqrt(m1_2 + d * d); + auto eff_1 = max(etmin_1, where(pt_tot2 < pt2_floor, et_corr, sqrt(max(m1_2, 0.)))); + auto eff_2 = max(etmin_2, sqrt(max(m2_2, 0.))); + auto base = piir0 * piir0 - eff_2 * eff_2 + eff_1 * eff_1; + auto root = (piir0 - eff_2 - eff_1) * (piir0 + eff_2 - eff_1) * + (piir0 - eff_2 + eff_1) * (piir0 + eff_2 + eff_1); + auto rootsq = sqrt(max(root, 0.)); + auto ratio = ppa0 / piir0; + // signed-t -> |t|: min uses (base - root), max uses (base + root); the + // angle->|t| offset carries the recoil mass^2 (m1). + return {ratio * (base - rootsq) - m1_2, ratio * (base + rootsq) - m1_2}; +} + +// Clamp the kinematic |t| range against the ETmin cut interval. Both kernels +// return the absolute (positive) t invariant. The ETmin bound is a sampling +// optimization, not the cut enforcement: the integrator applies the real cut on +// the final momenta (cf. AmpliCol's explicit ET-check on the generated +// solutions). +template +KERNELSPEC void kernel_t_inv_min_max_cut( + FIn pa, + FIn pb, + FIn m1, + FIn m2, + FIn etmin_1, + FIn etmin_2, + FOut t_min, + FOut t_max +) { + FourMom p_tot; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + } + auto s = lsquare(p_tot); + auto ma_2 = lsquare(load_mom(pa)); + auto mb_2 = lsquare(load_mom(pb)); + auto m1_2 = m1 * m1; + auto m2_2 = m2 * m2; + auto t_min_max = t_inv_min_max(s, ma_2, mb_2, m1_2, m2_2); + auto tmn = t_min_max.first; + auto tmx = t_min_max.second; + + // particle 1 = recoil (m1, etmin_1), particle 2 = peeled (m2, etmin_2) + auto cut = + t_cut_bounds_etmin(p_tot, load_mom(pa), m1_2, m2_2, etmin_1, etmin_2); + + // If the incoming particle pa is not along the beam axis, the ETmin cut is + // is too tight. In that case, we just use the kinematic bounds. + // auto tmn_c = max(tmn, cut.first); + // auto tmx_c = min(tmx, cut.second); + auto proj_pt2 = pa[1] * pa[1] + pa[2] * pa[2]; + auto proj_along_z = proj_pt2 <= (pa[0] * pa[0]) * 1e-9; + auto tmn_c = where(proj_along_z, max(tmn, cut.first), tmn); + auto tmx_c = where(proj_along_z, min(tmx, cut.second), tmx); + auto ok = tmx_c > tmn_c; + t_min = where(ok, tmn_c, tmn); + t_max = where(ok, tmx_c, tmx); +} + +template +KERNELSPEC void kernel_t_inv_value_and_min_max_cut( + FIn pa, + FIn pb, + FIn p1, + FIn p2, + FIn etmin_1, + FIn etmin_2, + FOut t_abs, + FOut t_min, + FOut t_max +) { + FourMom pa1, p_tot; + for (int i = 0; i < 4; ++i) { + pa1[i] = pa[i] - p1[i]; + p_tot[i] = pa[i] + pb[i]; + } + auto s = lsquare(p_tot); + auto t_temp = lsquare(pa1); + auto ma_2 = lsquare(load_mom(pa)); + auto mb_2 = lsquare(load_mom(pb)); + auto m1_2 = lsquare(load_mom(p1)); + auto m2_2 = lsquare(load_mom(p2)); + auto t_min_max = t_inv_min_max(s, ma_2, mb_2, m1_2, m2_2); + auto tmn = t_min_max.first; + auto tmx = t_min_max.second; + + // particle 1 = recoil (m1, etmin_1), particle 2 = peeled (m2, etmin_2) + auto cut = + t_cut_bounds_etmin(p_tot, load_mom(pa), m1_2, m2_2, etmin_1, etmin_2); + // auto tmn_c = max(tmn, cut.first); + // auto tmx_c = min(tmx, cut.second); + auto proj_pt2 = pa[1] * pa[1] + pa[2] * pa[2]; + auto proj_along_z = proj_pt2 <= (pa[0] * pa[0]) * 1e-9; + auto tmn_c = where(proj_along_z, max(tmn, cut.first), tmn); + auto tmx_c = where(proj_along_z, min(tmx, cut.second), tmx); + auto ok = tmx_c > tmn_c; + t_min = where(ok, tmn_c, tmn); + t_max = where(ok, tmx_c, tmx); + t_abs = -t_temp; +} + template KERNELSPEC void kernel_invariants_from_momenta( FIn p_ext, FIn factors, FOut invariants @@ -512,6 +690,8 @@ KERNELSPEC void kernel_t1_inv_min_max_doublet( FIn pb, FIn m1, FIn mir_min, + FIn etmin_i, + FIn etmin_ir, FOut t_min, FOut t_max ) { @@ -520,9 +700,13 @@ KERNELSPEC void kernel_t1_inv_min_max_doublet( p_tot[i] = pa[i] + pb[i]; } auto s = lsquare(p_tot); - auto bounds = t1_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min); - t_min = bounds.first; - t_max = bounds.second; + auto m1_2 = m1 * m1; + auto bounds = t1_inv_min_max_doublet(s, m1_2, mir_min * mir_min); + auto c = doublet_t1_etmin_clamp( + bounds.first, bounds.second, s, m1_2, etmin_i, etmin_ir + ); + t_min = c.first; + t_max = c.second; } template @@ -532,6 +716,8 @@ KERNELSPEC void kernel_t1_inv_value_and_min_max_doublet( FIn p1, FIn m1, FIn mir_min, + FIn etmin_i, + FIn etmin_ir, FOut t_abs, FOut t_min, FOut t_max @@ -542,9 +728,13 @@ KERNELSPEC void kernel_t1_inv_value_and_min_max_doublet( p_tot[i] = pa[i] + pb[i]; } auto s = lsquare(p_tot); - auto bounds = t1_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min); - t_min = bounds.first; - t_max = bounds.second; + auto m1_2 = m1 * m1; + auto bounds = t1_inv_min_max_doublet(s, m1_2, mir_min * mir_min); + auto c = doublet_t1_etmin_clamp( + bounds.first, bounds.second, s, m1_2, etmin_i, etmin_ir + ); + t_min = c.first; + t_max = c.second; t_abs = -lsquare(pa1); } @@ -555,6 +745,8 @@ KERNELSPEC void kernel_t2_inv_min_max_doublet( FIn m1, FIn mir_min, FIn t1_abs, + FIn etmin_i, + FIn etmin_ir, FOut t_min, FOut t_max ) { @@ -563,9 +755,13 @@ KERNELSPEC void kernel_t2_inv_min_max_doublet( p_tot[i] = pa[i] + pb[i]; } auto s = lsquare(p_tot); - auto bounds = t2_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min, t1_abs); - t_min = bounds.first; - t_max = bounds.second; + auto m1_2 = m1 * m1; + auto bounds = t2_inv_min_max_doublet(s, m1_2, mir_min * mir_min, t1_abs); + auto c = doublet_t2_etmin_clamp( + bounds.first, bounds.second, s, m1_2, t1_abs, etmin_i, etmin_ir + ); + t_min = c.first; + t_max = c.second; } template @@ -576,6 +772,8 @@ KERNELSPEC void kernel_t2_inv_value_and_min_max_doublet( FIn m1, FIn mir_min, FIn t1_abs, + FIn etmin_i, + FIn etmin_ir, FOut t_abs, FOut t_min, FOut t_max @@ -586,9 +784,13 @@ KERNELSPEC void kernel_t2_inv_value_and_min_max_doublet( p_tot[i] = pa[i] + pb[i]; } auto s = lsquare(p_tot); - auto bounds = t2_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min, t1_abs); - t_min = bounds.first; - t_max = bounds.second; + auto m1_2 = m1 * m1; + auto bounds = t2_inv_min_max_doublet(s, m1_2, mir_min * mir_min, t1_abs); + auto c = doublet_t2_etmin_clamp( + bounds.first, bounds.second, s, m1_2, t1_abs, etmin_i, etmin_ir + ); + t_min = c.first; + t_max = c.second; t_abs = -lsquare(pb1); } diff --git a/madspace/src/kernels/threeparticle.hpp b/madspace/src/kernels/threeparticle.hpp index be56503c3..9cebce51f 100644 --- a/madspace/src/kernels/threeparticle.hpp +++ b/madspace/src/kernels/threeparticle.hpp @@ -537,6 +537,178 @@ KERNELSPEC void kernel_s23_value_and_min_max( s_23 = lsquare(p_23); } +// AmpliCol block-B s-channel ETmin refinement (gen23 lines 745-762). Tightens +// the sampled s-channel mass s23 = (peeled + im1)^2 using the recoil's ETmin +// floor (upper bound, smax) and the peeled particle's ETmin + dR cut (lower +// bound, smin -- only for a massless peeled particle). +template +KERNELSPEC Pair, FVal> s23_etmin_clamp( + FVal smn, + FVal smx, + FourMom piir, + FourMom pim1, + FourMom pib, + FVal m1_2, + FVal m2_2, + FVal m3_2, + FVal t1_abs, + FVal etmin_1, + FVal etmin_2, + FVal drcut +) { + auto active = (etmin_2 > EPS) | (etmin_1 > EPS); + // Block B works in the frame where pz(im1)=0 (z-boost by the IM1 rapidity), + // with im1 rotated onto +x. mT(im1)=sqrt(E^2-pz^2)=sqrt(m3^2+pt(p3)^2). + auto Eim1 = pim1[0]; + auto pzim1 = pim1[3]; + auto mT_im1 = sqrt(max(Eim1 * Eim1 - pzim1 * pzim1, EPS)); + // Boost p_12 (=piir, the i+ir system) and pa (=pib) into that frame; im1's + // own energy there is mT(im1) and its x-momentum is pt(p3). + auto piir_0 = (piir[0] * Eim1 - piir[3] * pzim1) / mT_im1; + auto pim1_0 = mT_im1; + auto pim1_1 = sqrt(pim1[1] * pim1[1] + pim1[2] * pim1[2]); + auto pib_0 = max((pib[0] * Eim1 - pib[3] * pzim1) / mT_im1, EPS); + // Block-B recoil ET floor (t-dependent): invm(ir)-invm(ir+ib) = m1_2 + t1_abs. + auto denom = max(m1_2 + t1_abs, EPS); + auto etminir_b = + max(etmin_1, pib_0 * etmin_1 * etmin_1 / denom + denom / (4. * pib_0)); + // smax: the recoil's ETmin caps the s-channel mass from above. + auto e_eff = piir_0 - etminir_b; + auto disc = e_eff * e_eff - m2_2; + auto smax_b = m2_2 + m3_2 + 2. * e_eff * pim1_0 + 2. * sqrt(max(disc, 0.)) * pim1_1; + auto smax_ok = active & (e_eff > 0.) & (disc > 0.) & (smax_b > smn); + auto smx_new = where(smax_ok, min(smx, smax_b), smx); + // smin: a massless peeled particle needs ET -> minimum s-channel mass. + auto smin_b = 2. * etmin_2 * (pim1_0 - pim1_1 * cos(drcut)); + auto smin_ok = active & (m2_2 < EPS) & (smin_b > smn) & (smin_b < smx_new); + auto smn_new = where(smin_ok, max(smn, smin_b), smn); + smx_new = where(smx_new > smn_new, smx_new, smn_new + EPS); + return {smn_new, smx_new}; +} + +// Clamp the s23 invariant-mass range against the cut-derived lower bound. +// s23_min_cut is the minimum invariant mass^2 of the (2,3) system implied by +// the pt / m_inv_min / dR cuts (gen23's invm_min). +template +KERNELSPEC void kernel_s23_min_max_cut( + FIn pa, + FIn pb, + FIn p3, + FIn t1_abs, + FIn m1, + FIn m2, + FIn etmin_1, + FIn etmin_2, + FIn drcut, + FIn s23_min_cut, + FOut s23_min, + FOut s23_max +) { + FourMom p_tot, p_12, pt2; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + p_12[i] = pa[i] + pb[i] - p3[i]; + pt2[i] = pb[i] - p3[i]; + } + auto m0_2 = lsquare(p_tot); + auto ma_2 = lsquare(load_mom(pa)); + auto mb_2 = lsquare(load_mom(pb)); + auto m3_2 = lsquare(load_mom(p3)); + auto s12 = lsquare(p_12); + auto m1_2 = m1 * m1; + auto m2_2 = m2 * m2; + auto t2 = lsquare(pt2); + + auto s23_out = s23_min_max(m0_2, ma_2, mb_2, m1_2, m2_2, m3_2, t1_abs, t2, s12); + auto smn = s23_out.first; + auto smx = s23_out.second; + + FVal smin_cut(s23_min_cut); + smn = where(smin_cut > 0., max(smn, smin_cut), smn); + smx = where(smx > smn, smx, smn + EPS); + + // Block-B ETmin refinement: piir = p_12 (recoil system), pim1 = p3, pib = pa. + auto sb = s23_etmin_clamp( + smn, + smx, + p_12, + load_mom(p3), + load_mom(pa), + m1_2, + m2_2, + m3_2, + t1_abs, + etmin_1, + etmin_2, + drcut + ); + + s23_min = sb.first; + s23_max = sb.second; +} + +template +KERNELSPEC void kernel_s23_value_and_min_max_cut( + FIn pa, + FIn pb, + FIn p3, + FIn t1_abs, + FIn p1, + FIn p2, + FIn etmin_1, + FIn etmin_2, + FIn drcut, + FIn s23_min_cut, + FOut s_23, + FOut s23_min, + FOut s23_max +) { + FourMom p_tot, p_12, pt2, p_23; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + p_12[i] = p1[i] + p2[i]; + pt2[i] = pb[i] - p3[i]; + p_23[i] = p2[i] + p3[i]; + } + auto m0_2 = lsquare(p_tot); + auto ma_2 = lsquare(load_mom(pa)); + auto mb_2 = lsquare(load_mom(pb)); + auto m3_2 = lsquare(load_mom(p3)); + auto s12 = lsquare(p_12); + auto m1_2 = lsquare(load_mom(p1)); + auto m2_2 = lsquare(load_mom(p2)); + auto t2 = lsquare(pt2); + + auto s23_out = s23_min_max(m0_2, ma_2, mb_2, m1_2, m2_2, m3_2, t1_abs, t2, s12); + auto smn = s23_out.first; + auto smx = s23_out.second; + + FVal smin_cut(s23_min_cut); + smn = where(smin_cut > 0., max(smn, smin_cut), smn); + smx = where(smx > smn, smx, smn + EPS); + + // Block-B ETmin refinement (must mirror the forward kernel exactly so the + // sampled s23 range is identical and the round-trip stays invertible). + auto sb = s23_etmin_clamp( + smn, + smx, + p_12, + load_mom(p3), + load_mom(pa), + m1_2, + m2_2, + m3_2, + t1_abs, + etmin_1, + etmin_2, + drcut + ); + + s23_min = sb.first; + s23_max = sb.second; + s_23 = lsquare(p_23); +} + template KERNELSPEC void kernel_two_to_three_particle_scattering( IIn phi_index, diff --git a/madspace/src/phasespace/color_ordered_mapping.cpp b/madspace/src/phasespace/color_ordered_mapping.cpp index dce771ce6..14cc37317 100644 --- a/madspace/src/phasespace/color_ordered_mapping.cpp +++ b/madspace/src/phasespace/color_ordered_mapping.cpp @@ -1,6 +1,7 @@ #include "madspace/phasespace/color_ordered_mapping.hpp" #include +#include #include #include "madspace/util.hpp" @@ -29,7 +30,7 @@ split_sets_from_color_order(const std::vector& color_order) { std::vector rotated; rotated.reserve(n); for (std::size_t k = 0; k < n; ++k) { - rotated.push_back(color_order[(i0 + k) % n]); + rotated.push_back(color_order.at((i0 + k) % n)); } auto it1 = std::find(rotated.begin(), rotated.end(), 1u); if (it1 == rotated.end()) { @@ -38,14 +39,14 @@ split_sets_from_color_order(const std::vector& color_order) { std::size_t i1 = std::distance(rotated.begin(), it1); std::vector set1, set2; for (std::size_t k = 1; k < i1; ++k) { - std::size_t p = rotated[k]; + std::size_t p = rotated.at(k); if (p <= 1) { throw std::invalid_argument("invalid color_order"); } set1.push_back(p - 2); } for (std::size_t k = i1 + 1; k < n; ++k) { - std::size_t p = rotated[k]; + std::size_t p = rotated.at(k); if (p <= 1) { throw std::invalid_argument("invalid color_order"); } @@ -70,16 +71,91 @@ std::size_t n_block_randoms_for_set_size(std::size_t k) { if (k <= 1) { return 0; } - // First peel: 2->2 LAB (2 randoms); each subsequent peel: 2->3 (3 randoms). - return 2 + 3 * (k - 2); + // Continuous randoms only. First peel: 2->2 LAB (2 randoms). Each subsequent + // peel: 2->3 (2 randoms). + return 2 + 2 * (k - 2); +} + +std::size_t n_discrete_for_set_size(std::size_t k) { + // One discrete two-solution choice per 2->3 peel (all peels after the first). + return (k >= 2) ? (k - 2) : 0; +} + +double mat_at(const std::vector>& m, std::size_t i, std::size_t j) { + if (i < m.size() && j < m.at(i).size()) { + return m.at(i).at(j); + } + return 0.0; } } // namespace +// Whether any cut (pt, invariant-mass, or dR) is active. With no cut the cut +// kernels reduce to the plain ones, so we skip them to keep the map exactly +// invertible and avoid boosted-frame round-trip noise. +static bool has_any_cut( + const std::vector& pt_min, + const std::vector>& m_inv_min, + const std::vector>& dr_min +) { + for (double p : pt_min) { + if (p > 0.0) { + return true; + } + } + for (const auto& row : m_inv_min) { + for (double v : row) { + if (v > 0.0) { + return true; + } + } + } + for (const auto& row : dr_min) { + for (double v : row) { + if (v > 0.0) { + return true; + } + } + } + return false; +} + +double ColorOrderedMapping::pt2(std::size_t i) const { + double p = (i < _pt_min.size()) ? _pt_min.at(i) : 0.0; + return p * p; +} + +double ColorOrderedMapping::cut_floor(const std::vector& subset) const { + if (subset.size() < 2) { + return 0.0; + } + // Sum over distinct pairs of the per-pair invariant-mass floor implied by + // m_inv_min and the pt/dR cut, exactly as gen23's setup_PS_cuts. + double cut = 0.0; + for (std::size_t a = 0; a + 1 < subset.size(); ++a) { + for (std::size_t b = a + 1; b < subset.size(); ++b) { + std::size_t i = subset.at(a), j = subset.at(b); + double ss = mat_at(_m_inv_min, i, j); + double pti = (i < _pt_min.size()) ? _pt_min.at(i) : 0.0; + double ptj = (j < _pt_min.size()) ? _pt_min.at(j) : 0.0; + double dr = mat_at(_dr_min, i, j); + cut += std::max(ss * ss, 2.0 * pti * ptj * (1.0 - std::cos(dr))); + } + } + double npart = static_cast(subset.size()); + // gen23 uses npart/(npart-1) for the full final state and 1/2 for proper + // subsets. + double scaling = (subset.size() == _n_out) ? npart / (npart - 1.0) : 0.5; + return cut * scaling; +} + ColorOrderedMapping::ColorOrderedMapping( const std::vector& color_order, double t_invariant_power, - double s_invariant_power + double s_invariant_power, + const std::vector& pt_min, + const std::vector>& m_inv_min, + const std::vector>& dr_min ) : Mapping( "ColorOrderedMapping", @@ -102,10 +178,17 @@ ColorOrderedMapping::ColorOrderedMapping( n_block_randoms_for_set_size(s2.size()); std::size_t total = n_set_masses + n_intermediate_masses + n_central + n_walk; + std::size_t n_discrete = + n_discrete_for_set_size(s1.size()) + n_discrete_for_set_size(s2.size()); NamedVector input_types; for (std::size_t i = 0; i < total; ++i) { input_types.push_back(std::format("random{}", i), batch_float); } + // Opt-in discrete channel: one int per 2->3 peel (the two-solution + // choice), appended after the continuous randoms. + for (std::size_t j = 0; j < n_discrete; ++j) { + input_types.push_back(std::format("discrete{}", j), batch_int); + } return input_types; }(), [&] { @@ -126,10 +209,34 @@ ColorOrderedMapping::ColorOrderedMapping( }() ), _n_out(color_order.size() - 2), - _com_scattering(true, t_invariant_power), - _lab_scattering(false, t_invariant_power), - _two_to_three(t_invariant_power, 0., 0., s_invariant_power, 0., 0.), - _double_t(t_invariant_power, 0., 0., t_invariant_power, 0., 0.) { + _pt_min(pt_min), + _m_inv_min(m_inv_min), + _dr_min(dr_min), + _has_cut(has_any_cut(pt_min, m_inv_min, dr_min)), + _com_scattering( + true, t_invariant_power, 0., 0., has_any_cut(pt_min, m_inv_min, dr_min) + ), + _lab_scattering( + false, t_invariant_power, 0., 0., has_any_cut(pt_min, m_inv_min, dr_min) + ), + _two_to_three( + t_invariant_power, + 0., + 0., + s_invariant_power, + 0., + 0., + has_any_cut(pt_min, m_inv_min, dr_min) + ), + _double_t( + t_invariant_power, + 0., + 0., + t_invariant_power, + 0., + 0., + has_any_cut(pt_min, m_inv_min, dr_min) + ) { auto [s1, s2] = split_sets_from_color_order(color_order); _set1 = s1; _set2 = s2; @@ -144,6 +251,8 @@ ColorOrderedMapping::ColorOrderedMapping( std::size_t n_walk = n_block_randoms_for_set_size(s1.size()) + n_block_randoms_for_set_size(s2.size()); _random_dim = n_set_masses + n_intermediate_masses + n_central + n_walk; + _discrete_dim = + n_discrete_for_set_size(s1.size()) + n_discrete_for_set_size(s2.size()); } Mapping::Result ColorOrderedMapping::build_forward_impl( @@ -155,6 +264,8 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( ValueVec m_out(conditions.begin() + 1, conditions.end()); auto r = inputs.begin(); auto next_random = [&]() { return *(r++); }; + auto r_disc = inputs.begin() + _random_dim; + auto next_discrete = [&]() { return *(r_disc++); }; ValueVec dets; // Phase 1a + Phase 2: set composite masses and the central block. @@ -198,18 +309,37 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( Value mass_sum_set1 = masses_of(_set1); Value mass_sum_set2 = masses_of(_set2); if (_set1.size() == 1) { - m_set1 = m_out.at(_set1[0]); + m_set1 = m_out.at(_set1.at(0)); } if (_set2.size() == 1) { - m_set2 = m_out.at(_set2[0]); + m_set2 = m_out.at(_set2.at(0)); } bool single_is_set1 = (_set1.size() == 1); Value m_single = single_is_set1 ? m_set1 : m_set2; Value mir_min = single_is_set1 ? mass_sum_set2 : mass_sum_set1; + // ETmin floors (pt-cut tightening of the central t1/t2 bounds): the lone + // particle's ET and the summed ET of the recoil set. Zero (inert) with no cut. + Value etmin_i = Value(0.); + Value etmin_ir = Value(0.); + if (_has_cut) { + auto etmin_p = [&](std::size_t idx) { + return fb.sqrt(fb.add(fb.square(m_out.at(idx)), Value(pt2(idx)))); + }; + std::size_t single_idx = single_is_set1 ? _set1.at(0) : _set2.at(0); + const auto& recoil = single_is_set1 ? _set2 : _set1; + etmin_i = etmin_p(single_idx); + etmin_ir = etmin_p(recoil.at(0)); + for (std::size_t j = 1; j < recoil.size(); ++j) { + etmin_ir = fb.add(etmin_ir, etmin_p(recoil.at(j))); + } + } + ValueVec dt_cond{pa, pb, m_single, mir_min}; + if (_has_cut) { + dt_cond.push_back(etmin_i); + dt_cond.push_back(etmin_ir); + } auto central = _double_t.build_forward( - fb, - {next_random(), next_random(), next_random()}, - {pa, pb, m_single, mir_min} + fb, {next_random(), next_random(), next_random()}, dt_cond ); Value p_single = central.at(0); Value p_recoil = central.at(1); @@ -241,26 +371,59 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( Value mass_sum_set2 = masses_of(_set2); if (_set1.size() >= 2) { auto s_min = fb.square(mass_sum_set1); + // Cut floor: composite invariant mass must clear the cut-implied + // minimum (gen23 invm_min) in addition to the kinematic minimum. + double floor1 = cut_floor(_set1); + if (floor1 > 0.0) { + s_min = fb.max(s_min, Value(floor1)); + } auto s_max = fb.square(fb.sub(e_cm, mass_sum_set2)); auto res = _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); m_set1 = fb.sqrt(res["invariant"]); dets.push_back(res["det"]); } else { - m_set1 = m_out.at(_set1[0]); + m_set1 = m_out.at(_set1.at(0)); } if (_set2.size() >= 2) { auto s_min = fb.square(mass_sum_set2); + double floor2 = cut_floor(_set2); + if (floor2 > 0.0) { + s_min = fb.max(s_min, Value(floor2)); + } auto s_max = fb.square(fb.sub(e_cm, m_set1)); auto res = _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); m_set2 = fb.sqrt(res["invariant"]); dets.push_back(res["det"]); } else { - m_set2 = m_out.at(_set2[0]); + m_set2 = m_out.at(_set2.at(0)); + } + // ETmin floors for the central 2->2: each output system must carry enough + // transverse energy for its jets, so the floor is the summed ET of the set + // (a valid outer bound since ET_sys >= sum ET_constituents). Zero with no cut. + Value etmin_set1 = Value(0.); + Value etmin_set2 = Value(0.); + if (_has_cut) { + auto etmin_p = [&](std::size_t idx) { + return fb.sqrt(fb.add(fb.square(m_out.at(idx)), Value(pt2(idx)))); + }; + etmin_set1 = etmin_p(_set1.at(0)); + for (std::size_t j = 1; j < _set1.size(); ++j) { + etmin_set1 = fb.add(etmin_set1, etmin_p(_set1.at(j))); + } + etmin_set2 = etmin_p(_set2.at(0)); + for (std::size_t j = 1; j < _set2.size(); ++j) { + etmin_set2 = fb.add(etmin_set2, etmin_p(_set2.at(j))); + } + } + ValueVec com_cond{pa, pb}; + if (_has_cut) { + com_cond.push_back(etmin_set1); + com_cond.push_back(etmin_set2); } auto central = _com_scattering.build_forward( - fb, {next_random(), next_random(), m_set1, m_set2}, {pa, pb} + fb, {next_random(), next_random(), m_set1, m_set2}, com_cond ); P_set1 = central.at(0); P_set2 = central.at(1); @@ -288,12 +451,18 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( } Value prev_mass = m_set; for (std::size_t j = 0; j < k - 2; ++j) { - Value m_min = m_out.at(s[j + 1]); + Value m_min = m_out.at(s.at(j + 1)); for (std::size_t i = j + 2; i < k; ++i) { - m_min = fb.add(m_min, m_out.at(s[i])); + m_min = fb.add(m_min, m_out.at(s.at(i))); } auto s_min = fb.square(m_min); - auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s[j]))); + // Cut floor for the rest system {s[j+1], ..., s[k-1]}. + std::vector rest_sub(s.begin() + j + 1, s.end()); + double fl = cut_floor(rest_sub); + if (fl > 0.0) { + s_min = fb.max(s_min, Value(fl)); + } + auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s.at(j)))); auto r = _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); Value m_rest = fb.sqrt(r["invariant"]); @@ -321,20 +490,33 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( [&](const std::vector& s, Value P_set, Value R_b, + Value beam_other, const ValueVec& rest_masses) { std::size_t k = s.size(); if (k == 1) { - p_out[s[0]] = P_set; + p_out.at(s.at(0)) = P_set; return; } Value R_a = fb.sub(P_set, R_b); + auto etmin_particle = [&](std::size_t idx) { + return fb.sqrt(fb.add(fb.square(m_out.at(idx)), Value(pt2(idx)))); + }; + ValueVec etmin_suffix; + if (_has_cut) { + etmin_suffix.resize(k); + etmin_suffix.at(k - 1) = etmin_particle(s.at(k - 1)); + for (std::size_t j = k - 1; j-- > 0;) { + etmin_suffix.at(j) = + fb.add(etmin_particle(s.at(j)), etmin_suffix.at(j + 1)); + } + } Value im1; bool first = true; for (std::size_t j = 0; j < k - 1; ++j) { // New-rest mass: intermediate (rest_masses[j]) or final residual (j == // k-2). - Value m_rest = (j < k - 2) ? rest_masses[j] : m_out.at(s[k - 1]); - Value m_peel = m_out.at(s[j]); + Value m_rest = (j < k - 2) ? rest_masses[j] : m_out.at(s.at(k - 1)); + Value m_peel = m_out.at(s.at(j)); // Block convention: pa-side carries the chain (mass m_rest), pb-side // carries the peeled particle (mass m_peel). R_a is the active leg // and gets decremented by peeled; R_b is constant. @@ -343,38 +525,54 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( // has im1 subtracted from our previous step, so we pass pb = R_a + im1 // to recover p_12 = R_b + R_a inside the kernel. if (first) { + // First peel is a 2->2 LAB block, projecting the own-side on-z + // beam R_a (AmpliCol: gent_one_step beam i). Cut: ETmin of the + // recoil chain (etmin_1) then the peeled particle (etmin_2). + ValueVec cond{R_a, R_b}; + if (_has_cut) { + cond.push_back(etmin_suffix.at(j + 1)); + cond.push_back(etmin_particle(s.at(j))); + } auto ks = _lab_scattering.build_forward( - fb, {next_random(), next_random(), m_rest, m_peel}, {R_b, R_a} + fb, {next_random(), next_random(), m_rest, m_peel}, cond ); Value peeled = ks.at(1); - p_out[s[j]] = peeled; + p_out.at(s.at(j)) = peeled; R_a = fb.sub(R_a, peeled); im1 = peeled; dets.push_back(ks["det"]); first = false; } else { - Value pb_for_block = fb.add(R_a, im1); + Value S = fb.add(R_a, R_b); + Value pb_for_block = fb.add(fb.sub(S, beam_other), im1); + ValueVec cond{beam_other, pb_for_block, im1}; + if (_has_cut) { + cond.push_back(etmin_suffix.at(j + 1)); + cond.push_back(etmin_particle(s.at(j))); + cond.push_back(Value(mat_at(_dr_min, s.at(j - 1), s.at(j)))); + cond.push_back(Value(cut_floor({s.at(j - 1), s.at(j)}))); + } auto ks = _two_to_three.build_forward( fb, - {next_random(), next_random(), next_random(), m_rest, m_peel}, - {R_b, pb_for_block, im1} + {next_discrete(), next_random(), next_random(), m_rest, m_peel}, + cond ); Value peeled = ks.at(1); - p_out[s[j]] = peeled; + p_out.at(s.at(j)) = peeled; R_a = fb.sub(R_a, peeled); im1 = peeled; dets.push_back(ks["det"]); } } // Last particle = R_a + R_b (= what remains after all explicit peels). - p_out[s[k - 1]] = fb.add(R_a, R_b); + p_out.at(s.at(k - 1)) = fb.add(R_a, R_b); }; if (!_set1.empty()) { - walk(_set1, P_set1, R_b_for_set1, rest_masses_set1); + walk(_set1, P_set1, R_b_for_set1, pb, rest_masses_set1); } if (!_set2.empty()) { - walk(_set2, P_set2, R_b_for_set2, rest_masses_set2); + walk(_set2, P_set2, R_b_for_set2, pa, rest_masses_set2); } // Assemble outputs: momentum0, momentum1 = beams; momentum_{2+i} = outgoing i. @@ -383,7 +581,7 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( p_ext.push_back(pa); p_ext.push_back(pb); for (std::size_t i = 0; i < _n_out; ++i) { - p_ext.push_back(p_out[i]); + p_ext.push_back(p_out.at(i)); } return {{output_types().keys(), p_ext}, fb.product(dets)}; @@ -397,6 +595,7 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( Value e_cm = conditions.at(0); ValueVec m_out(conditions.begin() + 1, conditions.end()); ValueVec random_out; + ValueVec discrete_out; ValueVec dets; Value pa = inputs.at(0); @@ -447,9 +646,9 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( } auto sum_momenta = [&](const std::vector& s) -> Value { - Value p = p_outgoing(s[0]); + Value p = p_outgoing(s.at(0)); for (std::size_t k = 1; k < s.size(); ++k) { - p = fb.add(p, p_outgoing(s[k])); + p = fb.add(p, p_outgoing(s.at(k))); } return p; }; @@ -479,10 +678,10 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( } } else if (_use_double_t) { if (_set1.size() == 1) { - m_set1 = m_out.at(_set1[0]); + m_set1 = m_out.at(_set1.at(0)); m_set2 = fb.sqrt(invariants.at(idx_m2_set2)); } else { - m_set2 = m_out.at(_set2[0]); + m_set2 = m_out.at(_set2.at(0)); m_set1 = fb.sqrt(invariants.at(idx_m2_set1)); } } else { @@ -490,6 +689,10 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( Value mass_sum_set2 = masses_of(_set2); if (_set1.size() >= 2) { auto s_min = fb.square(mass_sum_set1); + double floor1 = cut_floor(_set1); + if (floor1 > 0.0) { + s_min = fb.max(s_min, Value(floor1)); + } auto s_max = fb.square(fb.sub(e_cm, mass_sum_set2)); auto m2_set1 = invariants.at(idx_m2_set1); auto res = _uniform_invariant.build_inverse(fb, {m2_set1}, {s_min, s_max}); @@ -497,10 +700,14 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( dets.push_back(res["det"]); m_set1 = fb.sqrt(m2_set1); } else { - m_set1 = m_out.at(_set1[0]); + m_set1 = m_out.at(_set1.at(0)); } if (_set2.size() >= 2) { auto s_min = fb.square(mass_sum_set2); + double floor2 = cut_floor(_set2); + if (floor2 > 0.0) { + s_min = fb.max(s_min, Value(floor2)); + } auto s_max = fb.square(fb.sub(e_cm, m_set1)); auto m2_set2 = invariants.at(idx_m2_set2); auto res = _uniform_invariant.build_inverse(fb, {m2_set2}, {s_min, s_max}); @@ -508,7 +715,7 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( dets.push_back(res["det"]); m_set2 = fb.sqrt(m2_set2); } else { - m_set2 = m_out.at(_set2[0]); + m_set2 = m_out.at(_set2.at(0)); } } @@ -523,15 +730,52 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( Value p_recoil = single_is_set1 ? P_set2 : P_set1; Value m_single = single_is_set1 ? m_set1 : m_set2; Value mir_min = single_is_set1 ? masses_of(_set2) : masses_of(_set1); - auto central = _double_t.build_inverse( - fb, {p_single, p_recoil}, {pa, pb, m_single, mir_min} - ); + Value etmin_i = Value(0.); + Value etmin_ir = Value(0.); + if (_has_cut) { + auto etmin_p = [&](std::size_t idx) { + return fb.sqrt(fb.add(fb.square(m_out.at(idx)), Value(pt2(idx)))); + }; + std::size_t single_idx = single_is_set1 ? _set1.at(0) : _set2.at(0); + const auto& recoil = single_is_set1 ? _set2 : _set1; + etmin_i = etmin_p(single_idx); + etmin_ir = etmin_p(recoil.at(0)); + for (std::size_t j = 1; j < recoil.size(); ++j) { + etmin_ir = fb.add(etmin_ir, etmin_p(recoil.at(j))); + } + } + ValueVec dt_cond{pa, pb, m_single, mir_min}; + if (_has_cut) { + dt_cond.push_back(etmin_i); + dt_cond.push_back(etmin_ir); + } + auto central = _double_t.build_inverse(fb, {p_single, p_recoil}, dt_cond); random_out.push_back(central.at(0)); random_out.push_back(central.at(1)); random_out.push_back(central.at(2)); dets.push_back(central["det"]); } else { - auto central = _com_scattering.build_inverse(fb, {P_set1, P_set2}, {pa, pb}); + Value etmin_set1 = Value(0.); + Value etmin_set2 = Value(0.); + if (_has_cut) { + auto etmin_p = [&](std::size_t idx) { + return fb.sqrt(fb.add(fb.square(m_out.at(idx)), Value(pt2(idx)))); + }; + etmin_set1 = etmin_p(_set1.at(0)); + for (std::size_t j = 1; j < _set1.size(); ++j) { + etmin_set1 = fb.add(etmin_set1, etmin_p(_set1.at(j))); + } + etmin_set2 = etmin_p(_set2.at(0)); + for (std::size_t j = 1; j < _set2.size(); ++j) { + etmin_set2 = fb.add(etmin_set2, etmin_p(_set2.at(j))); + } + } + ValueVec com_cond{pa, pb}; + if (_has_cut) { + com_cond.push_back(etmin_set1); + com_cond.push_back(etmin_set2); + } + auto central = _com_scattering.build_inverse(fb, {P_set1, P_set2}, com_cond); random_out.push_back(central.at(0)); random_out.push_back(central.at(1)); dets.push_back(central["det"]); @@ -548,12 +792,18 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( // the (sqrt of the) actual recovered m2 at each step. Value prev_mass = m_set; for (std::size_t j = 0; j < k - 2; ++j) { - Value m_min = m_out.at(s[j + 1]); + Value m_min = m_out.at(s.at(j + 1)); for (std::size_t i = j + 2; i < k; ++i) { - m_min = fb.add(m_min, m_out.at(s[i])); + m_min = fb.add(m_min, m_out.at(s.at(i))); } auto s_min = fb.square(m_min); - auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s[j]))); + // Identical cut floor to the forward pass. + std::vector rest_sub(s.begin() + j + 1, s.end()); + double fl = cut_floor(rest_sub); + if (fl > 0.0) { + s_min = fb.max(s_min, Value(fl)); + } + auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s.at(j)))); auto m2 = invariants.at(idx_start + j); auto res = _uniform_invariant.build_inverse(fb, {m2}, {s_min, s_max}); random_out.push_back(res["random"]); @@ -572,51 +822,86 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( Value R_b_for_set1 = fb.sub(pb, P_set2); Value R_b_for_set2 = fb.sub(pa, P_set1); - // Phase 3 inverse: peel-off walks - auto walk_inverse = [&](const std::vector& s, Value P_set, Value R_b) { - std::size_t k = s.size(); - if (k == 1) { - return; - } - Value R_a = fb.sub(P_set, R_b); - Value im1; - bool first = true; - for (std::size_t j = 0; j < k - 1; ++j) { - Value peeled = p_outgoing(s[j]); - // p1_out is the chain carrier (mass m_rest); by conservation - // p1_out = R_a + R_b - peeled. - Value p1_out = fb.sub(fb.add(R_a, R_b), peeled); - if (first) { - auto rs = - _lab_scattering.build_inverse(fb, {p1_out, peeled}, {R_b, R_a}); - random_out.push_back(rs.at(0)); - random_out.push_back(rs.at(1)); - dets.push_back(rs["det"]); - first = false; - } else { - // Mirror the forward's pb restoration: the 2->3 kernel - // internally subtracts p_3 = im1, so we pass pb = R_a + im1 - // to get p_12 = R_b + R_a (the remaining-to-produce system). - Value pb_for_block = fb.add(R_a, im1); - auto rs = _two_to_three.build_inverse( - fb, {p1_out, peeled}, {R_b, pb_for_block, im1} - ); - random_out.push_back(rs.at(0)); - random_out.push_back(rs.at(1)); - random_out.push_back(rs.at(2)); - dets.push_back(rs["det"]); + // Phase 3 inverse: peel-off walks (mirror the forward restructure: first peel + // projects the own-side beam R_a, 2->3 peels project beam_other = 3-i side). + auto walk_inverse = + [&]( + const std::vector& s, Value P_set, Value R_b, Value beam_other + ) { + std::size_t k = s.size(); + if (k == 1) { + return; } - R_a = fb.sub(R_a, peeled); - im1 = peeled; - } - }; + Value R_a = fb.sub(P_set, R_b); + // ETmin suffix sums for the recoil chain (see forward walk); only built + // when cuts are active. + auto etmin_particle = [&](std::size_t idx) { + return fb.sqrt(fb.add(fb.square(m_out.at(idx)), Value(pt2(idx)))); + }; + ValueVec etmin_suffix; + if (_has_cut) { + etmin_suffix.resize(k); + etmin_suffix.at(k - 1) = etmin_particle(s.at(k - 1)); + for (std::size_t j = k - 1; j-- > 0;) { + etmin_suffix.at(j) = + fb.add(etmin_particle(s.at(j)), etmin_suffix.at(j + 1)); + } + } + Value im1; + bool first = true; + for (std::size_t j = 0; j < k - 1; ++j) { + Value peeled = p_outgoing(s.at(j)); + // p1_out is the chain carrier (mass m_rest); by conservation + // p1_out = R_a + R_b - peeled. + Value p1_out = fb.sub(fb.add(R_a, R_b), peeled); + if (first) { + // Same projection as the forward 2->2 block: own-side beam R_a. + ValueVec cond{R_a, R_b}; + if (_has_cut) { + cond.push_back(etmin_suffix.at(j + 1)); + cond.push_back(etmin_particle(s.at(j))); + } + auto rs = _lab_scattering.build_inverse(fb, {p1_out, peeled}, cond); + random_out.push_back(rs.at(0)); + random_out.push_back(rs.at(1)); + dets.push_back(rs["det"]); + first = false; + } else { + // Mirror the forward 2->3 restructure: project beam_other, with + // S = R_a + R_b reconstructed as beam_other + (S - beam_other) and + // pb = (S - beam_other) + im1 (the kernel re-subtracts p_3 = im1). + Value S = fb.add(R_a, R_b); + Value pb_for_block = fb.add(fb.sub(S, beam_other), im1); + ValueVec cond{beam_other, pb_for_block, im1}; + if (_has_cut) { + cond.push_back(etmin_suffix.at(j + 1)); + cond.push_back(etmin_particle(s.at(j))); + cond.push_back(Value(mat_at(_dr_min, s.at(j - 1), s.at(j)))); + cond.push_back(Value(cut_floor({s.at(j - 1), s.at(j)}))); + } + auto rs = _two_to_three.build_inverse(fb, {p1_out, peeled}, cond); + // rs.at(0) is the discrete choice index (int) emitted by the + // 2->3 inverse; rs.at(1), rs.at(2) are the continuous r_s23, r_t1. + discrete_out.push_back(rs.at(0)); + random_out.push_back(rs.at(1)); + random_out.push_back(rs.at(2)); + dets.push_back(rs["det"]); + } + R_a = fb.sub(R_a, peeled); + im1 = peeled; + } + }; if (!_set1.empty()) { - walk_inverse(_set1, P_set1, R_b_for_set1); + walk_inverse(_set1, P_set1, R_b_for_set1, pb); } if (!_set2.empty()) { - walk_inverse(_set2, P_set2, R_b_for_set2); + walk_inverse(_set2, P_set2, R_b_for_set2, pa); } - return {{input_types().keys(), random_out}, fb.product(dets)}; + // input_types is [random0..random_{N-1}, discrete0..discrete_{M-1}], so the + // returned values must be the continuous randoms followed by the discrete ints. + ValueVec all_out = random_out; + all_out.insert(all_out.end(), discrete_out.begin(), discrete_out.end()); + return {{input_types().keys(), all_out}, fb.product(dets)}; } diff --git a/madspace/src/phasespace/cuts.cpp b/madspace/src/phasespace/cuts.cpp index d56a6da19..8fe09adca 100644 --- a/madspace/src/phasespace/cuts.cpp +++ b/madspace/src/phasespace/cuts.cpp @@ -91,3 +91,53 @@ std::vector Cuts::pt_min() const { } return pt_min; } + +std::vector> Cuts::pairwise_min( + Observable::ObservableOption obs, + const std::function< + std::vector>(const Observable&)>& pairs +) const { + std::size_t n = arg_types().at(0).shape.at(0) - 2; + std::vector> out(n, std::vector(n, 0.)); + for (auto& item : _cut_data) { + if (item.observable.observable() != obs) { + continue; + } + for (auto [i, j] : pairs(item.observable)) { + if (i < 2 || j < 2) { + continue; + } + i -= 2; + j -= 2; + if (i < n && j < n && item.min > out.at(i).at(j)) { + out.at(i).at(j) = item.min; + out.at(j).at(i) = item.min; + } + } + } + return out; +} + +std::vector> Cuts::m_inv_min() const { + return pairwise_min(Observable::obs_mass, [](const Observable& o) { + std::vector> pairs; + const auto& idx = o.indices(); + if (o.sum_momenta() && idx.size() == 1 && idx.at(0).size() == 2) { + pairs.emplace_back(idx.at(0).at(0), idx.at(0).at(1)); + } + return pairs; + }); +} + +std::vector> Cuts::dr_min() const { + return pairwise_min(Observable::obs_delta_r, [](const Observable& o) { + std::vector> pairs; + const auto& idx = o.indices(); + if (idx.size() == 2) { + for (std::size_t k = 0; k < idx.at(0).size(); ++k) { + pairs.emplace_back(idx.at(0).at(k), idx.at(1).at(k)); + } + } + return pairs; + }); +} diff --git a/madspace/src/phasespace/phasespace.cpp b/madspace/src/phasespace/phasespace.cpp index 16ca82db0..f95f2e165 100644 --- a/madspace/src/phasespace/phasespace.cpp +++ b/madspace/src/phasespace/phasespace.cpp @@ -1,6 +1,7 @@ #include "madspace/phasespace/phasespace.hpp" #include "madspace/constants.hpp" #include "madspace/util.hpp" +#include using namespace madspace; @@ -81,6 +82,42 @@ nested_vector2 invert_permutations(nested_vector2 perms_in) } // namespace +namespace { +// Chain (color) order for the t-channel ColorOrderedMapping: the externally +// supplied order if given, else the default single chain [0, 2, ..., n+1, 1]. +std::vector ps_chain_order( + const Topology& topology, const std::optional>& color_order +) { + if (color_order) { + return *color_order; + } + std::size_t n_t_out = topology.decays().at(0).child_indices.size(); + std::vector chain; + chain.reserve(n_t_out + 2); + chain.push_back(0); + for (std::size_t i = 0; i < n_t_out; ++i) { + chain.push_back(i + 2); + } + chain.push_back(1); + return chain; +} + +// Number of discrete two-solution choices for the t-channel (opt-in r_disc): +// non-zero only for color_ordered with at least one 2->3 peel. +std::size_t ps_discrete_dim( + const Topology& topology, + PhaseSpaceMapping::TChannelMode mode, + const std::optional>& color_order +) { + if (mode == PhaseSpaceMapping::color_ordered && + topology.t_propagator_count() >= 1) { + ColorOrderedMapping co(ps_chain_order(topology, color_order)); + return co.discrete_dim(); + } + return 0; +} +} // namespace + PhaseSpaceMapping::PhaseSpaceMapping( const Topology& topology, double cm_energy, @@ -88,14 +125,29 @@ PhaseSpaceMapping::PhaseSpaceMapping( double invariant_power, TChannelMode t_channel_mode, const std::optional& cuts, - const std::vector>& permutations + const std::vector>& permutations, + const std::optional>& color_order ) : Mapping( "PhaseSpaceMapping", - {{"random", - batch_float_array( - 3 * topology.outgoing_masses().size() - (leptonic ? 4 : 2) - )}}, + [&] { + NamedVector in{ + {"random", + batch_float_array( + 3 * topology.outgoing_masses().size() - (leptonic ? 4 : 2) + )} + }; + // Opt-in discrete channel: only declared when the t-channel strategy + // actually has discrete two-solution choices (color_ordered). + std::size_t nd = ps_discrete_dim(topology, t_channel_mode, color_order); + if (nd > 0) { + in.push_back( + "discrete", + Type{DataType::dt_int, batch_size, {static_cast(nd)}} + ); + } + return in; + }(), {{"momenta", batch_four_vec_array(topology.outgoing_masses().size() + 2)}, {"x1", batch_float}, {"x2", batch_float}}, @@ -178,26 +230,86 @@ PhaseSpaceMapping::PhaseSpaceMapping( double s_hat_min = std::max(total_mass * total_mass, sqrt_s_hat_min * sqrt_s_hat_min); if (has_t_channel) { + // Per-child pt_min (and eta_max), ordered to match the mass conditions + // handed to the t-channel mapping (leaf children carry their pt cut; + // composite children were reset to 0 above). + std::vector eta_max, pt_min; + for (std::size_t index : topology.decays().at(0).child_indices) { + auto& info = decay_info.at(index); + eta_max.push_back(info.eta_max); + pt_min.push_back(info.pt_min); + } if (t_channel_mode == PhaseSpaceMapping::chili) { // |y| <= |eta|, so we can pass y_max = eta_max - std::vector eta_max, pt_min; - for (std::size_t index : topology.decays().at(0).child_indices) { - auto& info = decay_info.at(index); - eta_max.push_back(info.eta_max); - pt_min.push_back(info.pt_min); - } _t_mapping = ChiliMapping(_topology.t_propagator_count() + 1, eta_max, pt_min); + } else if (t_channel_mode == PhaseSpaceMapping::color_ordered) { + // color_order is optional in general but REQUIRED here: the chain is + // built in the externally supplied color order so the t-channel + // topology matches the known color structure of the process. + if (!color_order) { + throw std::invalid_argument( + "PhaseSpaceMapping: color_ordered mode requires a color_order" + ); + } + // Reorder the per-pair cut matrices (indexed by raw outgoing index) + // into the child order in which masses/pt are handed to the chain, + // mirroring the pt_min reordering above. Composite (non-leaf) + // children carry no pairwise cut. + const auto& out_idx = topology.outgoing_indices(); + const auto& child_indices = topology.decays().at(0).child_indices; + std::vector child_to_out( + child_indices.size(), std::numeric_limits::max() + ); + for (std::size_t a = 0; a < child_indices.size(); ++a) { + auto it = + std::find(out_idx.begin(), out_idx.end(), child_indices.at(a)); + if (it != out_idx.end()) { + child_to_out.at(a) = std::distance(out_idx.begin(), it); + } + ++a; + } + auto m_inv_full = _cuts.m_inv_min(); + auto dr_full = _cuts.dr_min(); + std::size_t nc = child_to_out.size(); + std::vector> m_inv_co(nc, std::vector(nc, 0.)); + std::vector> dr_co(nc, std::vector(nc, 0.)); + for (std::size_t a = 0; a < nc; ++a) { + if (child_to_out.at(a) >= m_inv_full.size()) { + continue; + } + for (std::size_t b = 0; b < nc; ++b) { + if (child_to_out.at(b) >= m_inv_full.size()) { + continue; + } + m_inv_co.at(a).at(b) = + m_inv_full.at(child_to_out.at(a)).at(child_to_out.at(b)); + dr_co.at(a).at(b) = + dr_full.at(child_to_out.at(a)).at(child_to_out.at(b)); + } + } + _t_mapping = ColorOrderedMapping( + ps_chain_order(topology, color_order), + invariant_power, + invariant_power, + pt_min, + m_inv_co, + dr_co + ); } else if (t_channel_mode == PhaseSpaceMapping::propagator || topology.t_propagator_count() < 2) { - _t_mapping = - TPropagatorMapping(_topology.t_integration_order(), invariant_power); + _t_mapping = TPropagatorMapping( + _topology.t_integration_order(), invariant_power, pt_min + ); } else if (t_channel_mode == PhaseSpaceMapping::rambo) { // TODO: add massless special case _t_mapping = FastRamboMapping(_topology.t_propagator_count() + 1, false); } } + // Random-number budget: identical computation to the declared input shape. + _n_discrete = ps_discrete_dim(_topology, t_channel_mode, color_order); + for (auto& perm : permutations) { _permutations.emplace_back(perm.begin(), perm.end()); } @@ -209,7 +321,8 @@ PhaseSpaceMapping::PhaseSpaceMapping( bool leptonic, double invariant_power, TChannelMode mode, - const std::optional& cuts + const std::optional& cuts, + const std::optional>& color_order ) : PhaseSpaceMapping( Topology([&] { @@ -246,7 +359,9 @@ PhaseSpaceMapping::PhaseSpaceMapping( leptonic, invariant_power, mode, - cuts + cuts, + {}, + color_order ) {} Mapping::Result PhaseSpaceMapping::build_forward_impl( @@ -257,6 +372,15 @@ Mapping::Result PhaseSpaceMapping::build_forward_impl( auto random_numbers = fb.unstack(inputs.at(0)); auto r = random_numbers.begin(); auto next_random = [&]() { return *(r++); }; + // Opt-in discrete channel: present as inputs.at(1) only when the t-channel + // strategy declared discrete choices (color_ordered). These are passed + // through to the t-channel mapping after its continuous randoms. + ValueVec discrete_numbers; + if (inputs.size() > 1) { + discrete_numbers = fb.unstack(inputs.at(1)); + } + auto d = discrete_numbers.begin(); + auto next_discrete = [&]() { return *(d++); }; ValueVec dets{_pi_factors}; Value x1 = 1.0, x2 = 1.0; @@ -319,6 +443,11 @@ Mapping::Result PhaseSpaceMapping::build_forward_impl( for (std::size_t i = 0; i < t_mapping.random_dim(); ++i) { args.push_back(next_random()); } + // Discrete choices follow the continuous randoms, matching the + // t-channel mapping's input_types order [random..., discrete...]. + for (std::size_t j = 0; j < t_mapping.discrete_dim(); ++j) { + args.push_back(next_discrete()); + } conds.push_back(sqrt_s_hat); for (std::size_t index : decay_data.at(0).decay.child_indices) { conds.push_back(decay_data.at(index).mass.value()); @@ -453,6 +582,7 @@ Mapping::Result PhaseSpaceMapping::build_inverse_impl( // go through decays and recover random numbers from momenta ValueVec random_out_reversed; + ValueVec discrete_out; ValueVec dets{1. / _pi_factors}; for (std::size_t decay_map_index = 0; auto& data : std::views::reverse(decay_data)) { @@ -510,6 +640,11 @@ Mapping::Result PhaseSpaceMapping::build_inverse_impl( t_result.rend() - t_mapping.random_dim(), t_result.rend() ); + // Discrete choices sit at forward positions [nc, nc+nd) in + // t_result (after the continuous randoms, before "det"). + for (std::size_t j = 0; j < t_mapping.discrete_dim(); ++j) { + discrete_out.push_back(t_result.at(t_mapping.random_dim() + j)); + } dets.push_back(t_result["det"]); }, [&](std::monostate) {} @@ -553,5 +688,9 @@ Mapping::Result PhaseSpaceMapping::build_inverse_impl( random_out.insert( random_out.end(), random_out_reversed.rbegin(), random_out_reversed.rend() ); - return {{{"random", fb.stack(random_out)}}, fb.product(dets)}; + NamedVector result{{"random", fb.stack(random_out)}}; + if (!discrete_out.empty()) { + result.push_back("discrete", fb.stack(discrete_out)); + } + return {result, fb.product(dets)}; } diff --git a/madspace/src/phasespace/t_propagator_mapping.cpp b/madspace/src/phasespace/t_propagator_mapping.cpp index 3d8a23a2b..af1dc70b2 100644 --- a/madspace/src/phasespace/t_propagator_mapping.cpp +++ b/madspace/src/phasespace/t_propagator_mapping.cpp @@ -4,8 +4,25 @@ using namespace madspace; +double TPropagatorMapping::pt2(std::size_t i) const { + double p = (i < _pt_min.size()) ? _pt_min.at(i) : 0.0; + return p * p; +} + +// Only engage the cut kernel when there is an actual pt cut. +static bool has_pt_cut(const std::vector& pt_min) { + for (double p : pt_min) { + if (p > 0.0) { + return true; + } + } + return false; +} + TPropagatorMapping::TPropagatorMapping( - const std::vector& integration_order, double invariant_power + const std::vector& integration_order, + double invariant_power, + const std::vector& pt_min ) : Mapping( "TPropagatorMapping", @@ -32,8 +49,10 @@ TPropagatorMapping::TPropagatorMapping( }() ), _integration_order(integration_order), - _com_scattering(true, invariant_power), - _lab_scattering(false, invariant_power) { + _pt_min(pt_min), + _has_cut(has_pt_cut(pt_min)), + _com_scattering(true, invariant_power, 0., 0., has_pt_cut(pt_min)), + _lab_scattering(false, invariant_power, 0., 0., has_pt_cut(pt_min)) { std::size_t next_index_low = 0; std::size_t next_index_high = integration_order.size() - 1; for (std::size_t index : integration_order) { @@ -99,6 +118,17 @@ Mapping::Result TPropagatorMapping::build_forward_impl( p_ext.at(1) = p2; auto p1_rest = p1, p2_rest = p2; + auto etmin_particle = [&](std::size_t j) { + return fb.sqrt(fb.add(fb.square(m_out.at(j)), Value(pt2(j)))); + }; + Value total_etmin = Value(0.0), running_etmin = Value(0.0); + if (_has_cut) { + total_etmin = etmin_particle(0); + for (std::size_t j = 1; j < m_out.size(); ++j) { + total_etmin = fb.add(total_etmin, etmin_particle(j)); + } + } + // sample t-invariants and build momenta of t-channel part of the diagram Value k_rest; bool first = true; @@ -108,10 +138,15 @@ Mapping::Result TPropagatorMapping::build_forward_impl( first = false; std::size_t sampled_index = index + side; auto mass = m_out.at(sampled_index); + ValueVec cond{side ? p1_rest : p2_rest, side ? p2_rest : p1_rest}; + if (_has_cut) { + Value etmin_peeled = etmin_particle(sampled_index); + running_etmin = fb.add(running_etmin, etmin_peeled); + cond.push_back(fb.sub(total_etmin, running_etmin)); // etmin_1 (recoil) + cond.push_back(etmin_peeled); // etmin_2 (peeled) + } auto ks = scattering.build_forward( - fb, - {next_random(), next_random(), mass_sum, mass}, - {side ? p1_rest : p2_rest, side ? p2_rest : p1_rest} + fb, {next_random(), next_random(), mass_sum, mass}, cond ); k_rest = ks.at(0); auto k = ks.at(1); @@ -181,6 +216,18 @@ Mapping::Result TPropagatorMapping::build_inverse_impl( } } + // ETmin per particle (see forward); Only built when cuts are active. + auto etmin_particle = [&](std::size_t j) { + return fb.sqrt(fb.add(fb.square(m_out.at(j)), Value(pt2(j)))); + }; + Value total_etmin = Value(0.0), running_etmin = Value(0.0); + if (_has_cut) { + total_etmin = etmin_particle(0); + for (std::size_t j = 1; j < m_out.size(); ++j) { + total_etmin = fb.add(total_etmin, etmin_particle(j)); + } + } + // sample t-invariants and build momenta of t-channel part of the diagram Value k_rest = fb.add(inputs.at(0), inputs.at(1)); Value p1_rest = inputs.at(0), p2_rest = inputs.at(1); @@ -192,9 +239,14 @@ Mapping::Result TPropagatorMapping::build_inverse_impl( auto mass = m_out.at(sampled_index); auto k = inputs.at(sampled_index + 2); k_rest = fb.sub(k_rest, k); - auto rs = scattering.build_inverse( - fb, {k_rest, k}, {side ? p1_rest : p2_rest, side ? p2_rest : p1_rest} - ); + ValueVec cond{side ? p1_rest : p2_rest, side ? p2_rest : p1_rest}; + if (_has_cut) { + Value etmin_peeled = etmin_particle(sampled_index); + running_etmin = fb.add(running_etmin, etmin_peeled); + cond.push_back(fb.sub(total_etmin, running_etmin)); // etmin_1 (recoil) + cond.push_back(etmin_peeled); // etmin_2 (peeled) + } + auto rs = scattering.build_inverse(fb, {k_rest, k}, cond); random_out.push_back(rs.at(0)); random_out.push_back(rs.at(1)); dets.push_back(rs["det"]); diff --git a/madspace/src/phasespace/three_particle.cpp b/madspace/src/phasespace/three_particle.cpp index d94efa7fd..4528619c0 100644 --- a/madspace/src/phasespace/three_particle.cpp +++ b/madspace/src/phasespace/three_particle.cpp @@ -95,38 +95,66 @@ TwoToThreeParticleScattering::TwoToThreeParticleScattering( double t_width, double s_invariant_power, double s_mass, - double s_width + double s_width, + bool has_cut ) : Mapping( "TwoToThreeParticleScattering", - {{"random_choice", batch_float}, + {{"discrete_choice", batch_int}, {"random_s23", batch_float}, {"random_t1", batch_float}, {"mass1", batch_float}, {"mass2", batch_float}}, {{"momentum1", batch_four_vec}, {"momentum2", batch_four_vec}}, - {{"momentum_in1", batch_four_vec}, - {"momentum_in2", batch_four_vec}, - {"momentum3", batch_four_vec}} + [&] { + NamedVector cond{ + {"momentum_in1", batch_four_vec}, + {"momentum_in2", batch_four_vec}, + {"momentum3", batch_four_vec} + }; + if (has_cut) { + cond.push_back("etmin_1", batch_float); + cond.push_back("etmin_2", batch_float); + cond.push_back("drcut", batch_float); + cond.push_back("s23_min_cut", batch_float); + } + return cond; + }() ), _t_invariant(t_invariant_power, t_mass, t_width), - _s_invariant(s_invariant_power, s_mass, s_width) {} + _s_invariant(s_invariant_power, s_mass, s_width), + _has_cut(has_cut) {} Mapping::Result TwoToThreeParticleScattering::build_forward_impl( FunctionBuilder& fb, const NamedVector& inputs, const NamedVector& conditions ) const { - auto r_choice = inputs.at(0), r_s23 = inputs.at(1), r_t1 = inputs.at(2), + auto index_choice = inputs.at(0), r_s23 = inputs.at(1), r_t1 = inputs.at(2), m1 = inputs.at(3), m2 = inputs.at(4); auto p_a = conditions.at(0), p_b = conditions.at(1), p_3 = conditions.at(2); - auto [t1_min, t1_max] = fb.t_inv_min_max(p_a, fb.sub(p_b, p_3), m1, m2); + auto [t1_min, t1_max] = _has_cut + ? fb.t_inv_min_max_cut( + p_a, fb.sub(p_b, p_3), m1, m2, conditions.at(3), conditions.at(4) + ) + : fb.t_inv_min_max(p_a, fb.sub(p_b, p_3), m1, m2); auto t_inv_result = _t_invariant.build_forward(fb, {r_t1}, {t1_min, t1_max}); - auto [s23_min, s23_max] = - fb.s23_min_max(p_a, p_b, p_3, t_inv_result["invariant"], m1, m2); + auto [s23_min, s23_max] = _has_cut + ? fb.s23_min_max_cut( + p_a, + p_b, + p_3, + t_inv_result["invariant"], + m1, + m2, + conditions.at(3), + conditions.at(4), + conditions.at(5), + conditions.at(6) + ) + : fb.s23_min_max(p_a, p_b, p_3, t_inv_result["invariant"], m1, m2); auto s23_inv_result = _s_invariant.build_forward(fb, {r_s23}, {s23_min, s23_max}); auto det_inv = fb.mul(t_inv_result["det"], s23_inv_result["det"]); - auto [index_choice, index_det] = fb.sample_discrete(r_choice, 2); auto [p1, p2, det_scatter] = fb.two_to_three_particle_scattering( index_choice, p_a, @@ -137,8 +165,7 @@ Mapping::Result TwoToThreeParticleScattering::build_forward_impl( m1, m2 ); - auto det_scatter_23 = fb.mul(index_det, det_scatter); - return {{{"momentum1", p1}, {"momentum2", p2}}, fb.mul(det_inv, det_scatter_23)}; + return {{{"momentum1", p1}, {"momentum2", p2}}, fb.mul(det_inv, det_scatter)}; } Mapping::Result TwoToThreeParticleScattering::build_inverse_impl( @@ -148,23 +175,36 @@ Mapping::Result TwoToThreeParticleScattering::build_inverse_impl( ) const { auto p1 = inputs.at(0), p2 = inputs.at(1); auto p_a = conditions.at(0), p_b = conditions.at(1), p_3 = conditions.at(2); - auto [t1_abs, t1_min, t1_max] = - fb.t_inv_value_and_min_max(p_a, fb.sub(p_b, p_3), p1, p2); + auto [t1_abs, t1_min, t1_max] = _has_cut + ? fb.t_inv_value_and_min_max_cut( + p_a, fb.sub(p_b, p_3), p1, p2, conditions.at(3), conditions.at(4) + ) + : fb.t_inv_value_and_min_max(p_a, fb.sub(p_b, p_3), p1, p2); auto t_inv_result = _t_invariant.build_inverse(fb, {t1_abs}, {t1_min, t1_max}); - auto [s23, s23_min, s23_max] = - fb.s23_value_and_min_max(p_a, p_b, p_3, t1_abs, p1, p2); + auto [s23, s23_min, s23_max] = _has_cut + ? fb.s23_value_and_min_max_cut( + p_a, + p_b, + p_3, + t1_abs, + p1, + p2, + conditions.at(3), + conditions.at(4), + conditions.at(5), + conditions.at(6) + ) + : fb.s23_value_and_min_max(p_a, p_b, p_3, t1_abs, p1, p2); auto s23_inv_result = _s_invariant.build_inverse(fb, {s23}, {s23_min, s23_max}); auto det_inv = fb.mul(t_inv_result["det"], s23_inv_result["det"]); auto [m1, m2, index_choice, det_scatter] = fb.two_to_three_particle_scattering_inverse(p1, p2, p_3, p_a, p_b, t1_abs, s23); - auto [r_choice, index_det] = fb.sample_discrete_inverse(index_choice, 2); - auto det_scatter_23 = fb.mul(index_det, det_scatter); return { - {{"random_choice", r_choice}, + {{"discrete_choice", index_choice}, {"random_s23", s23_inv_result["random"]}, {"random_t1", t_inv_result["random"]}, {"mass1", m1}, {"mass2", m2}}, - fb.mul(det_inv, det_scatter_23) + fb.mul(det_inv, det_scatter) }; } diff --git a/madspace/src/phasespace/two_particle.cpp b/madspace/src/phasespace/two_particle.cpp index 8e8acf56e..4e118b66e 100644 --- a/madspace/src/phasespace/two_particle.cpp +++ b/madspace/src/phasespace/two_particle.cpp @@ -69,7 +69,7 @@ Mapping::Result TwoBodyDecay::build_inverse_impl( } TwoToTwoParticleScattering::TwoToTwoParticleScattering( - bool com, double invariant_power, double mass, double width + bool com, double invariant_power, double mass, double width, bool has_cut ) : Mapping( "TwoToTwoParticleScattering", @@ -78,10 +78,20 @@ TwoToTwoParticleScattering::TwoToTwoParticleScattering( {"mass1", batch_float}, {"mass2", batch_float}}, {{"momentum1", batch_four_vec}, {"momentum2", batch_four_vec}}, - {{"momentum_in1", batch_four_vec}, {"momentum_in2", batch_four_vec}} + [&] { + NamedVector cond{ + {"momentum_in1", batch_four_vec}, {"momentum_in2", batch_four_vec} + }; + if (has_cut) { + cond.push_back("etmin_1", batch_float); + cond.push_back("etmin_2", batch_float); + } + return cond; + }() ), _com(com), - _invariant(invariant_power, mass, width) {} + _invariant(invariant_power, mass, width), + _has_cut(has_cut) {} Mapping::Result TwoToTwoParticleScattering::build_forward_impl( FunctionBuilder& fb, @@ -91,7 +101,9 @@ Mapping::Result TwoToTwoParticleScattering::build_forward_impl( auto r_phi = inputs.at(0), r_inv = inputs.at(1), m1 = inputs.at(2), m2 = inputs.at(3); auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); - auto [t_min, t_max] = fb.t_inv_min_max(p_in1, p_in2, m1, m2); + auto [t_min, t_max] = _has_cut + ? fb.t_inv_min_max_cut(p_in1, p_in2, m1, m2, conditions.at(2), conditions.at(3)) + : fb.t_inv_min_max(p_in1, p_in2, m1, m2); auto t_result = _invariant.build_forward(fb, {r_inv}, {t_min, t_max}); auto [p1, p2, det_scatter] = _com ? fb.two_to_two_particle_scattering_com( @@ -112,7 +124,11 @@ Mapping::Result TwoToTwoParticleScattering::build_inverse_impl( ) const { auto p1 = inputs.at(0), p2 = inputs.at(1); auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); - auto [t_abs, t_min, t_max] = fb.t_inv_value_and_min_max(p_in1, p_in2, p1, p2); + auto [t_abs, t_min, t_max] = _has_cut + ? fb.t_inv_value_and_min_max_cut( + p_in1, p_in2, p1, p2, conditions.at(2), conditions.at(3) + ) + : fb.t_inv_value_and_min_max(p_in1, p_in2, p1, p2); auto t_result = _invariant.build_inverse(fb, {t_abs}, {t_min, t_max}); auto [r_phi, m1, m2, det_scatter] = _com ? fb.two_to_two_particle_scattering_com_inverse(p1, p2, p_in1, p_in2) @@ -132,7 +148,8 @@ DoubleT::DoubleT( double t1_width, double t2_invariant_power, double t2_mass, - double t2_width + double t2_width, + bool has_cut ) : Mapping( "DoubleT", @@ -140,13 +157,23 @@ DoubleT::DoubleT( {"random_t1", batch_float}, {"random_t2", batch_float}}, {{"momentum1", batch_four_vec}, {"momentum2", batch_four_vec}}, - {{"momentum_in1", batch_four_vec}, - {"momentum_in2", batch_four_vec}, - {"mass1", batch_float}, - {"mass_rest_min", batch_float}} + [&] { + NamedVector cond{ + {"momentum_in1", batch_four_vec}, + {"momentum_in2", batch_four_vec}, + {"mass1", batch_float}, + {"mass_rest_min", batch_float} + }; + if (has_cut) { + cond.push_back("etmin_i", batch_float); + cond.push_back("etmin_ir", batch_float); + } + return cond; + }() ), _t1_invariant(t1_invariant_power, t1_mass, t1_width), - _t2_invariant(t2_invariant_power, t2_mass, t2_width) {} + _t2_invariant(t2_invariant_power, t2_mass, t2_width), + _has_cut(has_cut) {} Mapping::Result DoubleT::build_forward_impl( FunctionBuilder& fb, @@ -157,12 +184,16 @@ Mapping::Result DoubleT::build_forward_impl( auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); auto m1 = conditions.at(2); auto mir_min = conditions.at(3); + auto etmin_i = _has_cut ? conditions.at(4) : Value(0.); + auto etmin_ir = _has_cut ? conditions.at(5) : Value(0.); - auto [t1_min, t1_max] = fb.t1_inv_min_max_doublet(p_in1, p_in2, m1, mir_min); + auto [t1_min, t1_max] = + fb.t1_inv_min_max_doublet(p_in1, p_in2, m1, mir_min, etmin_i, etmin_ir); auto t1_result = _t1_invariant.build_forward(fb, {r_t1}, {t1_min, t1_max}); - auto [t2_min, t2_max] = - fb.t2_inv_min_max_doublet(p_in1, p_in2, m1, mir_min, t1_result["invariant"]); + auto [t2_min, t2_max] = fb.t2_inv_min_max_doublet( + p_in1, p_in2, m1, mir_min, t1_result["invariant"], etmin_i, etmin_ir + ); auto t2_result = _t2_invariant.build_forward(fb, {r_t2}, {t2_min, t2_max}); auto [p1, p2, det_scatter] = fb.double_t_scattering( @@ -182,15 +213,19 @@ Mapping::Result DoubleT::build_inverse_impl( auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); auto m1 = conditions.at(2); auto mir_min = conditions.at(3); + auto etmin_i = _has_cut ? conditions.at(4) : Value(0.); + auto etmin_ir = _has_cut ? conditions.at(5) : Value(0.); auto [r_phi, det_scatter] = fb.double_t_scattering_inverse(p1, p2, p_in1, p_in2); - auto [t1_abs, t1_min, t1_max] = - fb.t1_inv_value_and_min_max_doublet(p_in1, p_in2, p1, m1, mir_min); + auto [t1_abs, t1_min, t1_max] = fb.t1_inv_value_and_min_max_doublet( + p_in1, p_in2, p1, m1, mir_min, etmin_i, etmin_ir + ); auto t1_result = _t1_invariant.build_inverse(fb, {t1_abs}, {t1_min, t1_max}); - auto [t2_abs, t2_min, t2_max] = - fb.t2_inv_value_and_min_max_doublet(p_in1, p_in2, p1, m1, mir_min, t1_abs); + auto [t2_abs, t2_min, t2_max] = fb.t2_inv_value_and_min_max_doublet( + p_in1, p_in2, p1, m1, mir_min, t1_abs, etmin_i, etmin_ir + ); auto t2_result = _t2_invariant.build_inverse(fb, {t2_abs}, {t2_min, t2_max}); auto det_inv = fb.mul(t1_result["det"], t2_result["det"]); diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index 1cd7bb9ec..c57544c09 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -82,12 +82,16 @@ void add_instructions(py::classh& fb) { fb.def("three_body_decay_inverse", &FunctionBuilder::three_body_decay_inverse, py::arg("p1"), py::arg("p2"), py::arg("p3")); fb.def("t_inv_min_max", &FunctionBuilder::t_inv_min_max, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("m2")); fb.def("t_inv_value_and_min_max", &FunctionBuilder::t_inv_value_and_min_max, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("p2")); - fb.def("t1_inv_min_max_doublet", &FunctionBuilder::t1_inv_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("mir_min")); - fb.def("t1_inv_value_and_min_max_doublet", &FunctionBuilder::t1_inv_value_and_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("m1"), py::arg("mir_min")); - fb.def("t2_inv_min_max_doublet", &FunctionBuilder::t2_inv_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("mir_min"), py::arg("t1_abs")); - fb.def("t2_inv_value_and_min_max_doublet", &FunctionBuilder::t2_inv_value_and_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("m1"), py::arg("mir_min"), py::arg("t1_abs")); + fb.def("t_inv_min_max_cut", &FunctionBuilder::t_inv_min_max_cut, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("m2"), py::arg("etmin_1"), py::arg("etmin_2")); + fb.def("t_inv_value_and_min_max_cut", &FunctionBuilder::t_inv_value_and_min_max_cut, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("p2"), py::arg("etmin_1"), py::arg("etmin_2")); + fb.def("t1_inv_min_max_doublet", &FunctionBuilder::t1_inv_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("mir_min"), py::arg("etmin_i"), py::arg("etmin_ir")); + fb.def("t1_inv_value_and_min_max_doublet", &FunctionBuilder::t1_inv_value_and_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("m1"), py::arg("mir_min"), py::arg("etmin_i"), py::arg("etmin_ir")); + fb.def("t2_inv_min_max_doublet", &FunctionBuilder::t2_inv_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("m1"), py::arg("mir_min"), py::arg("t1_abs"), py::arg("etmin_i"), py::arg("etmin_ir")); + fb.def("t2_inv_value_and_min_max_doublet", &FunctionBuilder::t2_inv_value_and_min_max_doublet, py::arg("pa"), py::arg("pb"), py::arg("p1"), py::arg("m1"), py::arg("mir_min"), py::arg("t1_abs"), py::arg("etmin_i"), py::arg("etmin_ir")); fb.def("s23_min_max", &FunctionBuilder::s23_min_max, py::arg("pa"), py::arg("pb"), py::arg("p3"), py::arg("t1_abs"), py::arg("m1"), py::arg("m2")); fb.def("s23_value_and_min_max", &FunctionBuilder::s23_value_and_min_max, py::arg("pa"), py::arg("pb"), py::arg("p3"), py::arg("t1_abs"), py::arg("p1"), py::arg("p2")); + fb.def("s23_min_max_cut", &FunctionBuilder::s23_min_max_cut, py::arg("pa"), py::arg("pb"), py::arg("p3"), py::arg("t1_abs"), py::arg("m1"), py::arg("m2"), py::arg("etmin_1"), py::arg("etmin_2"), py::arg("drcut"), py::arg("s23_min_cut")); + fb.def("s23_value_and_min_max_cut", &FunctionBuilder::s23_value_and_min_max_cut, py::arg("pa"), py::arg("pb"), py::arg("p3"), py::arg("t1_abs"), py::arg("p1"), py::arg("p2"), py::arg("etmin_1"), py::arg("etmin_2"), py::arg("drcut"), py::arg("s23_min_cut")); fb.def("invariants_from_momenta", &FunctionBuilder::invariants_from_momenta, py::arg("p_ext"), py::arg("factors")); fb.def("sde2_channel_weights", &FunctionBuilder::sde2_channel_weights, py::arg("invariants"), py::arg("masses"), py::arg("widths"), py::arg("indices")); fb.def("subchannel_weights", &FunctionBuilder::subchannel_weights, py::arg("invariants"), py::arg("masses"), py::arg("widths"), py::arg("indices"), py::arg("on_shell"), py::arg("group_sizes")); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 3d3313c8d..7cbe5b09d 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -432,11 +432,12 @@ PYBIND11_MODULE(_madspace_py, m) { py::classh(m, "TwoToTwoParticleScattering") .def( - py::init(), + py::init(), py::arg("com"), py::arg("invariant_power") = 0., py::arg("mass") = 0., - py::arg("width") = 0. + py::arg("width") = 0., + py::arg("has_cut") = false ); py::classh(m, "DoubleT") @@ -455,13 +456,14 @@ PYBIND11_MODULE(_madspace_py, m) { py::classh(m, "TwoToThreeParticleScattering") .def( - py::init(), + py::init(), py::arg("t_invariant_power") = 0., py::arg("t_mass") = 0., py::arg("t_width") = 0., py::arg("s_invariant_power") = 0., py::arg("s_mass") = 0., - py::arg("s_width") = 0. + py::arg("s_width") = 0., + py::arg("has_cut") = false ); py::classh(m, "Propagator") @@ -483,20 +485,31 @@ PYBIND11_MODULE(_madspace_py, m) { py::classh(m, "TPropagatorMapping") .def( - py::init, double>(), + py::init, double, std::vector>(), py::arg("integration_order"), - py::arg("invariant_power") = 0. + py::arg("invariant_power") = 0., + py::arg("pt_min") = std::vector{} ) .def("random_dim", &TPropagatorMapping::random_dim); py::classh(m, "ColorOrderedMapping") .def( - py::init, double, double>(), + py::init< + std::vector, + double, + double, + std::vector, + std::vector>, + std::vector>>(), py::arg("color_order"), py::arg("t_invariant_power") = 0., - py::arg("s_invariant_power") = 0. + py::arg("s_invariant_power") = 0., + py::arg("pt_min") = std::vector{}, + py::arg("m_inv_min") = std::vector>{}, + py::arg("dr_min") = std::vector>{} ) - .def("random_dim", &ColorOrderedMapping::random_dim); + .def("random_dim", &ColorOrderedMapping::random_dim) + .def("discrete_dim", &ColorOrderedMapping::discrete_dim); py::classh(m, "VegasHistogram") .def( @@ -602,7 +615,9 @@ PYBIND11_MODULE(_madspace_py, m) { .def(py::init(), py::arg("particle_count")) .def("sqrt_s_min", &Cuts::sqrt_s_min) .def("eta_max", &Cuts::eta_max) - .def("pt_min", &Cuts::pt_min); + .def("pt_min", &Cuts::pt_min) + .def("m_inv_min", &Cuts::m_inv_min) + .def("dr_min", &Cuts::dr_min); py::classh(m, "HistItem") .def( @@ -685,6 +700,7 @@ PYBIND11_MODULE(_madspace_py, m) { {"propagator", PhaseSpaceMapping::propagator}, {"rambo", PhaseSpaceMapping::rambo}, {"chili", PhaseSpaceMapping::chili}, + {"color_ordered", PhaseSpaceMapping::color_ordered}, } ); psmap @@ -696,14 +712,16 @@ PYBIND11_MODULE(_madspace_py, m) { double, PhaseSpaceMapping::TChannelMode, const std::optional&, - const nested_vector2&>(), + const nested_vector2&, + const std::optional>&>(), py::arg("topology"), py::arg("cm_energy"), py::arg("leptonic") = false, py::arg("invariant_power") = 0.8, py::arg("t_channel_mode") = PhaseSpaceMapping::propagator, py::arg("cuts") = std::nullopt, - py::arg("permutations") = std::vector{} + py::arg("permutations") = std::vector{}, + py::arg("color_order") = std::nullopt ) .def( py::init< @@ -712,15 +730,18 @@ PYBIND11_MODULE(_madspace_py, m) { bool, double, PhaseSpaceMapping::TChannelMode, - std::optional>(), + std::optional, + const std::optional>&>(), py::arg("masses"), py::arg("cm_energy"), py::arg("leptonic") = false, py::arg("invariant_power") = 0.8, py::arg("mode") = PhaseSpaceMapping::rambo, - py::arg("cuts") = std::nullopt + py::arg("cuts") = std::nullopt, + py::arg("color_order") = std::nullopt ) .def("random_dim", &PhaseSpaceMapping::random_dim) + .def("discrete_dim", &PhaseSpaceMapping::discrete_dim) .def("particle_count", &PhaseSpaceMapping::particle_count) .def("channel_count", &PhaseSpaceMapping::channel_count); diff --git a/madspace/tests/test_2to3_scattering.py b/madspace/tests/test_2to3_scattering.py index dfa52bed6..9885a0e3b 100644 --- a/madspace/tests/test_2to3_scattering.py +++ b/madspace/tests/test_2to3_scattering.py @@ -96,8 +96,8 @@ def fixed_input_points(rng, request): r2 = rng.random(N) p12, p3, det_22 = map_22.map_forward([r1, r2, m12, m3], [pA, pB]) - # Randoms for the 2->3 mapper - r_choice = rng.random(N) + # Discrete two-solution choice (int 0/1) + continuous randoms for the 2->3 + r_choice = rng.integers(0, 2, size=N).astype(np.int32) r_s23 = rng.random(N) r_t1 = rng.random(N) @@ -150,7 +150,7 @@ def input_points(rng, request): p12, p3, det_22 = map_22.map_forward([r1, r2, m12, m3], [pa, pb]) # Randoms for the 2->3 mapper - r_choice = rng.random(N) # decide branch (emitter choice) + r_choice = rng.integers(0, 2, size=N).astype(np.int32) r_s23 = rng.random(N) r_t1 = rng.random(N) @@ -204,7 +204,10 @@ def test_inverse(input_points): for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): if i == 0: - assert ((inp < 0.5) == (inv_inp < 0.5)).all() + # discrete two-solution choice: recovered exactly as an int + assert ( + np.asarray(inv_inp).astype(np.int64) == np.asarray(inp).astype(np.int64) + ).all() continue assert inp == approx(inv_inp), f"mismatch in input index {i}" @@ -254,6 +257,10 @@ def test_phase_space_compare(rng, input_points): p1, p2, det23 = mapping23.map_forward(inputs, conditions) p1s, p2s, det22 = mapping22.map_forward(inputs22, conditions22) + # det23 is now the per-branch Jacobian; the 2-solution multiplicity is owned + # externally, so apply it here to compare against the 2->2 phase-space element. + det23 = det23 * 2.0 + # Outgoing masses must match m1, m2; spectator stays whatever it was. std_error_23 = np.std(det23) / np.sqrt(N) assert np.mean(det23) == approx(np.mean(det22), abs=3 * std_error_23, rel=1e-6) @@ -273,6 +280,8 @@ def test_phase_space_volume(fixed_input_points): conditions = [fixed_input_points.pa, fixed_input_points.pb, fixed_input_points.p3] p1, p2, det = mapping23.map_forward(inputs, conditions) + # per-branch Jacobian; apply the 2-solution multiplicity externally + det = det * 2.0 s = fixed_input_points.m12**2 m1_2 = fixed_input_points.m1**2 diff --git a/madspace/tests/test_color_ordered_mapping.py b/madspace/tests/test_color_ordered_mapping.py deleted file mode 100644 index c64249254..000000000 --- a/madspace/tests/test_color_ordered_mapping.py +++ /dev/null @@ -1,212 +0,0 @@ -import numpy as np -import pytest -from pytest import approx - -import madspace as ms - -# ---------------------------- -# Fixtures -# ---------------------------- - - -@pytest.fixture -def rng(): - return np.random.default_rng(1234) - - -N = 10_000 # keep this moderate for CI speed; bump locally for tighter stats -CM_ENERGY = 13000.0 - - -# Mass set families. n_out is filled in at use. -def _massless(n_out): - return [0.0] * n_out, f"{n_out} massless" - - -def _w_like(n_out): - return [80.0] * n_out, f"{n_out} W-like" - - -def _mixed_top(n_out): - # one heavy, one W-like, rest massless; only meaningful for n_out >= 3. - return [173.0, 80.0] + [0.0] * (n_out - 2), f"{n_out} mixed" - - -def _mass_sets_for(n_out): - sets = [_massless(n_out), _w_like(n_out)] - if n_out >= 3: - sets.append(_mixed_top(n_out)) - return sets - - -# Each entry: (color_order, int_order_for_old_class, topology_label). -TOPOLOGIES = [ - ([0, 2, 3, 1, 4], [0, 1], "n=5: set1={2,3}, set2={4}"), - ([0, 1, 2, 3, 4], [0, 1], "n=5: set1={}, set2={2,3,4}"), - ([0, 2, 3, 4, 1, 5], [0, 1, 2], "n=6: set1={2,3,4}, set2={5}"), - ([0, 2, 3, 1, 4, 5], [0, 1, 2], "n=6: set1={2,3}, set2={4,5}"), - ([0, 1, 2, 3, 4, 5], [0, 1, 2], "n=6: set1={}, set2={2,3,4,5}"), - ([0, 2, 3, 4, 1, 5, 6], [0, 1, 2, 3], "n=7: set1={2,3,4}, set2={5,6}"), - ([0, 2, 3, 4, 5, 1, 6], [0, 1, 2, 3], "n=7: set1={2,3,4,5}, set2={6}"), - ([0, 2, 3, 4, 5, 6, 1, 7], [0, 1, 2, 3, 4], "n=8: set1={2,3,4,5,6}, set2={7}"), - ([0, 2, 3, 4, 5, 6, 7, 1], [0, 1, 2, 3, 4], "n=8: set1={2,3,4,5,6,7}, set2={}"), -] - - -def _expand(topologies): - """Flat list of (color_order, int_order, masses, label) pairs.""" - out = [] - for color_order, int_order, topo_label in topologies: - n_out = len(color_order) - 2 - for masses, mass_label in _mass_sets_for(n_out): - label = f"{mass_label}-{topo_label}" - out.append((color_order, int_order, masses, label)) - return out - - -CASES = _expand(TOPOLOGIES) -CASE_IDS = [c[3] for c in CASES] - - -def _physical_mask(det, p_ext): - """A physical event has finite det and finite momenta on every leg.""" - mask = np.isfinite(det) - for p in p_ext: - mask &= np.all(np.isfinite(p), axis=1) - return mask - - -def _filter_physical(det, p_ext): - """Keep physical events. Fail only if literally none survive.""" - mask = _physical_mask(det, p_ext) - if not mask.any(): - pytest.fail(f"No physical events out of {len(det)}") - return mask, [p[mask] for p in p_ext], det[mask] - - -# ---------------------------- -# Tests -# ---------------------------- - - -@pytest.mark.parametrize("color_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_momentum_conservation(rng, color_order, _int_order, masses, _label): - mapping = ms.ColorOrderedMapping(color_order) - r = rng.random((N, mapping.random_dim())) - e_cm = np.full(N, CM_ENERGY) - cond = [e_cm] + [np.full(N, m) for m in masses] - - *p_ext, det = mapping.map_forward(list(r.T), cond) - p_ext = list(p_ext) - _, p_ext, det = _filter_physical(det, p_ext) - - p_in = p_ext[0] + p_ext[1] - p_out_sum = sum(p_ext[i] for i in range(2, len(p_ext))) - assert p_in == approx(p_out_sum, abs=1e-3, rel=1e-6) - - -@pytest.mark.parametrize("color_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_on_shell_masses(rng, color_order, _int_order, masses, _label): - mapping = ms.ColorOrderedMapping(color_order) - r = rng.random((N, mapping.random_dim())) - e_cm = np.full(N, CM_ENERGY) - cond = [e_cm] + [np.full(N, m) for m in masses] - - *p_ext, det = mapping.map_forward(list(r.T), cond) - p_ext = list(p_ext) - _, p_ext, det = _filter_physical(det, p_ext) - - for i, m in enumerate(masses): - p = p_ext[i + 2] - m_check = np.sqrt(np.maximum(0, p[:, 0] ** 2 - np.sum(p[:, 1:] ** 2, axis=1))) - assert m_check == approx(m, abs=1e-2, rel=1e-3), f"particle {i} mass off" - - -@pytest.mark.parametrize("color_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_inverse(rng, color_order, _int_order, masses, _label): - """Forward o inverse = identity, to FP precision modulo a small tail. - - A handful of events at the very edge of the physical region produce a - forward det at the lower extreme of the distribution; for those, the - round-trip accumulates non-negligible FP error in atan2/sqrt/boost. We - drop the bottom 0.1% by |det| (genuine boundary events, near-zero - weight in any MC integral) and check the rest. - """ - n_events = N - mapping = ms.ColorOrderedMapping(color_order) - r = rng.random((n_events, mapping.random_dim())) - e_cm = np.full(n_events, CM_ENERGY) - cond = [e_cm] + [np.full(n_events, m) for m in masses] - - *p_ext, det = mapping.map_forward(list(r.T), cond) - p_ext = list(p_ext) - mask, p_ext_phys, det_phys = _filter_physical(det, p_ext) - cond_phys = [c[mask] for c in cond] - r_phys = r[mask] - - *r_inv, det_inv = mapping.map_inverse(p_ext_phys, cond_phys) - - floor = np.quantile(np.abs(det_phys), 0.001) - keep = np.isfinite(det_inv) & (np.abs(det_phys) > floor) - for ri in r_inv: - keep &= np.isfinite(ri) - - det_kept = det_phys[keep] - det_inv_kept = det_inv[keep] - # 0.1% tolerance absorbs accumulated FP error from chained - # boosts/rotations/sqrts in the deepest-walk topologies. Physics - # correctness is the job of test_phase_space_volume_matches_old. - assert det_inv_kept == approx(1 / det_kept, rel=1e-3) - - # Some randoms are discrete-branch choices (r_choice from 2->3 rungs); - # the round-trip recovers only the bin (< 0.5 vs >= 0.5), not the exact - # value. Bin-membership comparison handles those. - for i, (r_i, r_inv_i) in enumerate(zip(list(r_phys.T), r_inv)): - r_i_k = r_i[keep] - r_inv_i_k = r_inv_i[keep] - diff = np.abs(r_i_k - r_inv_i_k) - bin_ok = (r_i_k < 0.5) == (r_inv_i_k < 0.5) - ok = (diff < 1e-4) | bin_ok - assert np.all(ok), ( - f"random index {i}: {np.sum(~ok)} mismatches, " - f"max diff = {diff[~ok].max() if np.any(~ok) else 0:.3e}" - ) - - -@pytest.mark.parametrize("color_order,int_order,masses,_label", CASES, ids=CASE_IDS) -def test_phase_space_volume_matches_old(rng, color_order, int_order, masses, _label): - """The new class should integrate to the same phase-space volume as the - old TPropagatorMapping (matrix element = 1). - - This is *the* correctness test: both classes are MC estimators of the - same n-body integral, and their means must agree within a few sigma. - test_inverse is a precision check; this one is the physics check. - """ - n_events = N - new = ms.ColorOrderedMapping(color_order) - old = ms.TPropagatorMapping(int_order) - - r_new = rng.random((n_events, new.random_dim())) - r_old = rng.random((n_events, old.random_dim())) - e_cm = np.full(n_events, CM_ENERGY) - cond = [e_cm] + [np.full(n_events, m) for m in masses] - - *_, det_new = new.map_forward(list(r_new.T), cond) - *_, det_old = old.map_forward(list(r_old.T), cond) - - det_new = det_new[np.isfinite(det_new) & (det_new > 0)] - det_old = det_old[np.isfinite(det_old) & (det_old > 0)] - if not (len(det_new) and len(det_old)): - pytest.fail("no physical events in one of the estimators") - - mean_new, mean_old = np.mean(det_new), np.mean(det_old) - err_new = np.std(det_new) / np.sqrt(len(det_new)) - err_old = np.std(det_old) / np.sqrt(len(det_old)) - err = np.sqrt(err_new**2 + err_old**2) - - diff = abs(mean_new - mean_old) - assert diff < 5 * err, ( - f"phase-space volumes differ: new = {mean_new:.4e} +/- {err_new:.4e}, " - f"old = {mean_old:.4e} +/- {err_old:.4e}, " - f"diff = {diff:.4e} ({diff/err:.2f} sigma)" - ) diff --git a/madspace/tests/test_cuts.py b/madspace/tests/test_cuts.py new file mode 100644 index 000000000..952881057 --- /dev/null +++ b/madspace/tests/test_cuts.py @@ -0,0 +1,244 @@ +import math +from dataclasses import dataclass + +import numpy as np +import pytest + +import madspace as ms + +# from typing import List, Optional + + +PSM = ms.PhaseSpaceMapping + +CM_ENERGY = 13000.0 +N_SAMPLES = 300_000 +TOL_SIGMA = 5.0 +BASE_SEED = 2024 + +CUT_FRACTIONS = [0.05, 0.10, 0.20, 0.50] +CUT_DELTAR_MIN = 0.4 + +M_TOP, M_W, M_Z = 173.0, 80.4, 91.19 + + +VARIANTS = [ + ("ggg", [0.0, 0.0, 0.0], [21, 21, 21]), + ("gggg", [0.0, 0.0, 0.0, 0.0], [21, 21, 21, 21]), + ("ggggg", [0.0, 0.0, 0.0, 0.0, 0.0], [21, 21, 21, 21, 21]), + ("ttgg", [M_TOP, M_TOP, 0.0, 0.0], [6, -6, 21, 21]), + ("Wgg", [M_W, 0.0, 0.0], [24, 21, 21]), + ("Zgg", [M_Z, 0.0, 0.0], [23, 21, 21]), +] + + +@dataclass(frozen=True) +class Config: + topology: str + mode: any + color_order: list[int] | None + piped: bool + + +# def _mode_name(mode): +# return str(mode).rsplit(".", 1)[-1] + + +def _n_jets(pids_out): + return sum(1 for pid in pids_out if pid == 21) + + +def _pt_min(pids_out, fraction): + return fraction * CM_ENERGY / max(_n_jets(pids_out), 1) + + +def _color_orders(n_out): + """Each ordering is a distinct physical topology.""" + base = [i + 2 for i in range(n_out)] + + orders = [("empty-set (single chain)", [0, *base, 1])] + + for pos in range(1, n_out): + left, right = base[:pos], base[pos:] + orders.append( + (f"split topology | left={left} right={right}", [0, *left, 1, *right]) + ) + + return orders + + +def _inputs(mapping, rng, n): + inputs = [rng.random((n, mapping.random_dim()))] + dd = mapping.discrete_dim() + if dd: + inputs.append(rng.integers(0, 2, size=(n, dd)).astype(np.int32)) + return inputs + + +def _cuts(pids, pt_min): + O = ms.Observable + return ms.Cuts( + [ + ms.CutItem(O(pids, O.obs_pt, [O.jet_pids]), min=pt_min), + ms.CutItem(O(pids, O.obs_delta_r, [O.jet_pids]), min=CUT_DELTAR_MIN), + ] + ) + + +def _feats(p): + pt = np.sqrt(p[..., 1] ** 2 + p[..., 2] ** 2) + eta = np.arcsinh(np.divide(p[..., 3], pt, out=np.zeros_like(pt), where=pt > 0)) + phi = np.arctan2(p[..., 2], p[..., 1]) + return pt, eta, phi + + +def _passes(p_out, jet_idx, pt_min): + if not jet_idx: + return np.ones(p_out.shape[0], dtype=bool) + + pt, eta, phi = _feats(p_out) + + ok = np.all(pt[:, jet_idx] >= pt_min, axis=1) + + for a in range(len(jet_idx)): + for b in range(a + 1, len(jet_idx)): + i, j = jet_idx[a], jet_idx[b] + dphi = np.abs(phi[:, i] - phi[:, j]) + dphi = np.minimum(dphi, 2 * np.pi - dphi) + dR = np.sqrt((eta[:, i] - eta[:, j]) ** 2 + dphi**2) + ok = ok & (dR >= CUT_DELTAR_MIN) + + return ok + + +def _cut_volume( + masses_out, pids_out, mode, color_order, piped, pt_min, seed, n=N_SAMPLES +): + masses = [0.0, 0.0] + list(masses_out) + pids = [21, 21] + list(pids_out) + jet_idx = [i for i, pid in enumerate(pids_out) if pid == 21] + + mapping = PSM( + masses, + CM_ENERGY, + mode=mode, + leptonic=True, + color_order=color_order, + cuts=_cuts(pids, pt_min) if piped else None, + ) + + rng = np.random.default_rng(seed) + p_ext, _x1, _x2, det = mapping.map_forward(_inputs(mapping, rng, n)) + + det = np.asarray(det) * 2.0 ** mapping.discrete_dim() + p_out = np.asarray(p_ext)[:, 2:, :] + + passes = _passes(p_out, jet_idx, pt_min) + finite = np.isfinite(det) & np.all(np.isfinite(p_out), axis=(1, 2)) + + weights = np.nan_to_num(np.where(passes & finite, det, 0.0)) + + return float(weights.mean()), float(weights.std() / math.sqrt(n)) + + +def _configs(n_out): + cfgs = [] + + cfgs.append( + Config( + topology="propagator (external)", + mode=PSM.propagator, + color_order=None, + piped=False, + ) + ) + + cfgs.append( + Config( + topology="propagator (piped)", + mode=PSM.propagator, + color_order=None, + piped=True, + ) + ) + + for topo_name, order in _color_orders(n_out): + cfgs.append( + Config( + topology=f"color_ordered (external) | {topo_name}", + mode=PSM.color_ordered, + color_order=order, + piped=False, + ) + ) + + cfgs.append( + Config( + topology=f"color_ordered (piped) | {topo_name}", + mode=PSM.color_ordered, + color_order=order, + piped=True, + ) + ) + + return cfgs + + +_CASES = [ + (vi, fi, label, masses, pids, frac) + for vi, (label, masses, pids) in enumerate(VARIANTS) + for fi, frac in enumerate(CUT_FRACTIONS) +] + + +@pytest.mark.parametrize( + "vi, fi, label, masses_out, pids_out, fraction", + _CASES, + ids=[f"{c[2]}-f{c[5]:g}" for c in _CASES], +) +def test_cut_volume_consistent(vi, fi, label, masses_out, pids_out, fraction): + pt_min = _pt_min(pids_out, fraction) + seed0 = BASE_SEED + 1000 * vi + 100 * fi + + ref, ref_err = _cut_volume( + masses_out, pids_out, PSM.rambo, None, False, pt_min, seed0 + ) + + if ref <= 0.0 or not math.isfinite(ref): + pytest.skip(f"{label}: cut volume vanishes at pt_min={pt_min:.0f} GeV.") + + rows = [] + + for ci, cfg in enumerate(_configs(len(masses_out)), start=1): + val, err = _cut_volume( + masses_out, + pids_out, + cfg.mode, + cfg.color_order, + cfg.piped, + pt_min, + seed0 + ci, + ) + + sigma = math.sqrt(err**2 + ref_err**2) + pull = abs(val - ref) / sigma if sigma > 0 else math.inf + + rows.append((cfg.topology, val, err, pull)) + + failures = [r for r in rows if r[3] > TOL_SIGMA] + + if failures: + lines = [ + f"{label}: cut volume disagrees with rambo+external reference " + f"(tol {TOL_SIGMA:g} sigma, pt_min={pt_min:.0f} GeV, N={N_SAMPLES:_}).", + f" reference rambo external: {ref:.6e} +/- {ref_err:.2e}", + ] + + for topo, val, err, pull in rows: + mark = "FAIL" if pull > TOL_SIGMA else "ok " + lines.append( + f" {mark} {topo:45}: {val:.6e} +/- {err:.2e} " + f"(ratio {val/ref:5.3f}, {pull:6.1f} sigma)" + ) + + pytest.fail("\n".join(lines), pytrace=False) diff --git a/madspace/tests/test_t_channel.py b/madspace/tests/test_t_channel.py index fd78001b2..1040fd560 100644 --- a/madspace/tests/test_t_channel.py +++ b/madspace/tests/test_t_channel.py @@ -9,6 +9,25 @@ np.set_printoptions(linewidth=1000) +def _color_order(n_out, mode): + """Single-chain color order for color_ordered mode; None otherwise.""" + if mode != ms.PhaseSpaceMapping.color_ordered: + return None + return [0] + [i + 2 for i in range(n_out)] + [1] + + +def _fwd_inputs(mapping, rng, n): + """[random] (+ [discrete] when the mode declares discrete choices).""" + r = rng.random((n, mapping.random_dim())) + inputs = [r] + disc = None + dd = mapping.discrete_dim() + if dd: + disc = rng.integers(0, 2, size=(n, dd)).astype(np.int32) + inputs.append(disc) + return inputs, r, disc + + @pytest.fixture def rng(): return np.random.default_rng(1234) @@ -53,8 +72,9 @@ def masses(request): ms.PhaseSpaceMapping.propagator, ms.PhaseSpaceMapping.rambo, ms.PhaseSpaceMapping.chili, + ms.PhaseSpaceMapping.color_ordered, ], - ids=["propagator", "rambo", "chili"], + ids=["propagator", "rambo", "chili", "color_ordered"], ) def mode(request): return request.param @@ -65,9 +85,14 @@ def mode(request): def test_t_channel_masses(masses, rng, mode): - mapping = ms.PhaseSpaceMapping(masses, CM_ENERGY, mode=mode) - r = rng.random((BATCH_SIZE, mapping.random_dim())) - p_ext, x1, x2, det = mapping.map_forward([r]) + mapping = ms.PhaseSpaceMapping( + masses, + CM_ENERGY, + mode=mode, + color_order=_color_order(len(masses) - 2, mode), + ) + inputs, r, _ = _fwd_inputs(mapping, rng, BATCH_SIZE) + p_ext, x1, x2, det = mapping.map_forward(inputs) batch_phys = BATCH_SIZE if mode == ms.PhaseSpaceMapping.chili: @@ -79,13 +104,26 @@ def test_t_channel_masses(masses, rng, mode): m_ext = np.sqrt( np.maximum(0, p_ext[:, :, 0] ** 2 - np.sum(p_ext[:, :, 1:] ** 2, axis=2)) ) - assert m_ext == approx(m_ext_true, abs=1e-3, rel=1e-3) + # The massless on-shell reconstruction m = sqrt(E^2 - |p|^2) has an FP floor + # that scales with the particle energy (~2e-6 * E); the deep color_ordered + # boost chains reach it for a few high-energy boundary events. Shallower + # modes stay well under the flat 1e-3, so only color_ordered needs the + # energy-scaled term (a genuine kinematic error would exceed it ~1000x). + abs_tol = 1e-3 + if mode == ms.PhaseSpaceMapping.color_ordered: + abs_tol = 1e-3 + 2e-6 * np.abs(p_ext[:, :, 0]) + assert np.all(np.abs(m_ext - m_ext_true) <= abs_tol + 1e-3 * np.abs(m_ext_true)) def test_t_channel_incoming(masses, rng, mode): - mapping = ms.PhaseSpaceMapping(masses, CM_ENERGY, mode=mode) - r = rng.random((BATCH_SIZE, mapping.random_dim())) - p_ext, x1, x2, det = mapping.map_forward([r]) + mapping = ms.PhaseSpaceMapping( + masses, + CM_ENERGY, + mode=mode, + color_order=_color_order(len(masses) - 2, mode), + ) + inputs, r, _ = _fwd_inputs(mapping, rng, BATCH_SIZE) + p_ext, x1, x2, det = mapping.map_forward(inputs) batch_phys = BATCH_SIZE if mode == ms.PhaseSpaceMapping.chili: @@ -110,9 +148,14 @@ def test_t_channel_incoming(masses, rng, mode): def test_t_channel_momentum_conservation(masses, rng, mode): - mapping = ms.PhaseSpaceMapping(masses, CM_ENERGY, mode=mode) - r = rng.random((BATCH_SIZE, mapping.random_dim())) - p_ext, x1, x2, det = mapping.map_forward([r]) + mapping = ms.PhaseSpaceMapping( + masses, + CM_ENERGY, + mode=mode, + color_order=_color_order(len(masses) - 2, mode), + ) + inputs, r, _ = _fwd_inputs(mapping, rng, BATCH_SIZE) + p_ext, x1, x2, det = mapping.map_forward(inputs) if mode == ms.PhaseSpaceMapping.chili: physical_mask = det != 0.0 @@ -125,9 +168,15 @@ def test_t_channel_momentum_conservation(masses, rng, mode): def test_t_channel_inverse(masses, rng, mode): - mapping = ms.PhaseSpaceMapping(masses, CM_ENERGY, mode=mode, invariant_power=0.3) - r = rng.random((BATCH_SIZE, mapping.random_dim())) - p_ext, x1, x2, det = mapping.map_forward([r]) + mapping = ms.PhaseSpaceMapping( + masses, + CM_ENERGY, + mode=mode, + invariant_power=0.3, + color_order=_color_order(len(masses) - 2, mode), + ) + inputs, r, disc = _fwd_inputs(mapping, rng, BATCH_SIZE) + p_ext, x1, x2, det = mapping.map_forward(inputs) if mode == ms.PhaseSpaceMapping.chili: physical_mask = det != 0.0 @@ -135,12 +184,21 @@ def test_t_channel_inverse(masses, rng, mode): x1 = x1[physical_mask] x2 = x2[physical_mask] r = r[physical_mask] + if disc is not None: + disc = disc[physical_mask] det = det[physical_mask] - r_inv, det_inv = mapping.map_inverse((p_ext, x1, x2)) + out = mapping.map_inverse((p_ext, x1, x2)) + det_inv = out[-1] + r_inv = out[0] one_batch = np.ones_like(det) + # Continuous randoms round-trip; the discrete int channel (color_ordered) + # is checked for exact bin/solution recovery separately. assert r_inv == approx(r, abs=1e-3, rel=1e-3) assert det * det_inv == approx(one_batch, rel=1e-5) + if disc is not None: + disc_inv = np.asarray(out[1]).astype(np.int64) + assert np.array_equal(disc_inv, disc.astype(np.int64)) @pytest.mark.parametrize( @@ -152,17 +210,29 @@ def test_t_channel_inverse(masses, rng, mode): "energy", [10.0, 100.0, 1000.0], ids=["10GeV", "100GeV", "1TeV"] ) def test_t_channel_phase_space_volume(particle_count, energy, rng, mode): + co = _color_order(particle_count, mode) if mode == ms.PhaseSpaceMapping.chili: mapping = ms.PhaseSpaceMapping( - [0.0] * (particle_count + 2), energy, mode=mode, leptonic=False + [0.0] * (particle_count + 2), + energy, + mode=mode, + leptonic=False, + color_order=co, ) else: mapping = ms.PhaseSpaceMapping( - [0.0] * (particle_count + 2), energy, mode=mode, leptonic=True + [0.0] * (particle_count + 2), + energy, + mode=mode, + leptonic=True, + color_order=co, ) sample_count = 100000 - r = rng.random((sample_count, mapping.random_dim())) - *rest, det = mapping.map_forward([r]) + inputs, r, _ = _fwd_inputs(mapping, rng, sample_count) + *rest, det = mapping.map_forward(inputs) + # det is the per-branch Jacobian; the solution multiplicity (2 per 2->3 peel) + # is owned externally now, so apply 2^discrete_dim to recover the full volume. + det = det * 2.0 ** mapping.discrete_dim() ps_volume = ( (2 * math.pi) ** (4 - 3 * particle_count) * (math.pi / 2.0) ** (particle_count - 1) @@ -173,3 +243,71 @@ def test_t_channel_phase_space_volume(particle_count, energy, rng, mode): ps_volume /= (particle_count - 1) ** 2 # integration over x1*x2 in [0,1] std_error = np.std(det) / np.sqrt(sample_count) assert np.mean(det) == approx(ps_volume, abs=3 * std_error, rel=1e-6) + + +# --------------------------------------------------------------------------- +# Color-order variety (color_ordered mode, always via PhaseSpaceMapping). +# +# The mass/mode-fixture tests above only use the single-chain color order +# (beams cyclically adjacent -> all outgoing on one side). These cases also +# exercise *two-sided* splits (set1 and set2 both non-empty), which build a +# genuinely different t-channel topology, and a range of discrete-peel counts +# (discrete_dim grows by one for every set with > 2 outgoing partons). +# ColorOrderedMapping is never constructed directly here -- only through +# PhaseSpaceMapping, which is the only entry point it is meant to have. +# --------------------------------------------------------------------------- +_CO_VARIETY = [ + # (color_order, outgoing_masses, id) + ([0, 2, 3, 4, 1], [0.0, 0.0, 0.0], "n3 chain {2,3,4}|{}"), + ([0, 2, 3, 1, 4], [0.0, 0.0, 0.0], "n3 split {2,3}|{4}"), + ([0, 2, 1, 3, 4], [173.0, 80.0, 0.0], "n3 split {2}|{3,4} massive"), + ([0, 2, 3, 4, 5, 1], [0.0, 0.0, 0.0, 0.0], "n4 chain {2,3,4,5}|{}"), + ([0, 2, 3, 1, 4, 5], [0.0, 0.0, 0.0, 0.0], "n4 split {2,3}|{4,5}"), + ([0, 2, 3, 4, 1, 5], [80.0, 80.0, 80.0, 80.0], "n4 split {2,3,4}|{5} W"), + ([0, 2, 3, 4, 5, 6, 1], [0.0, 0.0, 0.0, 0.0, 0.0], "n5 chain {2..6}|{}"), + ([0, 2, 3, 4, 1, 5, 6], [0.0, 0.0, 0.0, 0.0, 0.0], "n5 split {2,3,4}|{5,6}"), + ([0, 2, 3, 1, 4, 5, 6], [173.0, 0.0, 0.0, 0.0, 0.0], "n5 split {2,3}|{4,5,6} top"), +] +_CO_IDS = [c[2] for c in _CO_VARIETY] + + +@pytest.mark.parametrize("color_order,out_masses,_label", _CO_VARIETY, ids=_CO_IDS) +def test_color_ordered_color_orders(rng, color_order, out_masses, _label): + masses = [0.0, 0.0, *out_masses] + mapping = ms.PhaseSpaceMapping( + masses, + CM_ENERGY, + mode=ms.PhaseSpaceMapping.color_ordered, + invariant_power=0.3, + color_order=color_order, + ) + # random_dim is color-order invariant; discrete_dim tracks the peel count. + assert mapping.random_dim() == 3 * len(out_masses) - 2 + inputs, r, disc = _fwd_inputs(mapping, rng, BATCH_SIZE) + + p_ext, x1, x2, det = mapping.map_forward(inputs) + + # (1) momentum conservation + p_in = np.sum(p_ext[:, :2], axis=1) + p_out = np.sum(p_ext[:, 2:], axis=1) + assert p_out == approx(p_in, rel=1e-6, abs=1e-9) + + # (2) on-shell external masses (energy-scaled FP floor for deep boost chains) + m_ext_true = np.full((BATCH_SIZE, len(masses)), masses) + m_ext = np.sqrt( + np.maximum(0, p_ext[:, :, 0] ** 2 - np.sum(p_ext[:, :, 1:] ** 2, axis=2)) + ) + abs_tol = 1e-3 + 2e-6 * np.abs(p_ext[:, :, 0]) + assert np.all(np.abs(m_ext - m_ext_true) <= abs_tol + 1e-3 * np.abs(m_ext_true)) + + # (3) inverse: continuous randoms round-trip, discrete choices exactly, + out = mapping.map_inverse((p_ext, x1, x2)) + r_inv, det_inv = out[0], out[-1] + assert r_inv == approx(r, abs=2e-3, rel=1e-3) + rt_err = np.abs(det * det_inv - 1.0) + assert np.quantile(rt_err, 0.99) < 1e-5 + assert np.isfinite(rt_err).all() + assert np.mean(rt_err > 1e-2) < 0.01 + if disc is not None: + disc_inv = np.asarray(out[1]).astype(np.int64) + assert np.array_equal(disc_inv, disc.astype(np.int64)) diff --git a/madspace/tests/test_two_particle.py b/madspace/tests/test_two_particle.py index 3af7b028d..5d3a0eb4f 100644 --- a/madspace/tests/test_two_particle.py +++ b/madspace/tests/test_two_particle.py @@ -18,12 +18,16 @@ def rng(): {"decay": True, "com": True}, {"decay": False, "com": False}, {"decay": False, "com": True}, + {"decay": False, "com": False, "has_cut": True}, + {"decay": False, "com": True, "has_cut": True}, ], ids=[ "1->2 decay", "1->2 decay, COM", "2->2 scattering", "2->2 scattering, COM", + "2->2 scattering, cut", + "2->2 scattering, COM, cut", ], ) def mapping_and_args(request): @@ -47,7 +51,9 @@ def make_args(point): ) else: - mapping = ms.TwoToTwoParticleScattering(com=com) + has_cut = request.param.get("has_cut", False) + mapping = ms.TwoToTwoParticleScattering(com=com, has_cut=has_cut) + cut = [zeros, zeros] if has_cut else [] if com: def make_args(point): @@ -55,14 +61,14 @@ def make_args(point): p0 = np.stack([point.m0, zeros, zeros, zeros], axis=1) pa = np.stack([e0, zeros, zeros, e0], axis=1) pb = np.stack([e0, zeros, zeros, -e0], axis=1) - return [point.r1, point.r2, point.m1, point.m2], [pa, pb], p0 + return [point.r1, point.r2, point.m1, point.m2], [pa, pb, *cut], p0 else: def make_args(point): return ( [point.r1, point.r2, point.m1, point.m2], - [point.pa, point.pb], + [point.pa, point.pb, *cut], point.p0, )