Feat madspace 2to3

madspace/CMakeLists.txt

@@ -137,6 +137,7 @@ add_library(
     src/phasespace/luminosity.cpp
     src/phasespace/topology.cpp
     src/phasespace/t_propagator_mapping.cpp
+    src/phasespace/color_ordered_mapping.cpp
     src/phasespace/phasespace.cpp
     src/phasespace/multichannel.cpp
     src/phasespace/madnis.cpp
@@ -189,6 +190,7 @@ add_library(
     include/madspace/phasespace/luminosity.hpp
     include/madspace/phasespace/topology.hpp
     include/madspace/phasespace/t_propagator_mapping.hpp
+    include/madspace/phasespace/color_ordered_mapping.hpp
     include/madspace/phasespace/phasespace.hpp
     include/madspace/phasespace/multichannel.hpp
     include/madspace/phasespace/madnis.hpp

madspace/generate_pyi.sh

@@ -2,4 +2,5 @@ pybind11-stubgen madspace._madspace_py \
     -o . \
     --enum-class-locations=TChannelMode:PhaseSpaceMapping.TChannelMode \
     --enum-class-locations=Activation:MLP.Activation \
-    --enum-class-locations=CutMode:Cuts.CutMode
+    --enum-class-locations=CutMode:Cuts.CutMode \
+    --enum-class-locations=LRSchedule:AdamOptimizer.LRSchedule

madspace/include/madspace/compgraphs/function_builder_mixin.inc

@@ -87,12 +87,12 @@ Value reduce_sum_vector(Value in1) {
     return instruction("reduce_sum_vector", {in1})[0];
 }
 
-Value batch_reduce_mean(Value in) {
-    return instruction("batch_reduce_mean", {in})[0];
+Value batch_reduce_mean(Value in1) {
+    return instruction("batch_reduce_mean", {in1})[0];
 }
 
-Value batch_reduce_mean_keepdim(Value in) {
-    return instruction("batch_reduce_mean_keepdim", {in})[0];
+Value batch_reduce_mean_keepdim(Value in1) {
+    return instruction("batch_reduce_mean_keepdim", {in1})[0];
 }
 
 Value reduce_product(Value in1) {
@@ -260,6 +260,16 @@ std::array<Value, 4> two_to_three_particle_scattering_inverse(Value p1, Value p2
     return {output_vector[0], output_vector[1], output_vector[2], output_vector[3]};
 }
 
+std::array<Value, 3> double_t_scattering(Value r_phi, Value pa, Value pb, Value t1_abs, Value t2_abs, Value m1) {
+    auto output_vector = instruction("double_t_scattering", {r_phi, pa, pb, t1_abs, t2_abs, m1});
+    return {output_vector[0], output_vector[1], output_vector[2]};
+}
+
+std::array<Value, 2> double_t_scattering_inverse(Value p1, Value p2, Value pa, Value pb) {
+    auto output_vector = instruction("double_t_scattering_inverse", {p1, p2, pa, pb});
+    return {output_vector[0], output_vector[1]};
+}
+
 std::array<Value, 4> three_body_decay_com(Value r_e1, Value r_e2, Value r_phi, Value r_cos_theta, Value r_beta, Value m0, Value m1, Value m2, Value m3) {
     auto output_vector = instruction("three_body_decay_com", {r_e1, r_e2, r_phi, r_cos_theta, r_beta, m0, m1, m2, m3});
     return {output_vector[0], output_vector[1], output_vector[2], output_vector[3]};
@@ -290,6 +300,26 @@ std::array<Value, 3> 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<Value, 2> 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});
+    return {output_vector[0], output_vector[1]};
+}
+
+std::array<Value, 3> 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});
+    return {output_vector[0], output_vector[1], output_vector[2]};
+}
+
+std::array<Value, 2> 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});
+    return {output_vector[0], output_vector[1]};
+}
+
+std::array<Value, 3> 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});
+    return {output_vector[0], output_vector[1], output_vector[2]};
+}
+
 std::array<Value, 2> s23_min_max(Value pa, Value pb, Value p3, Value t1_abs, Value m1, Value m2) {
     auto output_vector = instruction("s23_min_max", {pa, pb, p3, t1_abs, m1, m2});
     return {output_vector[0], output_vector[1]};

madspace/include/madspace/compgraphs/opcode_mixin.inc

@@ -59,86 +59,92 @@ two_to_two_particle_scattering = 57,
 two_to_two_particle_scattering_inverse = 58,
 two_to_three_particle_scattering = 59,
 two_to_three_particle_scattering_inverse = 60,
-three_body_decay_com = 61,
-three_body_decay_com_inverse = 62,
-three_body_decay = 63,
-three_body_decay_inverse = 64,
-t_inv_min_max = 65,
-t_inv_value_and_min_max = 66,
-s23_min_max = 67,
-s23_value_and_min_max = 68,
-invariants_from_momenta = 69,
-sde2_channel_weights = 70,
-subchannel_weights = 71,
-apply_subchannel_weights = 72,
-pt_eta_phi_x = 73,
-mirror_momenta = 74,
-momenta_to_x1x2 = 75,
-uniform_invariant = 76,
-uniform_invariant_inverse = 77,
-breit_wigner_invariant = 78,
-breit_wigner_invariant_inverse = 79,
-stable_invariant = 80,
-stable_invariant_inverse = 81,
-stable_invariant_nu = 82,
-stable_invariant_nu_inverse = 83,
-fast_rambo_massless = 84,
-fast_rambo_massless_inverse = 85,
-fast_rambo_massless_com = 86,
-fast_rambo_massive = 87,
-fast_rambo_massive_inverse = 88,
-fast_rambo_massive_com = 89,
-cut_unphysical = 90,
-cut_one = 91,
-cut_all = 92,
-cut_any = 93,
-scale_transverse_energy = 94,
-scale_transverse_mass = 95,
-scale_half_transverse_mass = 96,
-scale_partonic_energy = 97,
-chili_forward = 98,
-chili_inverse = 99,
-matrix_element = 100,
-collect_channel_weights = 101,
-interpolate_pdf = 102,
-interpolate_alpha_s = 103,
-matmul = 104,
-relu = 105,
-leaky_relu = 106,
-elu = 107,
-gelu = 108,
-sigmoid = 109,
-softplus = 110,
-rqs_reshape = 111,
-rqs_find_bin = 112,
-rqs_forward = 113,
-rqs_inverse = 114,
-softmax = 115,
-softmax_prior = 116,
-sample_discrete = 117,
-sample_discrete_inverse = 118,
-sample_discrete_probs = 119,
-sample_discrete_probs_inverse = 120,
-discrete_histogram = 121,
-permute_momenta = 122,
-gather = 123,
-gather_int = 124,
-select_int = 125,
-select = 126,
-select_vector = 127,
-argsort = 128,
-one_hot = 129,
-madnis_abs_weight = 130,
-madnis_softclip = 131,
-madnis_variance = 132,
-madnis_single_channel_variance = 133,
-madnis_multi_channel_variance = 134,
-nonzero = 135,
-batch_gather = 136,
-batch_scatter = 137,
-random = 138,
-unweight = 139,
-vegas_forward = 140,
-vegas_inverse = 141,
-vegas_histogram = 142,
-histogram = 143
+double_t_scattering = 61,
+double_t_scattering_inverse = 62,
+three_body_decay_com = 63,
+three_body_decay_com_inverse = 64,
+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,
+one_hot = 135,
+madnis_abs_weight = 136,
+madnis_softclip = 137,
+madnis_variance = 138,
+madnis_single_channel_variance = 139,
+madnis_multi_channel_variance = 140,
+nonzero = 141,
+batch_gather = 142,
+batch_scatter = 143,
+random = 144,
+unweight = 145,
+vegas_forward = 146,
+vegas_inverse = 147,
+vegas_histogram = 148,
+histogram = 149

madspace/include/madspace/phasespace.hpp

@@ -22,6 +22,7 @@
 #include "phasespace/rambo.hpp"
 #include "phasespace/scale.hpp"
 #include "phasespace/t_propagator_mapping.hpp"
+#include "phasespace/color_ordered_mapping.hpp"
 #include "phasespace/three_particle.hpp"
 #include "phasespace/topology.hpp"
 #include "phasespace/two_particle.hpp"

madspace/include/madspace/phasespace/color_ordered_mapping.hpp

@@ -0,0 +1,55 @@
+#pragma once
+
+#include <vector>
+
+#include "madspace/phasespace/base.hpp"
+#include "madspace/phasespace/invariants.hpp"
+#include "madspace/phasespace/three_particle.hpp"
+#include "madspace/phasespace/topology.hpp"
+#include "madspace/phasespace/two_particle.hpp"
+
+namespace madspace {
+
+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.
+    ColorOrderedMapping(
+        const std::vector<std::size_t>& color_order,
+        double t_invariant_power = 0.8,
+        double s_invariant_power = 0.8
+    );
+
+    std::size_t random_dim() const { return _random_dim; }
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb,
+        const NamedVector<Value>& inputs,
+        const NamedVector<Value>& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb,
+        const NamedVector<Value>& inputs,
+        const NamedVector<Value>& conditions
+    ) const override;
+
+    // 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.
+    std::vector<std::size_t> _set1;
+    std::vector<std::size_t> _set2;
+    std::size_t _n_out;
+    std::size_t _random_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;
+
+    Invariant _uniform_invariant;
+    TwoToTwoParticleScattering _com_scattering;
+    TwoToTwoParticleScattering _lab_scattering;
+    TwoToThreeParticleScattering _two_to_three;
+    DoubleT _double_t;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/two_particle.hpp

@@ -47,4 +47,31 @@ class TwoToTwoParticleScattering : public Mapping {
     Invariant _invariant;
 };
 
+class DoubleT : public Mapping {
+public:
+    DoubleT(
+        double t1_invariant_power = 0,
+        double t1_mass = 0,
+        double t1_width = 0,
+        double t2_invariant_power = 0,
+        double t2_mass = 0,
+        double t2_width = 0
+    );
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb,
+        const NamedVector<Value>& inputs,
+        const NamedVector<Value>& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb,
+        const NamedVector<Value>& inputs,
+        const NamedVector<Value>& conditions
+    ) const override;
+
+    Invariant _t1_invariant;
+    Invariant _t2_invariant;
+};
+
 } // namespace madspace