From 5096c25248ecf9e4abec59c5ab316de5285114dd Mon Sep 17 00:00:00 2001 From: Ramon Winterhalder Date: Sun, 10 May 2026 10:51:01 +0200 Subject: [PATCH 1/6] Fix inverse mapping of 2>3 mapping --- madspace/src/kernels/kinematics.hpp | 71 ++++++++++++++++++++++++++ madspace/src/kernels/threeparticle.hpp | 6 ++- madspace/tests/test_2to3_scattering.py | 46 ++++++++--------- 3 files changed, 97 insertions(+), 26 deletions(-) diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index 9150233e4..b5216ec0c 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -197,6 +197,77 @@ KERNELSPEC FourMom rotate_inverse(FourMom p, FourMom q) { }; } +template +KERNELSPEC FourMom +rotate_two_ref(FourMom p, FourMom q_z, FourMom q_x) { + // Forward rotation: take p from a canonical frame + // (e_z along q_z, e_x in the plane spanned by q_z and q_x with positive + // component along q_x's perpendicular part) + // into the world frame. Energy unchanged. + // + // Used by kernel_two_to_three_particle_scattering: q_z = pa_com, + // q_x = p3 boosted into the p_12 rest frame. Picks a unique azimuth + // so that the (q_z, q_x) plane is the phi=0 half-plane. + + // z_hat = q_z / |q_z| + auto qz_n2 = q_z[1]*q_z[1] + q_z[2]*q_z[2] + q_z[3]*q_z[3]; + auto qz_n = sqrt(max(qz_n2, EPS2)); + auto zx = q_z[1] / qz_n, zy = q_z[2] / qz_n, zz = q_z[3] / qz_n; + + // x_hat = (q_x perp to z_hat) / |...| + auto qx_dot_z = q_x[1]*zx + q_x[2]*zy + q_x[3]*zz; + auto rx = q_x[1] - qx_dot_z*zx; + auto ry = q_x[2] - qx_dot_z*zy; + auto rz = q_x[3] - qx_dot_z*zz; + auto rn2 = rx*rx + ry*ry + rz*rz; + auto rn = sqrt(max(rn2, EPS2)); + auto xx = rx / rn, xy = ry / rn, xz = rz / rn; + + // y_hat = z_hat x x_hat + auto yx = zy*xz - zz*xy; + auto yy = zz*xx - zx*xz; + auto yz = zx*xy - zy*xx; + + // world spatial = p[1]*x_hat + p[2]*y_hat + p[3]*z_hat + return FourMom{ + p[0], + p[1]*xx + p[2]*yx + p[3]*zx, + p[1]*xy + p[2]*yy + p[3]*zy, + p[1]*xz + p[2]*yz + p[3]*zz, + }; +} + +template +KERNELSPEC FourMom +rotate_two_ref_inverse(FourMom p, FourMom q_z, FourMom q_x) { + // Inverse of rotate_two_ref: take a vector p from the world frame into the + // canonical frame defined by (q_z, q_x). Energy unchanged. + + auto qz_n2 = q_z[1]*q_z[1] + q_z[2]*q_z[2] + q_z[3]*q_z[3]; + auto qz_n = sqrt(max(qz_n2, EPS2)); + auto zx = q_z[1] / qz_n, zy = q_z[2] / qz_n, zz = q_z[3] / qz_n; + + auto qx_dot_z = q_x[1]*zx + q_x[2]*zy + q_x[3]*zz; + auto rx = q_x[1] - qx_dot_z*zx; + auto ry = q_x[2] - qx_dot_z*zy; + auto rz = q_x[3] - qx_dot_z*zz; + auto rn2 = rx*rx + ry*ry + rz*rz; + auto rn = sqrt(max(rn2, EPS2)); + auto xx = rx / rn, xy = ry / rn, xz = rz / rn; + + auto yx = zy*xz - zz*xy; + auto yy = zz*xx - zx*xz; + auto yz = zx*xy - zy*xx; + + // canonical components = world spatial . (x_hat, y_hat, z_hat) + return FourMom{ + p[0], + p[1]*xx + p[2]*xy + p[3]*xz, + p[1]*yx + p[2]*yy + p[3]*yz, + p[1]*zx + p[2]*zy + p[3]*zz, + }; +} + template KERNELSPEC FourMom boost(FourMom k, FourMom p_boost, FVal sign) { // Perform the boost diff --git a/madspace/src/kernels/threeparticle.hpp b/madspace/src/kernels/threeparticle.hpp index 3f588a5e9..fc105f24e 100644 --- a/madspace/src/kernels/threeparticle.hpp +++ b/madspace/src/kernels/threeparticle.hpp @@ -567,7 +567,8 @@ KERNELSPEC void kernel_two_to_three_particle_scattering( auto p1_com = scatter_out.first; auto gram4 = bk_gram4(m0_2, ma_2, mb_2, m1_2, m2_2, m3_2, t1_abs, t2, s12, s23); auto det_2to3 = 1 / (8 * sqrt(max(-gram4, EPS2))); - auto p1_rot = rotate(p1_com, pa_com); + auto p3_p12 = boost(load_mom(p3), p_12, -1.); + auto p1_rot = rotate_two_ref(p1_com, pa_com, p3_p12); auto p1_lab = boost(p1_rot, p_12, 1.); store_mom(p1, p1_lab); for (int i = 0; i < 4; ++i) { @@ -605,7 +606,8 @@ KERNELSPEC void kernel_two_to_three_particle_scattering_inverse( auto pa_com = boost(load_mom(pa), p_12, -1.); auto p1_com = boost(load_mom(p1), p_12, -1.); - auto p1_rot = rotate_inverse(p1_com, pa_com); + auto p3_p12 = boost(load_mom(p3), p_12, -1.); + auto p1_rot = rotate_two_ref_inverse(p1_com, pa_com, p3_p12); auto m1_2 = lsquare(load_mom(p1)); auto m2_2 = lsquare(load_mom(p2)); diff --git a/madspace/tests/test_2to3_scattering.py b/madspace/tests/test_2to3_scattering.py index 6136bb343..147d5011f 100644 --- a/madspace/tests/test_2to3_scattering.py +++ b/madspace/tests/test_2to3_scattering.py @@ -186,30 +186,28 @@ def test_momentum_conservation(input_points): assert p1 + p2 == approx(input_points.p12) -# TODO: Find bug in inverse mapping -# def test_inverse(input_points): -# mapping = ms.TwoToThreeParticleScattering() - -# inputs = [ -# input_points.r_choice, -# input_points.r_s23, -# input_points.r_t1, -# input_points.m1, -# input_points.m2, -# ] -# conditions = [input_points.pa, input_points.pb, input_points.p3] - -# p1, p2, det = mapping.map_forward(inputs, conditions) -# inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) -# # r = inv_det * det -# # print("r stats:", np.min(r), np.median(r), np.max(r)) -# # print("abs(r-1) median:", np.median(np.abs(r-1))) -# # assert inv_det == approx(1 / det) -# for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): -# if i == 0: -# assert ((inp < 0.5) == (inv_inp < 0.5)).all() -# continue -# assert inp == approx(inv_inp), f"mismatch in input index {i}" +def test_inverse(input_points): + mapping = ms.TwoToThreeParticleScattering() + + inputs = [ + input_points.r_choice, + input_points.r_s23, + input_points.r_t1, + input_points.m1, + input_points.m2, + ] + conditions = [input_points.pa, input_points.pb, input_points.p3] + + p1, p2, det = mapping.map_forward(inputs, conditions) + *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) + + assert inv_det == approx(1 / det, rel=1e-5) + + for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): + if i == 0: + assert ((inp < 0.5) == (inv_inp < 0.5)).all() + continue + assert inp == approx(inv_inp), f"mismatch in input index {i}" def test_on_shell_masses(input_points): From cb90a166e1494dcc6ad738693ab1de2acb4853c1 Mon Sep 17 00:00:00 2001 From: Ramon Winterhalder Date: Mon, 11 May 2026 11:24:55 +0200 Subject: [PATCH 2/6] Add double_t mapping from rikkert --- .../compgraphs/function_builder_mixin.inc | 30 ++ .../madspace/compgraphs/opcode_mixin.inc | 172 ++++++------ .../madspace/phasespace/two_particle.hpp | 27 ++ madspace/instruction_set.yaml | 156 +++++++++++ .../src/compgraphs/instruction_set_mixin.inc | 172 ++++++------ madspace/src/cpu/runtime_backward_mixin.inc | 42 +-- madspace/src/cpu/runtime_mixin.inc | 182 +++++++------ madspace/src/gpu/runtime_backward_mixin.inc | 42 +-- madspace/src/gpu/runtime_mixin.inc | 182 +++++++------ madspace/src/kernels/kinematics.hpp | 112 ++++++++ madspace/src/kernels/twoparticle.hpp | 80 ++++++ madspace/src/phasespace/two_particle.cpp | 78 ++++++ madspace/src/python/instruction_set.hpp | 6 + madspace/src/python/madspace.cpp | 11 + madspace/tests/test_double_t.py | 257 ++++++++++++++++++ 15 files changed, 1177 insertions(+), 372 deletions(-) create mode 100644 madspace/tests/test_double_t.py diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc index ae220de13..d7a47e3c5 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/madspace/include/madspace/compgraphs/function_builder_mixin.inc @@ -260,6 +260,16 @@ std::array 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 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 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], output_vector[2]}; +} + std::array 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 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}); + return {output_vector[0], output_vector[1]}; +} + +std::array t1_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value mir_min) { + auto output_vector = instruction("t1_inv_value_and_min_max_doublet", {pa, pb, p1, mir_min}); + 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}); + return {output_vector[0], output_vector[1]}; +} + +std::array t2_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value mir_min, Value t1_abs) { + auto output_vector = instruction("t2_inv_value_and_min_max_doublet", {pa, pb, p1, mir_min, t1_abs}); + return {output_vector[0], output_vector[1], output_vector[2]}; +} + std::array 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]}; diff --git a/madspace/include/madspace/compgraphs/opcode_mixin.inc b/madspace/include/madspace/compgraphs/opcode_mixin.inc index 58a2cd229..39c4dadb3 100644 --- a/madspace/include/madspace/compgraphs/opcode_mixin.inc +++ b/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 diff --git a/madspace/include/madspace/phasespace/two_particle.hpp b/madspace/include/madspace/phasespace/two_particle.hpp index b038f2c4a..5faf94994 100644 --- a/madspace/include/madspace/phasespace/two_particle.hpp +++ b/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& inputs, + const NamedVector& conditions + ) const override; + Result build_inverse_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions + ) const override; + + Invariant _t1_invariant; + Invariant _t2_invariant; +}; + } // namespace madspace diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index ec98a4842..4e4d22dce 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -964,6 +964,62 @@ two_to_three_particle_scattering_inverse: type: [float] desc: +double_t_scattering: + inputs: + - name: r_phi + type: [float] + desc: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: t1_abs + type: [float] + desc: + - name: t2_abs + type: [float] + desc: + - name: m1 + type: [float] + desc: + outputs: + - name: p1 + type: [float, 4] + desc: + - name: p2 + type: [float, 4] + desc: + - name: det + type: [float] + desc: + +double_t_scattering_inverse: + inputs: + - name: p1 + type: [float, 4] + desc: + - name: p2 + type: [float, 4] + desc: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + outputs: + - name: r_phi + type: [float] + desc: + - name: m1 + type: [float] + desc: + - name: det + type: [float] + desc: + three_body_decay_com: inputs: - name: r_e1 @@ -1189,6 +1245,106 @@ t_inv_value_and_min_max: type: [float] desc: +t1_inv_min_max_doublet: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: m1 + type: [float] + desc: + - name: mir_min + type: [float] + desc: + outputs: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + +t1_inv_value_and_min_max_doublet: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: p1 + type: [float, 4] + desc: + - name: mir_min + type: [float] + desc: + outputs: + - name: t_abs + type: [float] + desc: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + +t2_inv_min_max_doublet: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: m1 + type: [float] + desc: + - name: mir_min + type: [float] + desc: + - name: t1_abs + type: [float] + desc: + outputs: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + +t2_inv_value_and_min_max_doublet: + inputs: + - name: pa + type: [float, 4] + desc: + - name: pb + type: [float, 4] + desc: + - name: p1 + type: [float, 4] + desc: + - name: mir_min + type: [float] + desc: + - name: t1_abs + type: [float] + desc: + outputs: + - name: t_abs + type: [float] + desc: + - name: t_min + type: [float] + desc: + - name: t_max + type: [float] + desc: + s23_min_max: inputs: - name: pa diff --git a/madspace/src/compgraphs/instruction_set_mixin.inc b/madspace/src/compgraphs/instruction_set_mixin.inc index 2ed08f15b..3587c7a45 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -76,87 +76,93 @@ InstructionOwner instructions[] { mi("two_to_two_particle_scattering_inverse", 58, 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}}), mi("two_to_three_particle_scattering", 59, true, {{DataType::dt_int, false, {}, 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, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), mi("two_to_three_particle_scattering_inverse", 60, 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, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("three_body_decay_com", 61, 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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("three_body_decay_com_inverse", 62, 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, {}, false}}), - mi("three_body_decay", 63, 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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {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}}), - mi("three_body_decay_inverse", 64, 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", 65, 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", 66, 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("s23_min_max", 67, 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", 68, 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", 69, 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", 70, 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", 71, 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", 72, 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", 73, 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", 74, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), - mi("momenta_to_x1x2", 75, 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", 76, 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", 77, 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", 78, 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", 79, 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", 80, 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", 81, 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", 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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("stable_invariant_nu_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}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("fast_rambo_massless", 84, 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", 85, 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", 86, 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", 87, 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", 88, 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", 89, 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", 90, 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", 91, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("cut_all", 92, 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", 93, 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", 94, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("scale_transverse_mass", 95, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("scale_half_transverse_mass", 96, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("scale_partonic_energy", 97, true, {{DataType::dt_float, false, {"n", 4}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("chili_forward", 98, 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", 99, 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(100, true)), - mi("collect_channel_weights", 101, 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", 102, 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", 103, 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", 104, 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", 105, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("leaky_relu", 106, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("elu", 107, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("gelu", 108, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("sigmoid", 109, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softplus", 110, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - InstructionOwner(new RqsReshapeInstruction(111, true)), - mi("rqs_find_bin", 112, 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", 113, 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", 114, 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", 115, true, {{DataType::dt_float, false, {std::monostate{}}, false}}, {{DataType::dt_float, false, {std::monostate{}}, false}}), - mi("softmax_prior", 116, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("sample_discrete", 117, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), - mi("sample_discrete_inverse", 118, 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", 119, 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", 120, 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", 121, 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", 122, 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", 123, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("gather_int", 124, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), - mi("select_int", 125, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), - mi("select", 126, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), - mi("select_vector", 127, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), - mi("argsort", 128, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), - mi("one_hot", 129, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), - mi("madnis_abs_weight", 130, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_softclip", 131, 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", 132, 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", 133, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), - mi("madnis_multi_channel_variance", 134, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), - InstructionOwner(new NonzeroInstruction(135, true)), - InstructionOwner(new BatchGatherInstruction(136, true)), - InstructionOwner(new BatchScatterInstruction(137, true)), - InstructionOwner(new RandomInstruction(138, true)), - InstructionOwner(new UnweightInstruction(139, true)), - mi("vegas_forward", 140, 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", 141, 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", 142, 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", 143, 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("double_t_scattering", 61, true, {{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, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("double_t_scattering_inverse", 62, 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("three_body_decay_com", 63, 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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), + mi("three_body_decay_com_inverse", 64, 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, {}, false}}), + mi("three_body_decay", 65, 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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {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}}), + 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}}), + 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}}), + 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, true, {}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_inverse", 124, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_probs", 125, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("sample_discrete_probs_inverse", 126, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}), + mi("discrete_histogram", 127, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, true, {"n"}, false}, {DataType::dt_int, true, {"n"}, false}}), + mi("permute_momenta", 128, true, {{DataType::dt_float, false, {"n", 4}, false}, {DataType::dt_int, true, {"m", "n"}, false}, {DataType::dt_int, false, {}, false}}, {{DataType::dt_float, false, {"n", 4}, false}}), + mi("gather", 129, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("gather_int", 130, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, false, {"n"}, false}}, {{DataType::dt_int, false, {}, false}}), + mi("select_int", 131, true, {{DataType::dt_int, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_int, false, {"m"}, false}}), + mi("select", 132, true, {{DataType::dt_float, false, {"n"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m"}, false}}), + mi("select_vector", 133, true, {{DataType::dt_float, false, {"n", "k"}, false}, {DataType::dt_int, false, {"m"}, false}}, {{DataType::dt_float, false, {"m", "k"}, false}}), + mi("argsort", 134, true, {{DataType::dt_float, false, {"n"}, false}}, {{DataType::dt_int, false, {"n"}, false}}), + mi("one_hot", 135, true, {{DataType::dt_int, false, {}, false}, {DataType::dt_int, true, {"n"}, true}}, {{DataType::dt_float, false, {"n"}, false}}), + mi("madnis_abs_weight", 136, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_softclip", 137, 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", 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_single_channel_variance", 139, true, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}}), + mi("madnis_multi_channel_variance", 140, true, {{DataType::dt_float, false, {"c"}, false}, {DataType::dt_float, false, {"c"}, false}}, {{DataType::dt_float, false, {}, false}}), + InstructionOwner(new NonzeroInstruction(141, true)), + InstructionOwner(new BatchGatherInstruction(142, true)), + InstructionOwner(new BatchScatterInstruction(143, true)), + InstructionOwner(new RandomInstruction(144, true)), + InstructionOwner(new UnweightInstruction(145, true)), + mi("vegas_forward", 146, 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", 147, 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", 148, 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", 149, 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 4474225a4..9a1327b88 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 104: +case 110: backward_op_matmul(instr, locals, local_grads, device); break; -case 105: +case 111: backward_batch_foreach, backward_kernel_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 106: +case 112: backward_batch_foreach, backward_kernel_leaky_relu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 107: +case 113: backward_batch_foreach, backward_kernel_elu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 108: +case 114: backward_batch_foreach, backward_kernel_gelu, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 109: +case 115: backward_batch_foreach, backward_kernel_sigmoid, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 110: +case 116: backward_batch_foreach, backward_kernel_softplus, 2, 1, DeviceType>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 111: +case 117: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 112: +case 118: 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 113: +case 119: backward_batch_foreach, backward_kernel_rqs_forward, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 114: +case 120: backward_batch_foreach, backward_kernel_rqs_inverse, 4, 2, 2, DeviceType>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 115: +case 121: backward_batch_foreach, backward_kernel_softmax, 2, 1, DeviceType>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 116: +case 122: backward_batch_foreach, backward_kernel_softmax_prior, 2, 2, 1, DeviceType>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 120: +case 126: 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 123: +case 129: backward_batch_foreach, backward_kernel_gather, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 126: +case 132: backward_batch_foreach, backward_kernel_select, 2, 2, 1, DeviceType>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 130: +case 136: 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 131: +case 137: 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 132: +case 138: 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 133: +case 139: 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 134: +case 140: 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 dea3f9934..669e0fa97 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -185,251 +185,269 @@ case 60: batch_foreach, kernel_two_to_three_particle_scattering_inverse, 7, 4, 1, DeviceType>, 7, 4>(instr, locals, device); break; case 61: - batch_foreach, kernel_three_body_decay_com, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device); + batch_foreach, kernel_double_t_scattering, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 62: - batch_foreach, kernel_three_body_decay_com_inverse, 3, 10, 1, DeviceType>, 3, 10>(instr, locals, device); + batch_foreach, kernel_double_t_scattering_inverse, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); break; case 63: - batch_foreach, kernel_three_body_decay, 10, 4, 1, DeviceType>, 10, 4>(instr, locals, device); + batch_foreach, kernel_three_body_decay_com, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device); break; case 64: - batch_foreach, kernel_three_body_decay_inverse, 3, 11, 1, DeviceType>, 3, 11>(instr, locals, device); + batch_foreach, kernel_three_body_decay_com_inverse, 3, 10, 1, DeviceType>, 3, 10>(instr, locals, device); break; case 65: - batch_foreach, kernel_t_inv_min_max, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_three_body_decay, 10, 4, 1, DeviceType>, 10, 4>(instr, locals, device); break; case 66: - batch_foreach, kernel_t_inv_value_and_min_max, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); + batch_foreach, kernel_three_body_decay_inverse, 3, 11, 1, DeviceType>, 3, 11>(instr, locals, device); break; case 67: - batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); + batch_foreach, kernel_t_inv_min_max, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 68: - batch_foreach, kernel_s23_value_and_min_max, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); + batch_foreach, kernel_t_inv_value_and_min_max, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); break; case 69: - batch_foreach, kernel_invariants_from_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_t1_inv_min_max_doublet, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 70: - batch_foreach, kernel_sde2_channel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_t1_inv_value_and_min_max_doublet, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); break; case 71: - batch_foreach, kernel_subchannel_weights, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_t2_inv_min_max_doublet, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 72: - batch_foreach, kernel_apply_subchannel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_t2_inv_value_and_min_max_doublet, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device); break; case 73: - batch_foreach, kernel_pt_eta_phi_x, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); break; case 74: - batch_foreach, kernel_mirror_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_s23_value_and_min_max, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 75: - batch_foreach, kernel_momenta_to_x1x2, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_invariants_from_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 76: - batch_foreach, kernel_uniform_invariant, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_sde2_channel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 77: - batch_foreach, kernel_uniform_invariant_inverse, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_subchannel_weights, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 78: - batch_foreach, kernel_breit_wigner_invariant, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_apply_subchannel_weights, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 79: - batch_foreach, kernel_breit_wigner_invariant_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_pt_eta_phi_x, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 80: - batch_foreach, kernel_stable_invariant, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_mirror_momenta, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 81: - batch_foreach, kernel_stable_invariant_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_momenta_to_x1x2, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 82: - batch_foreach, kernel_stable_invariant_nu, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_uniform_invariant, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 83: - batch_foreach, kernel_stable_invariant_nu_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_uniform_invariant_inverse, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 84: - batch_foreach, kernel_fast_rambo_massless, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_breit_wigner_invariant, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 85: - batch_foreach, kernel_fast_rambo_massless_inverse, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device); + batch_foreach, kernel_breit_wigner_invariant_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 86: - batch_foreach, kernel_fast_rambo_massless_com, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 87: - batch_foreach, kernel_fast_rambo_massive, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 88: - batch_foreach, kernel_fast_rambo_massive_inverse, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); + batch_foreach, kernel_stable_invariant_nu, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 89: - batch_foreach, kernel_fast_rambo_massive_com, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); + batch_foreach, kernel_stable_invariant_nu_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 90: - batch_foreach, kernel_cut_unphysical, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 91: - batch_foreach, kernel_cut_one, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless_inverse, 2, 3, 1, DeviceType>, 2, 3>(instr, locals, device); break; case 92: - batch_foreach, kernel_cut_all, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massless_com, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 93: - batch_foreach, kernel_cut_any, 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 94: - batch_foreach, kernel_scale_transverse_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive_inverse, 3, 3, 1, DeviceType>, 3, 3>(instr, locals, device); break; case 95: - batch_foreach, kernel_scale_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_fast_rambo_massive_com, 3, 2, 1, DeviceType>, 3, 2>(instr, locals, device); break; case 96: - batch_foreach, kernel_scale_half_transverse_mass, 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 97: - batch_foreach, kernel_scale_partonic_energy, 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 98: - batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_cut_all, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 99: - batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); + batch_foreach, kernel_cut_any, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 100: - op_matrix_element(instr, locals, device); + batch_foreach, kernel_scale_transverse_energy, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 101: - batch_foreach, kernel_collect_channel_weights, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_scale_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 102: - batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); + batch_foreach, kernel_scale_half_transverse_mass, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 103: - batch_foreach, kernel_interpolate_alpha_s, 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 104: - op_matmul(instr, locals, device); + batch_foreach, kernel_chili_forward, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 105: - batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_chili_inverse, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 106: - batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 107: - batch_foreach, kernel_elu, 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 108: - batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_interpolate_pdf, 6, 1, 1, DeviceType>, 6, 1>(instr, locals, device); break; case 109: - batch_foreach, kernel_sigmoid, 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 110: - batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 111: - op_rqs_reshape(instr, locals, device); + batch_foreach, kernel_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 112: - batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_leaky_relu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 113: - batch_foreach, kernel_rqs_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_elu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 114: - batch_foreach, kernel_rqs_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_gelu, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 115: - batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_sigmoid, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 116: - batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_softplus, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 117: - batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 118: - batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_rqs_find_bin, 4, 1, 2, DeviceType>, 4, 1>(instr, locals, device); break; case 119: - batch_foreach, kernel_sample_discrete_probs, 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 120: - batch_foreach, kernel_sample_discrete_probs_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 121: - op_discrete_histogram(instr, locals, device); + batch_foreach, kernel_softmax, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 122: - batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); + batch_foreach, kernel_softmax_prior, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 123: - batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 124: - batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_sample_discrete_inverse, 2, 2, 1, DeviceType>, 2, 2>(instr, locals, device); break; case 125: - batch_foreach, kernel_select_int, 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 126: - batch_foreach, kernel_select, 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 127: - batch_foreach, kernel_select_vector, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 128: - batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); + batch_foreach, kernel_permute_momenta, 3, 1, 1, DeviceType>, 3, 1>(instr, locals, device); break; case 129: - batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 130: - batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_gather_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 131: - batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_select_int, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 132: - batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); + batch_foreach, kernel_select, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 133: - batch_foreach, kernel_madnis_single_channel_variance, 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 134: - batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); + batch_foreach, kernel_argsort, 1, 1, 1, DeviceType>, 1, 1>(instr, locals, device); break; case 135: - op_nonzero(instr, locals, device); + batch_foreach, kernel_one_hot, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 136: - op_batch_gather(instr, locals, device); + batch_foreach, kernel_madnis_abs_weight, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 137: - op_batch_scatter(instr, locals, device); + batch_foreach, kernel_madnis_softclip, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 138: - op_random(instr, locals, device); + batch_foreach, kernel_madnis_variance, 4, 1, 1, DeviceType>, 4, 1>(instr, locals, device); break; case 139: - op_unweight(instr, locals, device); + batch_foreach, kernel_madnis_single_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 140: - batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + batch_foreach, kernel_madnis_multi_channel_variance, 2, 1, 1, DeviceType>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + op_nonzero(instr, locals, device); break; case 142: - op_vegas_histogram(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 143: + op_batch_scatter(instr, locals, device); + break; +case 144: + op_random(instr, locals, device); + break; +case 145: + op_unweight(instr, locals, device); + break; +case 146: + batch_foreach, kernel_vegas_forward, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 147: + batch_foreach, kernel_vegas_inverse, 2, 2, 2, DeviceType>, 2, 2>(instr, locals, device); + break; +case 148: + op_vegas_histogram(instr, locals, device); + break; +case 149: 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 c0253b0f7..030d74131 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 104: +case 110: backward_op_matmul(instr, locals, local_grads, device); break; -case 105: +case 111: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 106: +case 112: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 107: +case 113: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 108: +case 114: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 109: +case 115: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 110: +case 116: backward_batch_foreach, 2, 1>, 1, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 111: +case 117: backward_op_rqs_reshape(instr, locals, local_grads, device); break; -case 112: +case 118: backward_batch_foreach, 5, 4, 2>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 113: +case 119: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 114: +case 120: backward_batch_foreach, 4, 2, 2>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 115: +case 121: backward_batch_foreach, 2, 1>, 1, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 116: +case 122: backward_batch_foreach, 2, 2, 1>, 2, 1, 0, 1>(instr, locals, local_grads, {}, {0}, device); break; -case 120: +case 126: backward_batch_foreach, 4, 2, 1>, 2, 2, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 123: +case 129: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {0}, {}, device); break; -case 126: +case 132: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 130: +case 136: backward_batch_foreach, 3, 2, 1>, 2, 1, 2, 0>(instr, locals, local_grads, {0,1}, {}, device); break; -case 131: +case 137: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 132: +case 138: backward_batch_foreach, 5, 4, 1>, 4, 1, 4, 0>(instr, locals, local_grads, {0,1,2,3}, {}, device); break; -case 133: +case 139: backward_batch_foreach, 2, 2, 1>, 2, 1, 1, 0>(instr, locals, local_grads, {1}, {}, device); break; -case 134: +case 140: 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 220f3bad3..1068c00bc 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -185,251 +185,269 @@ case 60: batch_foreach, 7, 4, 1>, 7, 4>(instr, locals, device); break; case 61: - batch_foreach, 9, 4, 1>, 9, 4>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 62: - batch_foreach, 3, 10, 1>, 3, 10>(instr, locals, device); + batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); break; case 63: - batch_foreach, 10, 4, 1>, 10, 4>(instr, locals, device); + batch_foreach, 9, 4, 1>, 9, 4>(instr, locals, device); break; case 64: - batch_foreach, 3, 11, 1>, 3, 11>(instr, locals, device); + batch_foreach, 3, 10, 1>, 3, 10>(instr, locals, device); break; case 65: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 10, 4, 1>, 10, 4>(instr, locals, device); break; case 66: - batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); + batch_foreach, 3, 11, 1>, 3, 11>(instr, locals, device); break; case 67: - batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 68: - batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); + batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); break; case 69: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 70: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); break; case 71: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 72: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 5, 3, 1>, 5, 3>(instr, locals, device); break; case 73: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); break; case 74: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 75: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 76: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 77: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 78: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 79: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 80: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 81: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 82: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 83: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 84: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 85: - batch_foreach, 2, 3, 1>, 2, 3>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 86: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 87: - batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 88: - batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 89: - batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 90: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 91: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 3, 1>, 2, 3>(instr, locals, device); break; case 92: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 93: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 94: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 3, 1>, 3, 3>(instr, locals, device); break; case 95: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 2, 1>, 3, 2>(instr, locals, device); break; case 96: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 97: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 98: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 99: - batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 100: - op_matrix_element(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 101: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 102: - batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 103: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 104: - op_matmul(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 105: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 106: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matrix_element(instr, locals, device); break; case 107: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 108: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 6, 1, 1>, 6, 1>(instr, locals, device); break; case 109: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 110: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + op_matmul(instr, locals, device); break; case 111: - op_rqs_reshape(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 112: - batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 113: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 114: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 115: - batch_foreach, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 116: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 117: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + op_rqs_reshape(instr, locals, device); break; case 118: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 4, 1, 2>, 4, 1>(instr, locals, device); break; case 119: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 120: - batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); break; case 121: - op_discrete_histogram(instr, locals, device); + batch_foreach, 1, 1>, 1, 1>(instr, locals, device); break; case 122: - batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 123: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 124: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 125: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 126: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 2, 1>, 2, 2>(instr, locals, device); break; case 127: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + op_discrete_histogram(instr, locals, device); break; case 128: - batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); + batch_foreach, 3, 1, 1>, 3, 1>(instr, locals, device); break; case 129: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 130: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 131: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 132: - batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 133: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 134: - batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); + batch_foreach, 1, 1, 1>, 1, 1>(instr, locals, device); break; case 135: - op_nonzero(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 136: - op_batch_gather(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 137: - op_batch_scatter(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 138: - op_random(instr, locals, device); + batch_foreach, 4, 1, 1>, 4, 1>(instr, locals, device); break; case 139: - op_unweight(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 140: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + batch_foreach, 2, 1, 1>, 2, 1>(instr, locals, device); break; case 141: - batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + op_nonzero(instr, locals, device); break; case 142: - op_vegas_histogram(instr, locals, device); + op_batch_gather(instr, locals, device); break; case 143: + op_batch_scatter(instr, locals, device); + break; +case 144: + op_random(instr, locals, device); + break; +case 145: + op_unweight(instr, locals, device); + break; +case 146: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 147: + batch_foreach, 2, 2, 2>, 2, 2>(instr, locals, device); + break; +case 148: + op_vegas_histogram(instr, locals, device); + break; +case 149: op_histogram(instr, locals, device); break; diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index b5216ec0c..e62dd4e39 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -106,6 +106,31 @@ t_inv_min_max(FVal s, FVal ma_2, FVal mb_2, FVal m1_2, FVal m2_2) return {t_min, t_max}; } +template +KERNELSPEC Pair, FVal> +t1_inv_min_max_doublet(FVal s, FVal m1_2, FVal mir_min_2) { + // |t1| bounds for pa+pb -> p_i + p_ir with massless pa,pb, fixed m_i, + // and m_ir >= mir_min. Same shape as standard t bounds with + // ma=mb=0, m2=mir_min. + return t_inv_min_max(s, 0., 0., m1_2, mir_min_2); +} + +template +KERNELSPEC Pair, FVal> +t2_inv_min_max_doublet(FVal s, FVal m1_2, FVal mir_min_2, FVal t1_abs) { + // |t2| bounds given fixed s, m_i, mir_min, |t1|. + // Derivation (abs-value convention): from the signed-t formulas in the + // Fortran double_t, flip both signs. + // t2_min_abs = m_i^2 * (s - t1_abs - m_i^2) / (m_i^2 + t1_abs) + // t2_max_abs = s - t1_abs - m_i^2 - mir_min^2 + auto denom = m1_2 + t1_abs + EPS; + auto t2_min_raw = m1_2 * (s - t1_abs - m1_2) / denom; + auto t2_max_raw = s - t1_abs - m1_2 - mir_min_2; + auto t2_min = max(t2_min_raw, 0.); + auto t2_max = where(t2_max_raw > t2_min, t2_max_raw, t2_min + EPS); + return {t2_min, t2_max}; +} + template KERNELSPEC FVal lsquare(FourMom p) { return p[0] * p[0] - p[1] * p[1] - p[2] * p[2] - p[3] * p[3]; @@ -484,6 +509,93 @@ KERNELSPEC void kernel_invariants_from_momenta( } } +template +KERNELSPEC void kernel_t1_inv_min_max_doublet( + FIn pa, + FIn pb, + FIn m1, + FIn mir_min, + 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 bounds = t1_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min); + t_min = bounds.first; + t_max = bounds.second; +} + +template +KERNELSPEC void kernel_t1_inv_value_and_min_max_doublet( + FIn pa, + FIn pb, + FIn p1, + FIn mir_min, + 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 m1_2 = lsquare(load_mom(p1)); + auto bounds = t1_inv_min_max_doublet(s, m1_2, mir_min * mir_min); + t_min = bounds.first; + t_max = bounds.second; + t_abs = -lsquare(pa1); +} + +template +KERNELSPEC void kernel_t2_inv_min_max_doublet( + FIn pa, + FIn pb, + FIn m1, + FIn mir_min, + FIn t1_abs, + 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 bounds = + t2_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min, t1_abs); + t_min = bounds.first; + t_max = bounds.second; +} + +template +KERNELSPEC void kernel_t2_inv_value_and_min_max_doublet( + FIn pa, + FIn pb, + FIn p1, + FIn mir_min, + FIn t1_abs, + FOut t_abs, + FOut t_min, + FOut t_max +) { + FourMom pb1, p_tot; + for (int i = 0; i < 4; ++i) { + pb1[i] = pb[i] - p1[i]; + p_tot[i] = pa[i] + pb[i]; + } + auto s = lsquare(p_tot); + auto m1_2 = lsquare(load_mom(p1)); + auto bounds = t2_inv_min_max_doublet(s, m1_2, mir_min * mir_min, t1_abs); + t_min = bounds.first; + t_max = bounds.second; + t_abs = -lsquare(pb1); +} + template KERNELSPEC void kernel_pt_eta_phi_x(FIn p_ext, FIn x1, FIn x2, FOut output) { diff --git a/madspace/src/kernels/twoparticle.hpp b/madspace/src/kernels/twoparticle.hpp index 19463ca60..3ca1168c8 100644 --- a/madspace/src/kernels/twoparticle.hpp +++ b/madspace/src/kernels/twoparticle.hpp @@ -275,5 +275,85 @@ KERNELSPEC void kernel_two_to_two_particle_scattering_inverse( r_phi = phi / PI / 2. + 0.5; } +template +KERNELSPEC void kernel_double_t_scattering( + FIn r_phi, + FIn pa, + FIn pb, + FIn t1_abs, + FIn t2_abs, + FIn m1, + FOut p1, + FOut p2, + FOut det +) { + // pa + pb -> p1 + p2 with |t1| = -(pa-p1)^2, |t2| = -(pb-p1)^2, phi. + // Assumes pa, pb massless. m1 fixed; m2 derived. + // + // Canonical frame: (pa+pb) rest frame, with pa along +z. The formulas: + // E1 = (t1_abs + t2_abs + 2 m1^2) / (2 sqrt(s)) + // pz1 = (t2_abs - t1_abs) / (2 sqrt(s)) + // pt^2 = (t1_abs t2_abs + (t1_abs + t2_abs) m1^2 + m1^4 - s m1^2) / s + // Then rotate so canonical-z aligns with pa_com, boost back to lab. + FourMom p_tot; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + } + auto pa_com = boost(load_mom(pa), p_tot, -1.); + auto s = lsquare(p_tot); + auto sqrts = sqrt(max(s, EPS2)); + auto m1_2 = m1 * m1; + + auto e1 = ( t1_abs + t2_abs + 2. * m1_2) / (2. * sqrts); + auto pz1 = ( t2_abs - t1_abs) / (2. * sqrts); + auto pt2_raw = + (t1_abs * t2_abs + (t1_abs + t2_abs) * m1_2 + m1_2 * m1_2 - s * m1_2) / s; + auto pt2 = max(pt2_raw, 0.); + auto pt = sqrt(pt2); + auto phi = PI * (2. * r_phi - 1.); + + FourMom p1_canon{e1, pt * cos(phi), pt * sin(phi), pz1}; + + auto p1_rot = rotate(p1_canon, pa_com); + auto p1_lab = boost(p1_rot, p_tot, 1.); + store_mom(p1, p1_lab); + for (int i = 0; i < 4; ++i) { + p2[i] = p_tot[i] - p1_lab[i]; + } + + // Jacobian: 1 / (4 sqrt(lambda(s_ir, 0, 0))) with pa, pb massless. + // λ(s_ir, 0, 0) = s_ir^2, so √λ = s_ir. + FourMom p_ir; + for (int i = 0; i < 4; ++i) { + p_ir[i] = p_tot[i] - p1_lab[i]; + } + auto s_ir = lsquare(p_ir); + det = 1. / (4. * max(s_ir, EPS)); +} + +template +KERNELSPEC void kernel_double_t_scattering_inverse( + FIn p1, + FIn p2, + FIn pa, + FIn pb, + FOut r_phi, + FOut m1, + FOut det +) { + FourMom p_tot; + for (int i = 0; i < 4; ++i) { + p_tot[i] = pa[i] + pb[i]; + } + auto pa_com = boost(load_mom(pa), p_tot, -1.); + auto p1_com = boost(load_mom(p1), p_tot, -1.); + auto p1_canon = rotate_inverse(p1_com, pa_com); + auto phi = atan2(p1_canon[2], p1_canon[1]); + r_phi = phi / PI / 2. + 0.5; + m1 = sqrt(max(EPS2, lsquare(load_mom(p1)))); + auto s_ir = lsquare(load_mom(p2)); + det = 4. * s_ir; +} + } // namespace kernels } // namespace madspace diff --git a/madspace/src/phasespace/two_particle.cpp b/madspace/src/phasespace/two_particle.cpp index 634997c75..31c547f50 100644 --- a/madspace/src/phasespace/two_particle.cpp +++ b/madspace/src/phasespace/two_particle.cpp @@ -125,3 +125,81 @@ Mapping::Result TwoToTwoParticleScattering::build_inverse_impl( fb.mul(t_result["det"], det_scatter) }; } + +DoubleT::DoubleT( + double t1_invariant_power, double t1_mass, double t1_width, + double t2_invariant_power, double t2_mass, double t2_width +) : + Mapping( + "DoubleT", + {{"random_phi", batch_float}, + {"random_t1", batch_float}, + {"random_t2", batch_float}, + {"mass1", batch_float}}, + {{"momentum1", batch_four_vec}, {"momentum2", batch_four_vec}}, + {{"momentum_in1", batch_four_vec}, + {"momentum_in2", batch_four_vec}, + {"mass_rest_min", batch_float}} + ), + _t1_invariant(t1_invariant_power, t1_mass, t1_width), + _t2_invariant(t2_invariant_power, t2_mass, t2_width) {} + +Mapping::Result DoubleT::build_forward_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions +) const { + auto r_phi = inputs.at(0), r_t1 = inputs.at(1), r_t2 = inputs.at(2); + auto m1 = inputs.at(3); + auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); + auto mir_min = conditions.at(2); + + auto [t1_min, t1_max] = fb.t1_inv_min_max_doublet(p_in1, p_in2, m1, mir_min); + 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_result = _t2_invariant.build_forward(fb, {r_t2}, {t2_min, t2_max}); + + auto [p1, p2, det_scatter] = fb.double_t_scattering( + r_phi, p_in1, p_in2, + t1_result["invariant"], t2_result["invariant"], m1 + ); + + auto det_inv = fb.mul(t1_result["det"], t2_result["det"]); + return { + {{"momentum1", p1}, {"momentum2", p2}}, + fb.mul(det_inv, det_scatter) + }; +} + +Mapping::Result DoubleT::build_inverse_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions +) const { + auto p1 = inputs.at(0), p2 = inputs.at(1); + auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); + auto mir_min = conditions.at(2); + + auto [r_phi, m1, 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, mir_min); + 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, mir_min, t1_abs + ); + 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"]); + return { + {{"random_phi", r_phi}, + {"random_t1", t1_result["random"]}, + {"random_t2", t2_result["random"]}, + {"mass1", m1}}, + fb.mul(det_inv, det_scatter) + }; +} diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index ac48f4d7a..3b10ec7bf 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -74,12 +74,18 @@ void add_instructions(py::classh& fb) { fb.def("two_to_two_particle_scattering_inverse", &FunctionBuilder::two_to_two_particle_scattering_inverse, py::arg("p1"), py::arg("p2"), py::arg("pa"), py::arg("pb")); fb.def("two_to_three_particle_scattering", &FunctionBuilder::two_to_three_particle_scattering, py::arg("phi_choice"), py::arg("pa"), py::arg("pb"), py::arg("p3"), py::arg("s23"), py::arg("t1_abs"), py::arg("m1"), py::arg("m2")); fb.def("two_to_three_particle_scattering_inverse", &FunctionBuilder::two_to_three_particle_scattering_inverse, py::arg("p1"), py::arg("p2"), py::arg("p3"), py::arg("pa"), py::arg("pb"), py::arg("t1_abs"), py::arg("s23")); + fb.def("double_t_scattering", &FunctionBuilder::double_t_scattering, py::arg("r_phi"), py::arg("pa"), py::arg("pb"), py::arg("t1_abs"), py::arg("t2_abs"), py::arg("m1")); + fb.def("double_t_scattering_inverse", &FunctionBuilder::double_t_scattering_inverse, py::arg("p1"), py::arg("p2"), py::arg("pa"), py::arg("pb")); fb.def("three_body_decay_com", &FunctionBuilder::three_body_decay_com, py::arg("r_e1"), py::arg("r_e2"), py::arg("r_phi"), py::arg("r_cos_theta"), py::arg("r_beta"), py::arg("m0"), py::arg("m1"), py::arg("m2"), py::arg("m3")); fb.def("three_body_decay_com_inverse", &FunctionBuilder::three_body_decay_com_inverse, py::arg("p1"), py::arg("p2"), py::arg("p3")); fb.def("three_body_decay", &FunctionBuilder::three_body_decay, py::arg("r_e1"), py::arg("r_e2"), py::arg("r_phi"), py::arg("r_cos_theta"), py::arg("r_beta"), py::arg("m0"), py::arg("m1"), py::arg("m2"), py::arg("m3"), py::arg("p0")); 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("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("mir_min"), py::arg("t1_abs")); 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("invariants_from_momenta", &FunctionBuilder::invariants_from_momenta, py::arg("p_ext"), py::arg("factors")); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 8f5d08039..9d865bf77 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -438,6 +438,17 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("width") = 0. ); + py::classh(m, "DoubleT") + .def( + py::init(), + py::arg("t1_invariant_power") = 0., + py::arg("t1_mass") = 0., + py::arg("t1_width") = 0., + py::arg("t2_invariant_power") = 0., + py::arg("t2_mass") = 0., + py::arg("t2_width") = 0. + ); + py::classh(m, "ThreeBodyDecay") .def(py::init(), py::arg("com")); diff --git a/madspace/tests/test_double_t.py b/madspace/tests/test_double_t.py new file mode 100644 index 000000000..30ad266ff --- /dev/null +++ b/madspace/tests/test_double_t.py @@ -0,0 +1,257 @@ +from collections import namedtuple + +import numpy as np +import pytest +from pytest import approx + +import madspace as ms + +# ---------------------------- +# Utilities +# ---------------------------- + + +def inv_mass_sq(p): + return p[..., 0] ** 2 - np.sum(p[..., 1:] ** 2, axis=-1) + + +def mass(p): + return np.sqrt(np.maximum(inv_mass_sq(p), 0.0)) + + +# ---------------------------- +# Fixtures +# ---------------------------- + +N = 10_000 + +InputPoint = namedtuple( + "InputPoint", + [ + "r_phi", + "r_t1", + "r_t2", + "m1", + "mir_min", + "pa", + "pb", + ], +) + + +@pytest.fixture +def rng(): + return np.random.default_rng(1234) + + +@pytest.fixture( + params=[False, True], + ids=["LAB", "COM"], +) +def input_points(rng, request): + """Generate (pa, pb) and a set of forward inputs for the DoubleT mapping. + + Picks m1, mir_min so that the kinematic ranges of t1 and t2 are non-empty. + """ + com = request.param + + pz = rng.uniform(1000.0, 4000.0, N) + e1 = pz # massless beams in COM + e2 = pz + pa = np.stack([e1, np.zeros(N), np.zeros(N), +pz], axis=1) + pb = np.stack([e2, np.zeros(N), np.zeros(N), -pz], axis=1) + + if not com: + # boost to a random lab frame (the same recipe used in the 2->3 test) + pa3 = np.array([[0, 0, 4000.0]]) + rng.normal(0.0, 500.0, (N, 3)) + ea = np.sqrt(np.sum(pa3**2, axis=1)) # massless + pa = np.concatenate([ea[:, None], pa3], axis=1) + pb3 = np.array([[0, 0, -4000.0]]) + rng.normal(0.0, 500.0, (N, 3)) + eb = np.sqrt(np.sum(pb3**2, axis=1)) + pb = np.concatenate([eb[:, None], pb3], axis=1) + + # m1: the fixed single-particle mass + # mir_min: lower bound on recoil mass (e.g. sum of constituent masses). + # Pick mir_min well below the sqrt(s) so there's room for the recoil. + m1 = rng.uniform(1.0, 50.0, N) + mir_min = rng.uniform(50.0, 200.0, N) + + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + return InputPoint(r_phi, r_t1, r_t2, m1, mir_min, pa, pb) + + +@pytest.fixture( + params=[ + (np.full(N, 5.0), np.full(N, 80.0)), + (np.full(N, 50.0), np.full(N, 150.0)), + (np.full(N, 100.0), np.full(N, 250.0)), + ], + ids=[ + "light single, light recoil min", + "medium single, medium recoil min", + "heavy single, heavy recoil min", + ], +) +def fixed_input_points(rng, request): + """Fixed (m1, mir_min) used by the phase-space-volume test.""" + m1, mir_min = request.param + + pz = np.full(N, 3000.0) + e1 = pz + e2 = pz + pa = np.stack([e1, np.zeros(N), np.zeros(N), +pz], axis=1) + pb = np.stack([e2, np.zeros(N), np.zeros(N), -pz], axis=1) + + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + return InputPoint(r_phi, r_t1, r_t2, m1, mir_min, pa, pb) + + +# ---------------------------- +# Tests +# ---------------------------- + + +def test_momentum_conservation(input_points): + mapping = ms.DoubleT() + + inputs = [ + input_points.r_phi, + input_points.r_t1, + input_points.r_t2, + input_points.m1, + ] + conditions = [input_points.pa, input_points.pb, input_points.mir_min] + + p1, p2, det = mapping.map_forward(inputs, conditions) + assert p1 + p2 == approx(input_points.pa + input_points.pb) + + +def test_on_shell_mass(input_points): + mapping = ms.DoubleT() + + inputs = [ + input_points.r_phi, + input_points.r_t1, + input_points.r_t2, + input_points.m1, + ] + conditions = [input_points.pa, input_points.pb, input_points.mir_min] + + p1, p2, det = mapping.map_forward(inputs, conditions) + assert mass(p1) == approx(input_points.m1) + + +def test_recoil_mass_minimum(input_points): + """The recoil mass m_ir should be >= mir_min, since that bound goes + into the t1/t2 ranges.""" + mapping = ms.DoubleT() + + inputs = [ + input_points.r_phi, + input_points.r_t1, + input_points.r_t2, + input_points.m1, + ] + conditions = [input_points.pa, input_points.pb, input_points.mir_min] + + p1, p2, det = mapping.map_forward(inputs, conditions) + m_ir = mass(p2) + # allow a small numerical tolerance below mir_min + assert np.all(m_ir >= input_points.mir_min - 1e-6 * input_points.mir_min) + + +def test_inverse(input_points): + mapping = ms.DoubleT() + + inputs = [ + input_points.r_phi, + input_points.r_t1, + input_points.r_t2, + input_points.m1, + ] + conditions = [input_points.pa, input_points.pb, input_points.mir_min] + + p1, p2, det = mapping.map_forward(inputs, conditions) + *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) + + assert inv_det == approx(1 / det, rel=1e-5) + + for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): + assert inp == approx(inv_inp), f"mismatch in input index {i}" + + +def test_phase_space_volume(fixed_input_points): + """For massless beams and the double-t parameterization, the phase-space + volume over all (t1, t2, phi) at fixed (s, m1, mir_min) should match the + analytic 2-body phase-space volume integrated against the constraint + that m_ir >= mir_min. + + Simpler check: the integral of the det over the (r_phi, r_t1, r_t2) + cube should be independent of how we parameterize, i.e. equal the + same fixed kinematic volume. We compare it to the standard 2->2 volume + pi/(2s) * sqrt(lambda(s, m1^2, mir_min^2)) ... but the recoil mass is + NOT fixed in double_t, so this comparison is approximate. + + Instead: just check the determinant has a finite, frame-independent mean + with the expected scaling. The mean should be positive and finite.""" + mapping = ms.DoubleT() + + inputs = [ + fixed_input_points.r_phi, + fixed_input_points.r_t1, + fixed_input_points.r_t2, + fixed_input_points.m1, + ] + conditions = [ + fixed_input_points.pa, + fixed_input_points.pb, + fixed_input_points.mir_min, + ] + p1, p2, det = mapping.map_forward(inputs, conditions) + + assert np.all(np.isfinite(det)) + assert np.all(det > 0) + # variance should be finite and well below the mean (sanity check on numerics) + std_error = np.std(det) / np.sqrt(N) + assert std_error < np.mean(det) + + +def test_inverse_lab_specifically(rng): + """Explicit LAB test: round-trip should work in any frame, not just COM. + + This was the failing case for the 2->3 mapping before we fixed the + rotation; keep it as a guard for double_t too. + """ + mapping = ms.DoubleT() + + # boosted lab frame + pz = rng.uniform(1000.0, 4000.0, N) + pa3 = np.array([[0, 0, 4000.0]]) + rng.normal(0.0, 500.0, (N, 3)) + ea = np.sqrt(np.sum(pa3**2, axis=1)) + pa = np.concatenate([ea[:, None], pa3], axis=1) + pb3 = np.array([[0, 0, -4000.0]]) + rng.normal(0.0, 500.0, (N, 3)) + eb = np.sqrt(np.sum(pb3**2, axis=1)) + pb = np.concatenate([eb[:, None], pb3], axis=1) + + m1 = np.full(N, 30.0) + mir_min = np.full(N, 100.0) + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + inputs = [r_phi, r_t1, r_t2, m1] + conditions = [pa, pb, mir_min] + + p1, p2, det = mapping.map_forward(inputs, conditions) + *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) + + names = ["r_phi", "r_t1", "r_t2", "m1"] + for name, inp, inv_inp in zip(names, inputs, inv_inputs): + assert inp == approx(inv_inp), f"LAB mismatch in {name}" + assert inv_det == approx(1 / det, rel=1e-5) \ No newline at end of file From 8305bb4b4acfccee032882c8603e6994e1d4054f Mon Sep 17 00:00:00 2001 From: Ramon Winterhalder Date: Mon, 11 May 2026 22:08:54 +0200 Subject: [PATCH 3/6] Runs but needs mass fix --- madspace/CMakeLists.txt | 2 + .../compgraphs/function_builder_mixin.inc | 12 +- madspace/include/madspace/phasespace.hpp | 1 + .../phasespace/t_propagator_mapping_23.hpp | 55 ++ madspace/instruction_set.yaml | 9 +- .../src/compgraphs/instruction_set_mixin.inc | 6 +- madspace/src/cpu/runtime_mixin.inc | 6 +- madspace/src/gpu/runtime_mixin.inc | 6 +- madspace/src/kernels/kinematics.hpp | 8 +- madspace/src/kernels/twoparticle.hpp | 15 +- .../phasespace/t_propagator_mapping_23.cpp | 646 ++++++++++++++++++ madspace/src/phasespace/two_particle.cpp | 20 +- madspace/src/python/instruction_set.hpp | 4 +- madspace/src/python/madspace.cpp | 12 +- madspace/tests/test_2to3_scattering.py | 2 +- madspace/tests/test_double_t.py | 210 +++++- madspace/tests/test_t_propagator_23.py | 264 +++++++ 17 files changed, 1202 insertions(+), 76 deletions(-) create mode 100644 madspace/include/madspace/phasespace/t_propagator_mapping_23.hpp create mode 100644 madspace/src/phasespace/t_propagator_mapping_23.cpp create mode 100644 madspace/tests/test_t_propagator_23.py diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index d8231d185..029b0bcdd 100644 --- a/madspace/CMakeLists.txt +++ b/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/t_propagator_mapping_23.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/t_propagator_mapping_23.hpp include/madspace/phasespace/phasespace.hpp include/madspace/phasespace/multichannel.hpp include/madspace/phasespace/madnis.hpp diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc index d7a47e3c5..dea526d9e 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/madspace/include/madspace/compgraphs/function_builder_mixin.inc @@ -265,9 +265,9 @@ std::array double_t_scattering(Value r_phi, Value pa, Value pb, Value return {output_vector[0], output_vector[1], output_vector[2]}; } -std::array double_t_scattering_inverse(Value p1, Value p2, Value pa, Value pb) { +std::array 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], output_vector[2]}; + return {output_vector[0], output_vector[1]}; } std::array 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) { @@ -305,8 +305,8 @@ std::array t1_inv_min_max_doublet(Value pa, Value pb, Value m1, Value return {output_vector[0], output_vector[1]}; } -std::array t1_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value mir_min) { - auto output_vector = instruction("t1_inv_value_and_min_max_doublet", {pa, pb, p1, mir_min}); +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}); return {output_vector[0], output_vector[1], output_vector[2]}; } @@ -315,8 +315,8 @@ std::array t2_inv_min_max_doublet(Value pa, Value pb, Value m1, Value return {output_vector[0], output_vector[1]}; } -std::array t2_inv_value_and_min_max_doublet(Value pa, Value pb, Value p1, Value mir_min, Value t1_abs) { - auto output_vector = instruction("t2_inv_value_and_min_max_doublet", {pa, pb, p1, mir_min, t1_abs}); +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}); return {output_vector[0], output_vector[1], output_vector[2]}; } diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index 3746db95f..4ae860a54 100644 --- a/madspace/include/madspace/phasespace.hpp +++ b/madspace/include/madspace/phasespace.hpp @@ -22,6 +22,7 @@ #include "phasespace/rambo.hpp" #include "phasespace/scale.hpp" #include "phasespace/t_propagator_mapping.hpp" +#include "madspace/phasespace/t_propagator_mapping_23.hpp" #include "phasespace/three_particle.hpp" #include "phasespace/topology.hpp" #include "phasespace/two_particle.hpp" diff --git a/madspace/include/madspace/phasespace/t_propagator_mapping_23.hpp b/madspace/include/madspace/phasespace/t_propagator_mapping_23.hpp new file mode 100644 index 000000000..5b3d37cdb --- /dev/null +++ b/madspace/include/madspace/phasespace/t_propagator_mapping_23.hpp @@ -0,0 +1,55 @@ +#pragma once + +#include + +#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 TPropagatorMapping23 : public Mapping { +public: + // colour_order: 0-indexed permutation of {0, ..., n-1} (n = n_out + 2). + // Particles 0 and 1 are the two incoming beams. + TPropagatorMapping23( + const std::vector& colour_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& inputs, + const NamedVector& conditions + ) const override; + Result build_inverse_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& 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 _set1; + std::vector _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 diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 4e4d22dce..453e24c97 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -1013,9 +1013,6 @@ double_t_scattering_inverse: - name: r_phi type: [float] desc: - - name: m1 - type: [float] - desc: - name: det type: [float] desc: @@ -1278,6 +1275,9 @@ t1_inv_value_and_min_max_doublet: - name: p1 type: [float, 4] desc: + - name: m1 + type: [float] + desc: - name: mir_min type: [float] desc: @@ -1328,6 +1328,9 @@ t2_inv_value_and_min_max_doublet: - name: p1 type: [float, 4] desc: + - name: m1 + type: [float] + desc: - name: mir_min type: [float] desc: diff --git a/madspace/src/compgraphs/instruction_set_mixin.inc b/madspace/src/compgraphs/instruction_set_mixin.inc index 3587c7a45..cdaf72561 100644 --- a/madspace/src/compgraphs/instruction_set_mixin.inc +++ b/madspace/src/compgraphs/instruction_set_mixin.inc @@ -77,7 +77,7 @@ InstructionOwner instructions[] { mi("two_to_three_particle_scattering", 59, true, {{DataType::dt_int, false, {}, 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, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), mi("two_to_three_particle_scattering_inverse", 60, 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, {4}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {DataType::dt_int, false, {}, false}, {DataType::dt_float, false, {}, false}}), mi("double_t_scattering", 61, true, {{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, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), - mi("double_t_scattering_inverse", 62, 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("double_t_scattering_inverse", 62, 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}}), mi("three_body_decay_com", 63, 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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}}, {{DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {4}, false}, {DataType::dt_float, false, {}, false}}), mi("three_body_decay_com_inverse", 64, 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, {}, false}}), mi("three_body_decay", 65, 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}, {DataType::dt_float, false, {}, false}, {DataType::dt_float, false, {}, false}, {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}}), @@ -85,9 +85,9 @@ InstructionOwner instructions[] { 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}}), + 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}}), + 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}}), diff --git a/madspace/src/cpu/runtime_mixin.inc b/madspace/src/cpu/runtime_mixin.inc index 669e0fa97..90fdeac07 100644 --- a/madspace/src/cpu/runtime_mixin.inc +++ b/madspace/src/cpu/runtime_mixin.inc @@ -188,7 +188,7 @@ case 61: batch_foreach, kernel_double_t_scattering, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 62: - batch_foreach, kernel_double_t_scattering_inverse, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); + batch_foreach, kernel_double_t_scattering_inverse, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 63: batch_foreach, kernel_three_body_decay_com, 9, 4, 1, DeviceType>, 9, 4>(instr, locals, device); @@ -212,13 +212,13 @@ case 69: batch_foreach, kernel_t1_inv_min_max_doublet, 4, 2, 1, DeviceType>, 4, 2>(instr, locals, device); break; case 70: - batch_foreach, kernel_t1_inv_value_and_min_max_doublet, 4, 3, 1, DeviceType>, 4, 3>(instr, locals, device); + batch_foreach, kernel_t1_inv_value_and_min_max_doublet, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device); break; case 71: batch_foreach, kernel_t2_inv_min_max_doublet, 5, 2, 1, DeviceType>, 5, 2>(instr, locals, device); break; case 72: - batch_foreach, kernel_t2_inv_value_and_min_max_doublet, 5, 3, 1, DeviceType>, 5, 3>(instr, locals, device); + batch_foreach, kernel_t2_inv_value_and_min_max_doublet, 6, 3, 1, DeviceType>, 6, 3>(instr, locals, device); break; case 73: batch_foreach, kernel_s23_min_max, 6, 2, 1, DeviceType>, 6, 2>(instr, locals, device); diff --git a/madspace/src/gpu/runtime_mixin.inc b/madspace/src/gpu/runtime_mixin.inc index 1068c00bc..3f3927e29 100644 --- a/madspace/src/gpu/runtime_mixin.inc +++ b/madspace/src/gpu/runtime_mixin.inc @@ -188,7 +188,7 @@ case 61: batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 62: - batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); + batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 63: batch_foreach, 9, 4, 1>, 9, 4>(instr, locals, device); @@ -212,13 +212,13 @@ case 69: batch_foreach, 4, 2, 1>, 4, 2>(instr, locals, device); break; case 70: - batch_foreach, 4, 3, 1>, 4, 3>(instr, locals, device); + batch_foreach, 5, 3, 1>, 5, 3>(instr, locals, device); break; case 71: batch_foreach, 5, 2, 1>, 5, 2>(instr, locals, device); break; case 72: - batch_foreach, 5, 3, 1>, 5, 3>(instr, locals, device); + batch_foreach, 6, 3, 1>, 6, 3>(instr, locals, device); break; case 73: batch_foreach, 6, 2, 1>, 6, 2>(instr, locals, device); diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index e62dd4e39..a4f3a38f1 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -533,6 +533,7 @@ KERNELSPEC void kernel_t1_inv_value_and_min_max_doublet( FIn pa, FIn pb, FIn p1, + FIn m1, FIn mir_min, FOut t_abs, FOut t_min, @@ -544,8 +545,7 @@ KERNELSPEC void kernel_t1_inv_value_and_min_max_doublet( p_tot[i] = pa[i] + pb[i]; } auto s = lsquare(p_tot); - auto m1_2 = lsquare(load_mom(p1)); - auto bounds = t1_inv_min_max_doublet(s, m1_2, mir_min * mir_min); + auto bounds = t1_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min); t_min = bounds.first; t_max = bounds.second; t_abs = -lsquare(pa1); @@ -577,6 +577,7 @@ KERNELSPEC void kernel_t2_inv_value_and_min_max_doublet( FIn pa, FIn pb, FIn p1, + FIn m1, FIn mir_min, FIn t1_abs, FOut t_abs, @@ -589,8 +590,7 @@ KERNELSPEC void kernel_t2_inv_value_and_min_max_doublet( p_tot[i] = pa[i] + pb[i]; } auto s = lsquare(p_tot); - auto m1_2 = lsquare(load_mom(p1)); - auto bounds = t2_inv_min_max_doublet(s, m1_2, mir_min * mir_min, t1_abs); + auto bounds = t2_inv_min_max_doublet(s, m1 * m1, mir_min * mir_min, t1_abs); t_min = bounds.first; t_max = bounds.second; t_abs = -lsquare(pb1); diff --git a/madspace/src/kernels/twoparticle.hpp b/madspace/src/kernels/twoparticle.hpp index 3ca1168c8..39a46885d 100644 --- a/madspace/src/kernels/twoparticle.hpp +++ b/madspace/src/kernels/twoparticle.hpp @@ -321,14 +321,7 @@ KERNELSPEC void kernel_double_t_scattering( p2[i] = p_tot[i] - p1_lab[i]; } - // Jacobian: 1 / (4 sqrt(lambda(s_ir, 0, 0))) with pa, pb massless. - // λ(s_ir, 0, 0) = s_ir^2, so √λ = s_ir. - FourMom p_ir; - for (int i = 0; i < 4; ++i) { - p_ir[i] = p_tot[i] - p1_lab[i]; - } - auto s_ir = lsquare(p_ir); - det = 1. / (4. * max(s_ir, EPS)); + det = PI / (2. * max(s, EPS)); } template @@ -338,7 +331,6 @@ KERNELSPEC void kernel_double_t_scattering_inverse( FIn pa, FIn pb, FOut r_phi, - FOut m1, FOut det ) { FourMom p_tot; @@ -350,9 +342,8 @@ KERNELSPEC void kernel_double_t_scattering_inverse( auto p1_canon = rotate_inverse(p1_com, pa_com); auto phi = atan2(p1_canon[2], p1_canon[1]); r_phi = phi / PI / 2. + 0.5; - m1 = sqrt(max(EPS2, lsquare(load_mom(p1)))); - auto s_ir = lsquare(load_mom(p2)); - det = 4. * s_ir; + auto s = lsquare(p_tot); + det = 2. * max(s, EPS) / PI; } } // namespace kernels diff --git a/madspace/src/phasespace/t_propagator_mapping_23.cpp b/madspace/src/phasespace/t_propagator_mapping_23.cpp new file mode 100644 index 000000000..d539e3cc3 --- /dev/null +++ b/madspace/src/phasespace/t_propagator_mapping_23.cpp @@ -0,0 +1,646 @@ +#include "madspace/phasespace/t_propagator_mapping_23.hpp" + +#include +#include + +#include "madspace/util.hpp" + +using namespace madspace; + +namespace { + +// Cyclically rotate the colour order so that particle 0 is first, then +// split the rest into two sets: particles strictly between 0 and 1 in the +// cyclic order go into set1, particles after 1 (wrapping around) go into +// set2. Both sets are returned as 0-indexed outgoing-particle indices, +// i.e. particle p (0-indexed external) becomes (p - 2) in our 0-indexed +// convention for outgoing particles. (Particles 0 and 1 are the beams +// and not in either set.) +std::pair, std::vector> +split_sets_from_colour_order(const std::vector& colour_order) { + std::size_t n = colour_order.size(); + if (n < 4) { + throw std::invalid_argument( + "TPropagatorMapping23 requires at least 4 particles (2 beams + 2 outgoing)" + ); + } + // Find index of particle 0 and rotate + auto it = std::find(colour_order.begin(), colour_order.end(), 0u); + if (it == colour_order.end()) { + throw std::invalid_argument("colour_order must contain particle 0"); + } + std::size_t i0 = std::distance(colour_order.begin(), it); + std::vector rotated; + rotated.reserve(n); + for (std::size_t k = 0; k < n; ++k) { + rotated.push_back(colour_order[(i0 + k) % n]); + } + // rotated[0] is now particle 0. Find particle 1. + auto it1 = std::find(rotated.begin(), rotated.end(), 1u); + if (it1 == rotated.end()) { + throw std::invalid_argument("colour_order must contain particle 1"); + } + std::size_t i1 = std::distance(rotated.begin(), it1); + // set1 = rotated[1..i1-1], set2 = rotated[i1+1..n-1] + std::vector set1, set2; + for (std::size_t k = 1; k < i1; ++k) { + std::size_t p = rotated[k]; // 0-indexed external + if (p <= 1) throw std::invalid_argument("invalid colour_order"); + set1.push_back(p - 2); // 0-indexed outgoing + } + for (std::size_t k = i1 + 1; k < n; ++k) { + std::size_t p = rotated[k]; + if (p <= 1) throw std::invalid_argument("invalid colour_order"); + set2.push_back(p - 2); + } + if (set1.empty() || set2.empty()) { + throw std::invalid_argument( + "TPropagatorMapping23 requires both sets to be non-empty " + "(particles 0 and 1 must not be adjacent in the colour order)" + ); + } + return {set1, set2}; +} + +// Number of intermediate rest masses to sample for a walk of size k: +// max(0, k - 2) +std::size_t n_intermediate_masses_for_set_size(std::size_t k) { + return (k >= 2) ? (k - 2) : 0; +} + +// Number of randoms a walk of size k consumes for its t-channel-block calls +// (excluding mass samples). +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); +} + +} // namespace + +TPropagatorMapping23::TPropagatorMapping23( + const std::vector& colour_order, + double t_invariant_power, + double s_invariant_power +) : + Mapping( + "TPropagatorMapping23", + [&] { + auto [s1, s2] = split_sets_from_colour_order(colour_order); + bool use_double_t = (s1.size() == 1) != (s2.size() == 1); + // When DoubleT is used: neither set-mass is sampled (the + // single-particle side has fixed mass, the multi-particle side's + // mass is derived by DoubleT), and the central block consumes + // 3 randoms (r_phi, r_t1, r_t2). Otherwise: 1 set-mass per + // multi-particle side, plus 2 randoms for the central 2->2. + std::size_t n_set_masses = use_double_t + ? 0u + : (s1.size() >= 2 ? 1u : 0u) + (s2.size() >= 2 ? 1u : 0u); + std::size_t n_central = use_double_t ? 3u : 2u; + std::size_t n_intermediate_masses = + n_intermediate_masses_for_set_size(s1.size()) + + n_intermediate_masses_for_set_size(s2.size()); + std::size_t n_walk = + n_block_randoms_for_set_size(s1.size()) + + n_block_randoms_for_set_size(s2.size()); + std::size_t total = n_set_masses + n_intermediate_masses + + n_central + n_walk; + NamedVector input_types; + for (std::size_t i = 0; i < total; ++i) { + input_types.push_back(std::format("random{}", i), batch_float); + } + return input_types; + }(), + [&] { + std::size_t n_out = colour_order.size() - 2; + NamedVector output_types; + for (std::size_t i = 0; i < n_out + 2; ++i) { + output_types.push_back(std::format("momentum{}", i), batch_four_vec); + } + return output_types; + }(), + [&] { + std::size_t n_out = colour_order.size() - 2; + NamedVector cond_types{{"com_energy", batch_float}}; + for (std::size_t i = 0; i < n_out; ++i) { + cond_types.push_back(std::format("mass{}", i), batch_float); + } + return cond_types; + }() + ), + _n_out(colour_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.) { + auto [s1, s2] = split_sets_from_colour_order(colour_order); + _set1 = s1; + _set2 = s2; + _use_double_t = (s1.size() == 1) != (s2.size() == 1); + std::size_t n_set_masses = _use_double_t + ? 0u + : (s1.size() >= 2 ? 1u : 0u) + (s2.size() >= 2 ? 1u : 0u); + std::size_t n_central = _use_double_t ? 3u : 2u; + std::size_t n_intermediate_masses = + n_intermediate_masses_for_set_size(s1.size()) + + n_intermediate_masses_for_set_size(s2.size()); + 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; +} + +// === Helpers used in both build_forward and build_inverse === +namespace { + +// Sum of masses of particles at indices `idxs` (0-indexed outgoing) in m_out. +Value sum_of_masses( + FunctionBuilder& fb, const ValueVec& m_out, const std::vector& idxs +) { + Value s = m_out.at(idxs[0]); + for (std::size_t k = 1; k < idxs.size(); ++k) { + s = fb.add(s, m_out.at(idxs[k])); + } + return s; +} + +} // namespace + +Mapping::Result TPropagatorMapping23::build_forward_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions +) const { + Value e_cm = conditions.at(0); + ValueVec m_out(conditions.begin() + 1, conditions.end()); + auto r = inputs.begin(); + auto next_random = [&]() { return *(r++); }; + ValueVec dets; + + // ============================================================ + // Phase 1a: pre-sample set composite masses. + // When _use_double_t: neither set mass is sampled here. The single- + // particle side has fixed mass; the multi-particle side's mass is + // derived later from the DoubleT central block. + // Otherwise: sample multi-particle set masses, fix single-particle ones. + // ============================================================ + Value m_set1, m_set2; + Value mass_sum_set1 = sum_of_masses(fb, m_out, _set1); + Value mass_sum_set2 = sum_of_masses(fb, m_out, _set2); + if (_use_double_t) { + // Both set masses set later (multi-side from DoubleT output, + // single-side fixed). + if (_set1.size() == 1) m_set1 = m_out.at(_set1[0]); + if (_set2.size() == 1) m_set2 = m_out.at(_set2[0]); + // The multi-side m_set will be filled in after the DoubleT call, + // before sampling intermediate rest masses. + } else { + if (_set1.size() >= 2) { + auto s_min = fb.square(mass_sum_set1); + 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]); + } + if (_set2.size() >= 2) { + auto s_min = fb.square(mass_sum_set2); + 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]); + } + } + + // ============================================================ + // Phase 2: central block producing (P_set1, P_set2). + // ============================================================ + auto [pa, pb] = fb.com_p_in(e_cm); + Value P_set1, P_set2; + if (_use_double_t) { + // DoubleT: pa, pb -> (single, recoil). The single particle's mass + // is m_single (fixed); the recoil mass must be >= mir_min (sum of + // recoil-side outgoing masses). + 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; + auto central = _double_t.build_forward( + fb, + {next_random(), next_random(), next_random()}, + {pa, pb, m_single, mir_min} + ); + Value p_single = central.at(0); // mass = m_single + Value p_recoil = central.at(1); // recoil; mass derived from t1, t2 + dets.push_back(central["det"]); + if (single_is_set1) { + P_set1 = p_single; + P_set2 = p_recoil; + // m_set1 already set above. Derive m_set2 from p_recoil for + // downstream intermediate-mass sampling. invariants_from_momenta + // returns m^2 for the summed momenta indicated by the factor row. + nested_vector2 factors_recoil{{1.0}}; + auto m2 = fb.unstack( + fb.invariants_from_momenta(fb.stack({p_recoil}), factors_recoil) + ).at(0); + m_set2 = fb.sqrt(m2); + } else { + P_set2 = p_single; + P_set1 = p_recoil; + nested_vector2 factors_recoil{{1.0}}; + auto m2 = fb.unstack( + fb.invariants_from_momenta(fb.stack({p_recoil}), factors_recoil) + ).at(0); + m_set1 = fb.sqrt(m2); + } + } else { + auto central = _com_scattering.build_forward( + fb, + {next_random(), next_random(), m_set1, m_set2}, + {pa, pb} + ); + P_set1 = central.at(0); + P_set2 = central.at(1); + dets.push_back(central["det"]); + } + + // The "other beam after central block" for each set's walk: + // by 4-momentum conservation, p_other_beam - P_other_set is what is + // (kinematically) left on that beam line after the central block emitted + // the other set. + Value R_b_for_set1 = fb.sub(pb, P_set2); + Value R_b_for_set2 = fb.sub(pa, P_set1); + + // ============================================================ + // Phase 1b: pre-sample intermediate rest masses for each walk. + // For a walk peeling particles set[0], set[1], ..., set[k-1]: + // - mass of rest after peeling set[0] = m_rest_1 (sampled) + // - mass of rest after peeling set[0..1] = m_rest_2 (sampled) + // - ... + // - mass of rest after peeling set[0..k-3] = m_rest_{k-2} (sampled) + // - mass of rest after peeling set[0..k-2] = m of set[k-1] (fixed, residual) + // Bounds: + // m_min = sum of masses of [set[j+1], ..., set[k-1]] + // m_max = m_set - sum of masses of [set[0], ..., set[j]] + // Note: we sample these AFTER the central block, because for DoubleT + // we don't know m_set (multi-side) until then. + // ============================================================ + auto sample_intermediate_masses = + [&](const std::vector& s, Value m_set) -> ValueVec { + std::size_t k = s.size(); + ValueVec res; + if (k <= 2) { + return res; // no intermediate masses needed + } + // The walk peels particles s[0], s[1], ..., s[k-2] (s[k-1] is the + // final residual). After peeling s[0..j], the "rest" has mass + // m_rest_{j+1}. The chain is monotonic: + // m_set > m_rest_1 > m_rest_2 > ... > m_rest_{k-2} > mass(s[k-1]) + // We sample m_rest_{j+1} from + // m_min = sum of masses [s[j+1]..s[k-1]] (kinematic lower bound) + // m_max = m_rest_j - mass[s[j]] (monotonicity + peel-off room) + // where m_rest_0 = m_set. + // Sampling independently from [m_min, m_set - sum_of_already_peeled] + // would over-count by (k-2)! because it ignores the ordering. + Value prev_mass = m_set; + for (std::size_t j = 0; j < k - 2; ++j) { + Value m_min = m_out.at(s[j + 1]); + for (std::size_t i = j + 2; i < k; ++i) { + m_min = fb.add(m_min, m_out.at(s[i])); + } + auto s_min = fb.square(m_min); + auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s[j]))); + auto r = _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); + Value m_rest = fb.sqrt(r["invariant"]); + res.push_back(m_rest); + dets.push_back(r["det"]); + prev_mass = m_rest; + } + return res; + }; + + ValueVec rest_masses_set1 = sample_intermediate_masses(_set1, m_set1); + ValueVec rest_masses_set2 = sample_intermediate_masses(_set2, m_set2); + + // ============================================================ + // Phase 3: peel-off walks + // ============================================================ + ValueVec p_out(_n_out); // 0-indexed by outgoing particle index + + auto walk = [&](const std::vector& s, + Value P_set, + Value R_b, + const ValueVec& rest_masses) { + std::size_t k = s.size(); + if (k == 1) { + p_out[s[0]] = P_set; + 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) { + // Peel particle s[j]. The "new rest" after this peel has either + // mass = rest_masses[j] (intermediate) or mass = m_out[s[k-1]] + // (the final residual, when 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]); + // Block input: (R_b, R_a). Block output: p1_out + p2_out = R_b + R_a, + // where p1_out has mass m_rest (the chain carrier, sits on the pa=R_b + // side per the block's convention) and p2_out has mass m_peel (the + // newly peeled particle, sits on the pb=R_a side -- which is what we + // want since R_a is the "active" leg that emits). + // R_a is updated as R_a -= peeled. R_b stays constant. + // + // For the 2->3 block, the kernel internally computes + // p_12 = pa + pb - p_3, where p_3 = im1 (previous peel on this side). + // We want p_12 = remaining-system-left-to-produce = R_a + R_b + // (using R_a *after* the previous peel was subtracted from it). + // Since the kernel subtracts im1 internally, and im1 is already + // subtracted from R_a in our bookkeeping, we restore it: pass + // pb = R_a + im1, so p_12 = R_b + (R_a + im1) - im1 = R_b + R_a. ✓ + if (first) { + auto ks = _lab_scattering.build_forward( + fb, + {next_random(), next_random(), m_rest, m_peel}, + {R_b, R_a} + ); + Value peeled = ks.at(1); + p_out[s[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); + auto ks = _two_to_three.build_forward( + fb, + {next_random(), next_random(), next_random(), m_rest, m_peel}, + {R_b, pb_for_block, im1} + ); + Value peeled = ks.at(1); + p_out[s[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); + }; + + walk(_set1, P_set1, R_b_for_set1, rest_masses_set1); + walk(_set2, P_set2, R_b_for_set2, rest_masses_set2); + + // Assemble outputs: momentum0, momentum1 = beams; momentum_{2+i} = outgoing i. + ValueVec p_ext; + p_ext.reserve(_n_out + 2); + 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]); + } + + return {{output_types().keys(), p_ext}, fb.product(dets)}; +} + +Mapping::Result TPropagatorMapping23::build_inverse_impl( + FunctionBuilder& fb, + const NamedVector& inputs, + const NamedVector& conditions +) const { + Value e_cm = conditions.at(0); + ValueVec m_out(conditions.begin() + 1, conditions.end()); + ValueVec random_out; + ValueVec dets; + + // Inputs: pa, pb, p_out[0], p_out[1], ..., p_out[n_out-1] + Value pa = inputs.at(0); + Value pb = inputs.at(1); + auto p_outgoing = [&](std::size_t i) { return inputs.at(2 + i); }; + + // ============================================================ + // Build a factor matrix to extract all needed invariant masses in + // one shot. Rows (in order): + // row 0: m^2(P_set1) (if set1.size() >= 2) + // row 1: m^2(P_set2) (if set2.size() >= 2) + // then for set1 if size >= 3: + // k_1 - 2 rows, each summing momenta s1[j+1..k_1-1] for j=0..k_1-3 + // then for set2 if size >= 3: similarly + // Column = index into inputs (n_out + 2 columns). + // We record the layout so we can pull invariants out by name afterwards. + // ============================================================ + nested_vector2 invariant_factors; + auto n_inputs = _n_out + 2; + auto factor_set = [&](const std::vector& idxs) { + std::vector row(n_inputs, 0.0); + for (auto idx : idxs) row.at(idx + 2) = 1.0; + return row; + }; + // Track positions in the stacked invariants output: + int idx_m2_set1 = -1, idx_m2_set2 = -1; + int idx_rest_set1_start = -1, idx_rest_set2_start = -1; + if (_set1.size() >= 2) { + idx_m2_set1 = invariant_factors.size(); + invariant_factors.push_back(factor_set(_set1)); + } + if (_set2.size() >= 2) { + idx_m2_set2 = invariant_factors.size(); + invariant_factors.push_back(factor_set(_set2)); + } + if (_set1.size() >= 3) { + idx_rest_set1_start = invariant_factors.size(); + for (std::size_t j = 0; j < _set1.size() - 2; ++j) { + std::vector sub(_set1.begin() + j + 1, _set1.end()); + invariant_factors.push_back(factor_set(sub)); + } + } + if (_set2.size() >= 3) { + idx_rest_set2_start = invariant_factors.size(); + for (std::size_t j = 0; j < _set2.size() - 2; ++j) { + std::vector sub(_set2.begin() + j + 1, _set2.end()); + invariant_factors.push_back(factor_set(sub)); + } + } + ValueVec invariants; + if (!invariant_factors.empty()) { + invariants = fb.unstack( + fb.invariants_from_momenta(fb.stack(inputs.values()), invariant_factors) + ); + } + + // ============================================================ + // Recover composites P_set1, P_set2 from the momenta. + // ============================================================ + auto sum_momenta = [&](const std::vector& s) -> Value { + Value p = p_outgoing(s[0]); + for (std::size_t k = 1; k < s.size(); ++k) { + p = fb.add(p, p_outgoing(s[k])); + } + return p; + }; + Value P_set1 = sum_momenta(_set1); + Value P_set2 = sum_momenta(_set2); + + // ============================================================ + // Phase 1a inverse: recover set-mass randoms (skipped in DoubleT case). + // Even when randoms are skipped, we still need m_set1/m_set2 for the + // Phase 1b intermediate-mass-bound calculation; those come from the + // invariants_from_momenta result for multi-particle sides. + // ============================================================ + Value m_set1, m_set2; + Value mass_sum_set1 = sum_of_masses(fb, m_out, _set1); + Value mass_sum_set2 = sum_of_masses(fb, m_out, _set2); + if (_use_double_t) { + // No set-mass randoms to recover; just set m_set1/m_set2. + if (_set1.size() == 1) { + m_set1 = m_out.at(_set1[0]); + m_set2 = fb.sqrt(invariants.at(idx_m2_set2)); + } else { + m_set2 = m_out.at(_set2[0]); + m_set1 = fb.sqrt(invariants.at(idx_m2_set1)); + } + } else { + if (_set1.size() >= 2) { + auto s_min = fb.square(mass_sum_set1); + 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}); + random_out.push_back(res["random"]); + dets.push_back(res["det"]); + m_set1 = fb.sqrt(m2_set1); + } else { + m_set1 = m_out.at(_set1[0]); + } + if (_set2.size() >= 2) { + auto s_min = fb.square(mass_sum_set2); + 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}); + random_out.push_back(res["random"]); + dets.push_back(res["det"]); + m_set2 = fb.sqrt(m2_set2); + } else { + m_set2 = m_out.at(_set2[0]); + } + } + + // ============================================================ + // Phase 2 inverse: central block from (pa, pb). Comes BEFORE Phase 1b + // because the forward emits central randoms before intermediate-mass + // randoms (so the multi-side m_set is known by Phase 1b). + // ============================================================ + if (_use_double_t) { + bool single_is_set1 = (_set1.size() == 1); + Value p_single = single_is_set1 ? P_set1 : P_set2; + 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 ? mass_sum_set2 : mass_sum_set1; + auto central = _double_t.build_inverse( + fb, + {p_single, p_recoil}, + {pa, pb, m_single, mir_min} + ); + // DoubleT inverse returns (r_phi, r_t1, r_t2, det). m1 is no longer + // an output (it's a condition now). + 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} + ); + random_out.push_back(central.at(0)); + random_out.push_back(central.at(1)); + dets.push_back(central["det"]); + } + + // ============================================================ + // Phase 1b inverse: recover intermediate rest-mass randoms + // ============================================================ + auto recover_intermediate_masses = + [&](const std::vector& s, Value m_set, int idx_start) { + std::size_t k = s.size(); + if (k <= 2) return; + // Mirror the forward chaining: prev_mass starts at m_set, then becomes + // 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]); + for (std::size_t i = j + 2; i < k; ++i) { + m_min = fb.add(m_min, m_out.at(s[i])); + } + auto s_min = fb.square(m_min); + auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s[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"]); + dets.push_back(res["det"]); + prev_mass = fb.sqrt(m2); + } + }; + + recover_intermediate_masses(_set1, m_set1, idx_rest_set1_start); + recover_intermediate_masses(_set2, m_set2, idx_rest_set2_start); + + 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]); + // The block at this rung had inputs (R_a, R_b) and produced + // (p1_out, p2_out) where p1_out has mass m_rest (carrier of + // remaining) and p2_out = peeled. 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"]); + } + R_a = fb.sub(R_a, peeled); + im1 = peeled; + } + }; + + walk_inverse(_set1, P_set1, R_b_for_set1); + walk_inverse(_set2, P_set2, R_b_for_set2); + + return {{input_types().keys(), random_out}, fb.product(dets)}; +} diff --git a/madspace/src/phasespace/two_particle.cpp b/madspace/src/phasespace/two_particle.cpp index 31c547f50..fd85ae860 100644 --- a/madspace/src/phasespace/two_particle.cpp +++ b/madspace/src/phasespace/two_particle.cpp @@ -134,11 +134,11 @@ DoubleT::DoubleT( "DoubleT", {{"random_phi", batch_float}, {"random_t1", batch_float}, - {"random_t2", batch_float}, - {"mass1", 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}} ), _t1_invariant(t1_invariant_power, t1_mass, t1_width), @@ -150,9 +150,9 @@ Mapping::Result DoubleT::build_forward_impl( const NamedVector& conditions ) const { auto r_phi = inputs.at(0), r_t1 = inputs.at(1), r_t2 = inputs.at(2); - auto m1 = inputs.at(3); auto p_in1 = conditions.at(0), p_in2 = conditions.at(1); - auto mir_min = conditions.at(2); + auto m1 = conditions.at(2); + auto mir_min = conditions.at(3); auto [t1_min, t1_max] = fb.t1_inv_min_max_doublet(p_in1, p_in2, m1, mir_min); auto t1_result = _t1_invariant.build_forward(fb, {r_t1}, {t1_min, t1_max}); @@ -180,17 +180,18 @@ Mapping::Result DoubleT::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 mir_min = conditions.at(2); + auto m1 = conditions.at(2); + auto mir_min = conditions.at(3); - auto [r_phi, m1, det_scatter] = + 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, mir_min); + fb.t1_inv_value_and_min_max_doublet(p_in1, p_in2, p1, m1, mir_min); 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, mir_min, t1_abs + p_in1, p_in2, p1, m1, mir_min, t1_abs ); auto t2_result = _t2_invariant.build_inverse(fb, {t2_abs}, {t2_min, t2_max}); @@ -198,8 +199,7 @@ Mapping::Result DoubleT::build_inverse_impl( return { {{"random_phi", r_phi}, {"random_t1", t1_result["random"]}, - {"random_t2", t2_result["random"]}, - {"mass1", m1}}, + {"random_t2", t2_result["random"]}}, fb.mul(det_inv, det_scatter) }; } diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index 3b10ec7bf..10dc81741 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -83,9 +83,9 @@ void add_instructions(py::classh& fb) { 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("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("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("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("invariants_from_momenta", &FunctionBuilder::invariants_from_momenta, py::arg("p_ext"), py::arg("factors")); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 9d865bf77..a58736f48 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -485,7 +485,17 @@ PYBIND11_MODULE(_madspace_py, m) { py::init, double>(), py::arg("integration_order"), py::arg("invariant_power") = 0. - ); + ) + .def("random_dim", &TPropagatorMapping::random_dim); + + py::classh(m, "TPropagatorMapping23") + .def( + py::init, double, double>(), + py::arg("colour_order"), + py::arg("t_invariant_power") = 0.8, + py::arg("s_invariant_power") = 0.8 + ) + .def("random_dim", &TPropagatorMapping23::random_dim); py::classh(m, "VegasHistogram") .def( diff --git a/madspace/tests/test_2to3_scattering.py b/madspace/tests/test_2to3_scattering.py index 147d5011f..ba96e8909 100644 --- a/madspace/tests/test_2to3_scattering.py +++ b/madspace/tests/test_2to3_scattering.py @@ -41,7 +41,7 @@ def boost(k: np.ndarray, p_boost: np.ndarray, inverse: bool = False) -> np.ndarr # ---------------------------- # N = 50_000 # keep this moderate for CI speed; bump locally for tighter stats -N = 10_000 # keep this moderate for CI speed; bump locally for tighter stats +N = 100_000 # keep this moderate for CI speed; bump locally for tighter stats InputPoint = namedtuple( "InputPoint", diff --git a/madspace/tests/test_double_t.py b/madspace/tests/test_double_t.py index 30ad266ff..b6ca9ce38 100644 --- a/madspace/tests/test_double_t.py +++ b/madspace/tests/test_double_t.py @@ -72,7 +72,7 @@ def input_points(rng, request): # m1: the fixed single-particle mass # mir_min: lower bound on recoil mass (e.g. sum of constituent masses). - # Pick mir_min well below the sqrt(s) so there's room for the recoil. + # Pick mir_min well below the sqrt(s) so there\'s room for the recoil. m1 = rng.uniform(1.0, 50.0, N) mir_min = rng.uniform(50.0, 200.0, N) @@ -124,9 +124,13 @@ def test_momentum_conservation(input_points): input_points.r_phi, input_points.r_t1, input_points.r_t2, + ] + conditions = [ + input_points.pa, + input_points.pb, input_points.m1, + input_points.mir_min, ] - conditions = [input_points.pa, input_points.pb, input_points.mir_min] p1, p2, det = mapping.map_forward(inputs, conditions) assert p1 + p2 == approx(input_points.pa + input_points.pb) @@ -139,9 +143,13 @@ def test_on_shell_mass(input_points): input_points.r_phi, input_points.r_t1, input_points.r_t2, + ] + conditions = [ + input_points.pa, + input_points.pb, input_points.m1, + input_points.mir_min, ] - conditions = [input_points.pa, input_points.pb, input_points.mir_min] p1, p2, det = mapping.map_forward(inputs, conditions) assert mass(p1) == approx(input_points.m1) @@ -156,9 +164,13 @@ def test_recoil_mass_minimum(input_points): input_points.r_phi, input_points.r_t1, input_points.r_t2, + ] + conditions = [ + input_points.pa, + input_points.pb, input_points.m1, + input_points.mir_min, ] - conditions = [input_points.pa, input_points.pb, input_points.mir_min] p1, p2, det = mapping.map_forward(inputs, conditions) m_ir = mass(p2) @@ -173,9 +185,13 @@ def test_inverse(input_points): input_points.r_phi, input_points.r_t1, input_points.r_t2, + ] + conditions = [ + input_points.pa, + input_points.pb, input_points.m1, + input_points.mir_min, ] - conditions = [input_points.pa, input_points.pb, input_points.mir_min] p1, p2, det = mapping.map_forward(inputs, conditions) *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) @@ -186,40 +202,64 @@ def test_inverse(input_points): assert inp == approx(inv_inp), f"mismatch in input index {i}" -def test_phase_space_volume(fixed_input_points): - """For massless beams and the double-t parameterization, the phase-space - volume over all (t1, t2, phi) at fixed (s, m1, mir_min) should match the - analytic 2-body phase-space volume integrated against the constraint - that m_ir >= mir_min. +def test_phase_space_volume_analytic(fixed_input_points): + """The integral of DoubleT\'s det over (r_phi, r_t1, r_t2) should equal + the 2-body phase-space volume integrated over m_ir^2 in [mir_min^2, + (sqrt(s)-m1)^2]: - Simpler check: the integral of the det over the (r_phi, r_t1, r_t2) - cube should be independent of how we parameterize, i.e. equal the - same fixed kinematic volume. We compare it to the standard 2->2 volume - pi/(2s) * sqrt(lambda(s, m1^2, mir_min^2)) ... but the recoil mass is - NOT fixed in double_t, so this comparison is approximate. + V = int_{mir_min^2}^{(sqrt(s)-m1)^2} dm_ir^2 * (pi/(2s)) * sqrt(lambda(s, m1^2, mir^2)) - Instead: just check the determinant has a finite, frame-independent mean - with the expected scaling. The mean should be positive and finite.""" - mapping = ms.DoubleT() + This matches the 2->2 volume convention (no (2pi)^n factors). + """ + from scipy.integrate import quad + mapping = ms.DoubleT() inputs = [ fixed_input_points.r_phi, fixed_input_points.r_t1, fixed_input_points.r_t2, - fixed_input_points.m1, ] conditions = [ fixed_input_points.pa, fixed_input_points.pb, + fixed_input_points.m1, fixed_input_points.mir_min, ] p1, p2, det = mapping.map_forward(inputs, conditions) - assert np.all(np.isfinite(det)) - assert np.all(det > 0) - # variance should be finite and well below the mean (sanity check on numerics) - std_error = np.std(det) / np.sqrt(N) - assert std_error < np.mean(det) + # All entries of fixed_input_points have the same s, m1, mir_min by construction. + pa = fixed_input_points.pa[0] + pb = fixed_input_points.pb[0] + s = (pa[0] + pb[0]) ** 2 - np.sum((pa[1:] + pb[1:]) ** 2) + m1_val = fixed_input_points.m1[0] + mir_min_val = fixed_input_points.mir_min[0] + + def lam(s, m1_sq, m2_sq): + return (s - m1_sq - m2_sq) ** 2 - 4 * m1_sq * m2_sq + + def integrand(mir_sq): + l = lam(s, m1_val ** 2, mir_sq) + if l <= 0: + return 0.0 + return np.pi / (2 * s) * np.sqrt(l) + + analytic_volume, _ = quad( + integrand, + mir_min_val ** 2, + (np.sqrt(s) - m1_val) ** 2, + ) + + mc_volume = np.mean(det) + mc_err = np.std(det) / np.sqrt(N) + + # 5-sigma agreement + diff = abs(mc_volume - analytic_volume) + assert diff < 5 * mc_err, ( + f"DoubleT volume mismatch: MC = {mc_volume:.6e} +/- {mc_err:.6e}, " + f"analytic = {analytic_volume:.6e}, " + f"diff = {diff:.4e} ({diff/mc_err:.2f} sigma), " + f"ratio MC/analytic = {mc_volume/analytic_volume:.4f}" + ) def test_inverse_lab_specifically(rng): @@ -245,13 +285,127 @@ def test_inverse_lab_specifically(rng): r_t1 = rng.random(N) r_t2 = rng.random(N) - inputs = [r_phi, r_t1, r_t2, m1] - conditions = [pa, pb, mir_min] + inputs = [r_phi, r_t1, r_t2] + conditions = [pa, pb, m1, mir_min] p1, p2, det = mapping.map_forward(inputs, conditions) *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) - names = ["r_phi", "r_t1", "r_t2", "m1"] + names = ["r_phi", "r_t1", "r_t2"] for name, inp, inv_inp in zip(names, inputs, inv_inputs): assert inp == approx(inv_inp), f"LAB mismatch in {name}" - assert inv_det == approx(1 / det, rel=1e-5) \ No newline at end of file + assert inv_det == approx(1 / det, rel=1e-5) + + +# ---------------------------- +# Massless m1 stress tests (regression for the t-bound precision issue) +# ---------------------------- + + +@pytest.fixture( + params=[ + ("m1=0", 0.0), + ("m1=1e-9", 1e-9), + ("m1=1e-6", 1e-6), + ("m1=1e-3", 1e-3), + ], + ids=["m1=0", "m1=1e-9", "m1=1e-6", "m1=1e-3"], +) +def massless_input_points(rng, request): + """Near-massless m1 stresses the inverse: lsquare(p1) suffers + catastrophic cancellation when m1 is recovered from it. With m1 as a + condition (Option B fix), this should round-trip cleanly.""" + _, m1_val = request.param + + pz = np.full(N, 6500.0) # 13 TeV + pa = np.stack([pz, np.zeros(N), np.zeros(N), +pz], axis=1) + pb = np.stack([pz, np.zeros(N), np.zeros(N), -pz], axis=1) + + m1 = np.full(N, m1_val) + mir_min = np.full(N, 50.0) + + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + return InputPoint(r_phi, r_t1, r_t2, m1, mir_min, pa, pb) + + +def test_inverse_massless(massless_input_points): + """The DoubleT inverse must round-trip even for m1 ~= 0.""" + mapping = ms.DoubleT() + + inputs = [ + massless_input_points.r_phi, + massless_input_points.r_t1, + massless_input_points.r_t2, + ] + conditions = [ + massless_input_points.pa, + massless_input_points.pb, + massless_input_points.m1, + massless_input_points.mir_min, + ] + + p1, p2, det = mapping.map_forward(inputs, conditions) + *inv_inputs, inv_det = mapping.map_inverse([p1, p2], conditions) + + # det round-trip + assert inv_det == approx(1 / det, rel=1e-5) + + # random round-trip; r_t2 used to be the most sensitive (its bounds + # depend on t1_abs which depended on the lossy m1^2 recompute). + for i, (inp, inv_inp) in enumerate(zip(inputs, inv_inputs)): + assert inp == approx(inv_inp, rel=1e-5, abs=1e-8), \ + f"mismatch in input index {i}" + + +def test_phase_space_volume_massless(rng): + """Volume integral should still match analytic for massless m1.""" + from scipy.integrate import quad + + m1_val = 0.0 + mir_min_val = 80.0 + s_val = 13000.0 ** 2 + + pa = np.tile([6500.0, 0, 0, 6500.0], (N, 1)) + pb = np.tile([6500.0, 0, 0, -6500.0], (N, 1)) + + m1 = np.full(N, m1_val) + mir_min = np.full(N, mir_min_val) + + r_phi = rng.random(N) + r_t1 = rng.random(N) + r_t2 = rng.random(N) + + mapping = ms.DoubleT() + p1, p2, det = mapping.map_forward( + [r_phi, r_t1, r_t2], + [pa, pb, m1, mir_min], + ) + + def lam(s, m1_sq, m2_sq): + return (s - m1_sq - m2_sq) ** 2 - 4 * m1_sq * m2_sq + + def integrand(mir_sq): + l = lam(s_val, m1_val ** 2, mir_sq) + if l <= 0: + return 0.0 + return np.pi / (2 * s_val) * np.sqrt(l) + + analytic_volume, _ = quad( + integrand, + mir_min_val ** 2, + (np.sqrt(s_val) - m1_val) ** 2, + ) + + mc_volume = np.mean(det) + mc_err = np.std(det) / np.sqrt(N) + + diff = abs(mc_volume - analytic_volume) + assert diff < 5 * mc_err, ( + f"DoubleT massless volume mismatch: MC = {mc_volume:.6e} +/- " + f"{mc_err:.6e}, analytic = {analytic_volume:.6e}, " + f"diff = {diff:.4e} ({diff/mc_err:.2f} sigma), " + f"ratio MC/analytic = {mc_volume/analytic_volume:.4f}" + ) \ No newline at end of file diff --git a/madspace/tests/test_t_propagator_23.py b/madspace/tests/test_t_propagator_23.py new file mode 100644 index 000000000..07d1f1934 --- /dev/null +++ b/madspace/tests/test_t_propagator_23.py @@ -0,0 +1,264 @@ +"""Tests for TPropagatorMapping23 (per-set, colour-ordered topology). + +Colour order convention: 0-indexed permutation of {0, ..., n-1}. +Particles 0 and 1 are the two incoming beams; 2, 3, ..., n-1 are outgoing. + +Note on non-physical events: For inner-rung 2->3 peels, the kinematic +constraint imposed by the 2->3 block (Gram-determinant condition, s23/t1 +range non-empty) can be violated for some sampled (rest-mass, t, s) combinations. +These events produce non-finite momenta/det and should be discarded in any +Monte Carlo integration -- they correspond to zero-weight points outside the +physical region. The tests therefore filter on `np.isfinite(det)` before +checking conservation, on-shell, and inverse round-trip. + +Note on numerical tails: events near kinematic boundaries (very small forward +det) have accumulated floating-point error from chained sqrts, boosts, and +rotations. A 10th-percentile det floor filters these without biasing the +integrand (they have near-zero weight in any MC integral anyway). +""" +import numpy as np +import pytest +from pytest import approx + +import madspace as ms + + +@pytest.fixture +def rng(): + return np.random.default_rng(1234) + + +N = 100_000 +CM_ENERGY = 13000.0 +MIN_PHYSICAL_FRACTION = 0.10 # at least 10% of events should be physical +LOW_DET_QUANTILE = 0.10 # drop bottom 10% by |det| in stability-sensitive tests + + +# Mass set families. The labels are templates; 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 defined for n_out >= 3. + return [173.0, 80.0] + [0.0] * (n_out - 2), f"{n_out} mixed" + + +def _mass_sets_for(n_out): + """Mass families that make sense at this n_out.""" + sets = [_massless(n_out), _w_like(n_out)] + if n_out >= 3: + sets.append(_mixed_top(n_out)) + return sets + + +# Each entry: (colour_order, int_order_for_old_class, topology_label). +# The topology determines n_out = len(colour_order) - 2. +TOPOLOGIES = [ + ([0, 2, 3, 1, 4], [0, 1], "n=5: set1={2,3}, set2={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, 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}"), +] + + +def _expand(topologies): + """Produce a flat list of (colour_order, int_order, masses, label) tuples + where each topology is paired only with its compatible mass families. + This replaces the cross-product parametrize that produced many skips.""" + out = [] + for colour_order, int_order, topo_label in topologies: + n_out = len(colour_order) - 2 + for masses, mass_label in _mass_sets_for(n_out): + label = f"{mass_label}-{topo_label}" + out.append((colour_order, int_order, masses, label)) + return out + + +CASES = _expand(TOPOLOGIES) +CASE_IDS = [c[3] for c in CASES] + + +def _physical_mask(det, p_ext): + """Events are physical iff det is finite AND every momentum is finite.""" + 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): + mask = _physical_mask(det, p_ext) + if mask.sum() < MIN_PHYSICAL_FRACTION * len(det): + pytest.fail( + f"Too few physical events: {mask.sum()}/{len(det)} " + f"({100*mask.sum()/len(det):.1f}%), expected >= " + f"{100*MIN_PHYSICAL_FRACTION:.0f}%" + ) + return mask, [p[mask] for p in p_ext], det[mask] + + +def _drop_low_det_tail(det, *arrays, q=LOW_DET_QUANTILE, median_relative=1e-6): + """Drop the bottom-q quantile AND any events below median_relative * median + of events by |det|. These boundary events accumulate floating-point error + from chained operations and have near-zero weight in MC anyway. The + median-relative floor matters for heavy-tailed DoubleT topologies where + the quantile floor alone is still many decades below the median.""" + median = np.median(np.abs(det)) + floor = max(np.quantile(np.abs(det), q), median_relative * median) + keep = np.abs(det) > floor + return (det[keep], *(a[keep] for a in arrays)) + + +@pytest.mark.parametrize("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_momentum_conservation(rng, colour_order, _int_order, masses, _label): + mapping = ms.TPropagatorMapping23(colour_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("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_on_shell_masses(rng, colour_order, _int_order, masses, _label): + mapping = ms.TPropagatorMapping23(colour_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) + + # Drop boundary events that accumulate FP error from chained kinematics. + det, *p_ext = _drop_low_det_tail(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("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_inverse(rng, colour_order, _int_order, masses, _label): + mapping = ms.TPropagatorMapping23(colour_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) + 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) + + # Drop the lowest-det events: boundary events whose forward det is many + # orders of magnitude below the median have inverse det near infinity, + # and the round-trip |det * inv_det - 1| can hit a few percent (or more) + # from accumulated FP error in atan2/sqrt/boosts. They have negligible + # weight in any MC integral. We use both a quantile floor (handles + # well-behaved distributions) and a median-relative floor (handles the + # heavy-tailed DoubleT-using topologies where the 10th percentile is + # still many decades below the median). + median = np.median(np.abs(det_phys)) + floor_quantile = np.quantile(np.abs(det_phys), LOW_DET_QUANTILE) + floor_median = 1e-6 * median + floor = max(floor_quantile, floor_median) + keep = ( + np.isfinite(det_inv) + & np.isfinite(det_phys) + & (np.abs(det_phys) > floor) + ) + for ri in r_inv: + keep &= np.isfinite(ri) + if keep.sum() < MIN_PHYSICAL_FRACTION * len(det): + pytest.fail( + f"Too few invertible events: {keep.sum()}/{len(det)} " + f"({100*keep.sum()/len(det):.1f}%)" + ) + + det_kept = det_phys[keep] + det_inv_kept = det_inv[keep] + # Inverse-roundtrip tolerance: 0.1% accounts for accumulated FP error + # in chained kinematic operations (boosts, rotations, sqrts) for the + # deeper-walk topologies. Correctness is established by + # 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); + # round-trip only recovers which-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("colour_order,int_order,masses,_label", CASES, ids=CASE_IDS) +def test_phase_space_volume_matches_old(rng, colour_order, int_order, masses, _label): + """The new class should integrate to the same phase-space volume as + the old TPropagatorMapping (with matrix element = 1). + + This is the actual correctness check: both classes are MC estimators of + the same n-body phase-space integral, and their averages must agree + within a few sigma. (The inverse-roundtrip test is a precision check; + this one is the physics check.) + """ + new = ms.TPropagatorMapping23(colour_order) + old = ms.TPropagatorMapping(int_order) + + r_new = rng.random((N, new.random_dim())) + r_old = rng.random((N, old.random_dim())) + e_cm = np.full(N, CM_ENERGY) + cond_new = [e_cm] + [np.full(N, m) for m in masses] + cond_old = cond_new + + *p_new, det_new = new.map_forward(list(r_new.T), cond_new) + *p_old, det_old = old.map_forward(list(r_old.T), cond_old) + + mask_new = np.isfinite(det_new) & (det_new > 0) + mask_old = np.isfinite(det_old) & (det_old > 0) + det_new = det_new[mask_new] + det_old = det_old[mask_old] + + assert len(det_new) > MIN_PHYSICAL_FRACTION * N, \ + f"new class: {len(det_new)}/{N} physical, too few" + assert len(det_old) > MIN_PHYSICAL_FRACTION * N, \ + f"old class: {len(det_old)}/{N} physical, too few" + + mean_new = np.mean(det_new) + mean_old = 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)" + ) \ No newline at end of file From 86abb2df0922a8fbd2558be923ecf7abba02df18 Mon Sep 17 00:00:00 2001 From: Ramon Winterhalder Date: Tue, 19 May 2026 15:05:22 +0200 Subject: [PATCH 4/6] Fix the color ordered mapping and rename properly --- madspace/CMakeLists.txt | 4 +- madspace/generate_pyi.sh | 3 +- .../compgraphs/function_builder_mixin.inc | 8 +- madspace/include/madspace/phasespace.hpp | 2 +- ...pping_23.hpp => color_ordered_mapping.hpp} | 4 +- madspace/instruction_set.yaml | 4 +- madspace/madspace/_madspace_py.pyi | 786 +++++++++++++----- madspace/src/kernels/kernels.hpp | 1 + madspace/src/kernels/kinematics.hpp | 6 +- madspace/src/kernels/lup_det.hpp | 371 +++++++++ madspace/src/kernels/threeparticle.hpp | 89 +- ...pping_23.cpp => color_ordered_mapping.cpp} | 200 ++--- madspace/src/python/instruction_set.hpp | 4 +- madspace/src/python/madspace.cpp | 158 ++-- madspace/tests/test_2to3_scattering.py | 3 +- madspace/tests/test_color_ordered_mapping.py | 209 +++++ madspace/tests/test_t_propagator_23.py | 264 ------ 17 files changed, 1354 insertions(+), 762 deletions(-) rename madspace/include/madspace/phasespace/{t_propagator_mapping_23.hpp => color_ordered_mapping.hpp} (95%) create mode 100644 madspace/src/kernels/lup_det.hpp rename madspace/src/phasespace/{t_propagator_mapping_23.cpp => color_ordered_mapping.cpp} (67%) create mode 100644 madspace/tests/test_color_ordered_mapping.py delete mode 100644 madspace/tests/test_t_propagator_23.py diff --git a/madspace/CMakeLists.txt b/madspace/CMakeLists.txt index 029b0bcdd..b16193713 100644 --- a/madspace/CMakeLists.txt +++ b/madspace/CMakeLists.txt @@ -137,7 +137,7 @@ add_library( src/phasespace/luminosity.cpp src/phasespace/topology.cpp src/phasespace/t_propagator_mapping.cpp - src/phasespace/t_propagator_mapping_23.cpp + src/phasespace/color_ordered_mapping.cpp src/phasespace/phasespace.cpp src/phasespace/multichannel.cpp src/phasespace/madnis.cpp @@ -190,7 +190,7 @@ add_library( include/madspace/phasespace/luminosity.hpp include/madspace/phasespace/topology.hpp include/madspace/phasespace/t_propagator_mapping.hpp - include/madspace/phasespace/t_propagator_mapping_23.hpp + include/madspace/phasespace/color_ordered_mapping.hpp include/madspace/phasespace/phasespace.hpp include/madspace/phasespace/multichannel.hpp include/madspace/phasespace/madnis.hpp diff --git a/madspace/generate_pyi.sh b/madspace/generate_pyi.sh index 82bc3082c..a1b009f8d 100755 --- a/madspace/generate_pyi.sh +++ b/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 diff --git a/madspace/include/madspace/compgraphs/function_builder_mixin.inc b/madspace/include/madspace/compgraphs/function_builder_mixin.inc index dea526d9e..b1050b404 100644 --- a/madspace/include/madspace/compgraphs/function_builder_mixin.inc +++ b/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) { diff --git a/madspace/include/madspace/phasespace.hpp b/madspace/include/madspace/phasespace.hpp index 4ae860a54..e2f1ade96 100644 --- a/madspace/include/madspace/phasespace.hpp +++ b/madspace/include/madspace/phasespace.hpp @@ -22,7 +22,7 @@ #include "phasespace/rambo.hpp" #include "phasespace/scale.hpp" #include "phasespace/t_propagator_mapping.hpp" -#include "madspace/phasespace/t_propagator_mapping_23.hpp" +#include "phasespace/color_ordered_mapping.hpp" #include "phasespace/three_particle.hpp" #include "phasespace/topology.hpp" #include "phasespace/two_particle.hpp" diff --git a/madspace/include/madspace/phasespace/t_propagator_mapping_23.hpp b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp similarity index 95% rename from madspace/include/madspace/phasespace/t_propagator_mapping_23.hpp rename to madspace/include/madspace/phasespace/color_ordered_mapping.hpp index 5b3d37cdb..dedc5d6e6 100644 --- a/madspace/include/madspace/phasespace/t_propagator_mapping_23.hpp +++ b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp @@ -10,11 +10,11 @@ namespace madspace { -class TPropagatorMapping23 : public Mapping { +class ColorOrderedMapping : public Mapping { public: // colour_order: 0-indexed permutation of {0, ..., n-1} (n = n_out + 2). // Particles 0 and 1 are the two incoming beams. - TPropagatorMapping23( + ColorOrderedMapping( const std::vector& colour_order, double t_invariant_power = 0.8, double s_invariant_power = 0.8 diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 453e24c97..116d298f2 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -255,7 +255,7 @@ reduce_sum_vector: batch_reduce_mean: inputs: - - name: in + - name: in1 type: [float] desc: outputs: @@ -267,7 +267,7 @@ batch_reduce_mean: batch_reduce_mean_keepdim: inputs: - - name: in + - name: in1 type: [float] desc: outputs: diff --git a/madspace/madspace/_madspace_py.pyi b/madspace/madspace/_madspace_py.pyi index 963c150df..8b0251fe5 100644 --- a/madspace/madspace/_madspace_py.pyi +++ b/madspace/madspace/_madspace_py.pyi @@ -1,12 +1,14 @@ from __future__ import annotations - import collections.abc import typing -__all__ = [ +__all__: list[str] = [ + "AdamOptimizer", "AlphaSGrid", "BatchSize", + "ChannelEventGenerator", "ChannelWeightNetwork", + "ColorOrderedMapping", "Context", "CutItem", "Cuts", @@ -19,19 +21,19 @@ __all__ = [ "DiscreteHistogram", "DiscreteOptimizer", "DiscreteSampler", + "DoubleT", "EnergyScale", "EventGenerator", - "EventGeneratorConfig", - "EventGeneratorHistogram", - "EventGeneratorStatus", - "EventGeneratorVerbosity", "FastRamboMapping", "Flow", "Function", "FunctionBuilder", "FunctionGenerator", "FunctionRuntime", + "GeneratorConfig", + "GeneratorStatus", "HistItem", + "Histogram", "Instruction", "InstructionCall", "Integrand", @@ -48,6 +50,9 @@ __all__ = [ "Logger", "Luminosity", "MLP", + "MadnisConfig", + "MadnisLoss", + "MadnisTraining", "Mapping", "MatrixElement", "MatrixElementApi", @@ -55,6 +60,9 @@ __all__ = [ "MultiChannelFunction", "MultiChannelIntegrand", "MultiChannelMapping", + "MultiMadnisTraining", + "NamedTypes", + "NamedValues", "Observable", "ObservableHistograms", "PartonDensity", @@ -79,6 +87,7 @@ __all__ = [ "VegasGridOptimizer", "VegasHistogram", "VegasMapping", + "Verbosity", "batch_float", "batch_float_array", "batch_four_vec", @@ -103,12 +112,64 @@ __all__ = [ "pretty", "set_lib_path", "set_simd_vector_size", - "set_thread_count", "silent", "single_float", "single_int", ] +class AdamOptimizer: + class LRSchedule: + """ + Members: + + none + + cosine + """ + + __members__: typing.ClassVar[ + dict[str, AdamOptimizer.LRSchedule] + ] # value = {'none': , 'cosine': } + cosine: typing.ClassVar[ + AdamOptimizer.LRSchedule + ] # value = + none: typing.ClassVar[AdamOptimizer.LRSchedule] # value = + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + @typing.overload + def __init__(self, value: typing.SupportsInt) -> None: ... + @typing.overload + def __init__(self, name: str) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: typing.SupportsInt) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... + + cosine: typing.ClassVar[AdamOptimizer.LRSchedule] # value = + none: typing.ClassVar[AdamOptimizer.LRSchedule] # value = + def __init__( + self, + function: Function, + context: Context, + learning_rate: typing.SupportsFloat, + schedule: AdamOptimizer.LRSchedule = AdamOptimizer.LRSchedule.LRSchedule.none, + step_count: typing.SupportsInt = 0, + beta1: typing.SupportsFloat = 0.9, + beta2: typing.SupportsFloat = 0.999, + eps: typing.SupportsFloat = 1e-08, + ) -> None: ... + def context(self) -> Context: ... + def input_types(self) -> list[Type]: ... + def learning_rate(self) -> float: ... + def step(self, inputs: collections.abc.Sequence[typing.Any]) -> list[Tensor]: ... + class AlphaSGrid: def __init__(self, file: str) -> None: ... def coefficients_shape(self, batch_dim: bool = False) -> list[int]: ... @@ -134,6 +195,30 @@ class BatchSize: def __repr__(self) -> str: ... def __str__(self) -> str: ... +class ChannelEventGenerator: + integrand_flags: typing.ClassVar[int] = 1077 + @staticmethod + def load( + channel_file: str, + contexts: collections.abc.Sequence[Context], + event_file: str, + weight_file: str, + config: GeneratorConfig, + ) -> ChannelEventGenerator: ... + def __init__( + self, + contexts: collections.abc.Sequence[Context], + integrand: Integrand, + event_file: str, + weight_file: str, + config: GeneratorConfig, + subprocess_index: typing.SupportsInt, + name: str, + histograms: madspace._madspace_py.ObservableHistograms | None, + ) -> None: ... + def save(self, save: str) -> None: ... + def status(self) -> GeneratorStatus: ... + class ChannelWeightNetwork(FunctionGenerator): def __init__( self, @@ -143,17 +228,30 @@ class ChannelWeightNetwork(FunctionGenerator): layers: typing.SupportsInt = 3, activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, prefix: str = "", + include_preprocessing: bool = True, ) -> None: ... def initialize_globals(self, context: Context) -> None: ... def mask_name(self) -> str: ... def mlp(self) -> MLP: ... def preprocessing(self) -> MomentumPreprocessing: ... +class ColorOrderedMapping(Mapping): + def __init__( + self, + colour_order: collections.abc.Sequence[typing.SupportsInt], + t_invariant_power: typing.SupportsFloat = 0.0, + s_invariant_power: typing.SupportsFloat = 0.0, + ) -> None: ... + def random_dim(self) -> int: ... + class Context: @typing.overload - def __init__(self) -> None: ... + def __init__(self, thread_count: typing.SupportsInt = -1) -> None: ... @typing.overload - def __init__(self, device: Device) -> None: ... + def __init__( + self, device: Device, thread_count: typing.SupportsInt = -1 + ) -> None: ... + def copy_globals_from(self, context: Context) -> None: ... def define_global( self, name: str, @@ -161,14 +259,16 @@ class Context: shape: collections.abc.Sequence[typing.SupportsInt], requires_grad: bool = False, ) -> Tensor: ... + def delete_global(self, name: str) -> None: ... def device(self) -> Device: ... def get_global(self, name: str) -> Tensor: ... def global_exists(self, name: str) -> bool: ... + def global_names(self) -> list[str]: ... def global_requires_grad(self, name: str) -> bool: ... - def load(self, file: str) -> None: ... + def load_globals(self, dir: str) -> None: ... def load_matrix_element(self, file: str, param_card: str) -> MatrixElementApi: ... def matrix_element(self, index: typing.SupportsInt) -> MatrixElementApi: ... - def save(self, file: str) -> None: ... + def save_globals(self, dir: str) -> None: ... class CutItem: def __init__( @@ -354,7 +454,9 @@ class DiscreteHistogram(FunctionGenerator): class DiscreteOptimizer: def __init__( - self, context: Context, prob_names: collections.abc.Sequence[str] + self, + contexts: collections.abc.Sequence[Context], + prob_names: collections.abc.Sequence[str], ) -> None: ... def add_data( self, values_and_counts: collections.abc.Sequence[typing.Any] @@ -372,6 +474,17 @@ class DiscreteSampler(Mapping): def option_counts(self) -> list[int]: ... def prob_names(self) -> list[str]: ... +class DoubleT(Mapping): + def __init__( + self, + t1_invariant_power: typing.SupportsFloat = 0.0, + t1_mass: typing.SupportsFloat = 0.0, + t1_width: typing.SupportsFloat = 0.0, + t2_invariant_power: typing.SupportsFloat = 0.0, + t2_mass: typing.SupportsFloat = 0.0, + t2_width: typing.SupportsFloat = 0.0, + ) -> None: ... + class EnergyScale(FunctionGenerator): class DynamicalScaleType: """ @@ -455,181 +568,27 @@ class EnergyScale(FunctionGenerator): class EventGenerator: default_config: typing.ClassVar[ - EventGeneratorConfig - ] # value = - integrand_flags: typing.ClassVar[int] = 1077 + GeneratorConfig + ] # value = def __init__( self, - context: Context, - channels: collections.abc.Sequence[Integrand], - temp_file_prefix: str, + contexts: collections.abc.Sequence[Context], + channels: collections.abc.Sequence[ChannelEventGenerator], status_file: str = "", - config: EventGeneratorConfig = ..., - channel_subprocesses: collections.abc.Sequence[typing.SupportsInt] = [], - channel_names: collections.abc.Sequence[str] = [], - channel_histograms: collections.abc.Sequence[ObservableHistograms] = [], + config: GeneratorConfig = ..., ) -> None: ... - def channel_status(self) -> list[EventGeneratorStatus]: ... + def channel_status(self) -> list[GeneratorStatus]: ... + def channels(self) -> list[ChannelEventGenerator]: ... def combine_to_compact_npy(self, file_name: str) -> None: ... - def combine_to_lhe(self, file_name: str, lhe_completer: ...) -> None: ... - def combine_to_lhe_npy(self, file_name: str, lhe_completer: ...) -> None: ... + def combine_to_lhe(self, file_name: str, lhe_completer: LHECompleter) -> None: ... + def combine_to_lhe_npy( + self, file_name: str, lhe_completer: LHECompleter + ) -> None: ... def generate(self) -> None: ... - def histograms(self) -> list[EventGeneratorHistogram]: ... - def status(self) -> EventGeneratorStatus: ... + def histograms(self) -> list[Histogram]: ... + def status(self) -> GeneratorStatus: ... def survey(self) -> None: ... - -class EventGeneratorConfig: - verbosity: EventGeneratorVerbosity - def __init__(self) -> None: ... - @property - def batch_size(self) -> int: ... - @batch_size.setter - def batch_size(self, arg0: typing.SupportsInt) -> None: ... - @property - def freeze_max_weight_after(self) -> int: ... - @freeze_max_weight_after.setter - def freeze_max_weight_after(self, arg0: typing.SupportsInt) -> None: ... - @property - def max_batch_size(self) -> int: ... - @max_batch_size.setter - def max_batch_size(self, arg0: typing.SupportsInt) -> None: ... - @property - def max_overweight_truncation(self) -> float: ... - @max_overweight_truncation.setter - def max_overweight_truncation(self, arg0: typing.SupportsFloat) -> None: ... - @property - def optimization_patience(self) -> int: ... - @optimization_patience.setter - def optimization_patience(self, arg0: typing.SupportsInt) -> None: ... - @property - def optimization_threshold(self) -> float: ... - @optimization_threshold.setter - def optimization_threshold(self, arg0: typing.SupportsFloat) -> None: ... - @property - def start_batch_size(self) -> int: ... - @start_batch_size.setter - def start_batch_size(self, arg0: typing.SupportsInt) -> None: ... - @property - def survey_max_iters(self) -> int: ... - @survey_max_iters.setter - def survey_max_iters(self, arg0: typing.SupportsInt) -> None: ... - @property - def survey_min_iters(self) -> int: ... - @survey_min_iters.setter - def survey_min_iters(self, arg0: typing.SupportsInt) -> None: ... - @property - def survey_target_precision(self) -> float: ... - @survey_target_precision.setter - def survey_target_precision(self, arg0: typing.SupportsFloat) -> None: ... - @property - def target_count(self) -> int: ... - @target_count.setter - def target_count(self, arg0: typing.SupportsInt) -> None: ... - @property - def vegas_damping(self) -> float: ... - @vegas_damping.setter - def vegas_damping(self, arg0: typing.SupportsFloat) -> None: ... - -class EventGeneratorHistogram: - @property - def bin_errors(self) -> list[float]: ... - @property - def bin_values(self) -> list[float]: ... - @property - def max(self) -> float: ... - @property - def min(self) -> float: ... - @property - def name(self) -> str: ... - -class EventGeneratorStatus: - done: bool - def __init__(self) -> None: ... - @property - def count(self) -> int: ... - @count.setter - def count(self, arg0: typing.SupportsInt) -> None: ... - @property - def count_after_cuts(self) -> int: ... - @count_after_cuts.setter - def count_after_cuts(self, arg0: typing.SupportsInt) -> None: ... - @property - def count_after_cuts_opt(self) -> int: ... - @count_after_cuts_opt.setter - def count_after_cuts_opt(self, arg0: typing.SupportsInt) -> None: ... - @property - def count_opt(self) -> int: ... - @count_opt.setter - def count_opt(self, arg0: typing.SupportsInt) -> None: ... - @property - def count_target(self) -> float: ... - @count_target.setter - def count_target(self, arg0: typing.SupportsFloat) -> None: ... - @property - def count_unweighted(self) -> float: ... - @count_unweighted.setter - def count_unweighted(self, arg0: typing.SupportsFloat) -> None: ... - @property - def error(self) -> float: ... - @error.setter - def error(self, arg0: typing.SupportsFloat) -> None: ... - @property - def index(self) -> int: ... - @index.setter - def index(self, arg0: typing.SupportsInt) -> None: ... - @property - def iterations(self) -> int: ... - @iterations.setter - def iterations(self, arg0: typing.SupportsInt) -> None: ... - @property - def mean(self) -> float: ... - @mean.setter - def mean(self, arg0: typing.SupportsFloat) -> None: ... - @property - def rel_std_dev(self) -> float: ... - @rel_std_dev.setter - def rel_std_dev(self, arg0: typing.SupportsFloat) -> None: ... - -class EventGeneratorVerbosity: - """ - Members: - - silent - - log - - pretty - """ - - __members__: typing.ClassVar[ - dict[str, EventGeneratorVerbosity] - ] # value = {'silent': , 'log': , 'pretty': } - log: typing.ClassVar[ - EventGeneratorVerbosity - ] # value = - pretty: typing.ClassVar[ - EventGeneratorVerbosity - ] # value = - silent: typing.ClassVar[ - EventGeneratorVerbosity - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... - @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... - @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... - @property - def name(self) -> str: ... - @property - def value(self) -> int: ... + def used_globals(self) -> set[str]: ... class FastRamboMapping(Mapping): def __init__(self, n_particles: typing.SupportsInt, massless: bool) -> None: ... @@ -654,27 +613,24 @@ class Flow(Mapping): class Function: @staticmethod def load(file: str) -> Function: ... - def __call__(self, *args): ... + def __call__(self, *args, **kwargs): ... def __repr__(self) -> str: ... def __str__(self) -> str: ... def save(self, file: str) -> None: ... @property - def globals(self) -> dict[str, Value]: ... + def globals(self) -> list[tuple[str, Value]]: ... @property - def inputs(self) -> list[Value]: ... + def inputs(self) -> NamedValues: ... @property def instructions(self) -> list[InstructionCall]: ... @property def locals(self) -> list[Value]: ... @property - def outputs(self) -> list[Value]: ... + def outputs(self) -> NamedValues: ... class FunctionBuilder: - def __init__( - self, - input_types: collections.abc.Sequence[Type], - output_types: collections.abc.Sequence[Type], - ) -> None: ... + def __init__(self, input_types: NamedTypes, output_types: NamedTypes) -> None: ... + def accept_norm(self, accepted_batch: Value, full_batch: Value) -> Value: ... def add(self, in1: Value, in2: Value) -> Value: ... def add_int(self, in1: Value, in2: Value) -> Value: ... def apply_subchannel_weights( @@ -689,6 +645,8 @@ class FunctionBuilder: self, args: collections.abc.Sequence[Value] ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... def batch_gather(self, indices: Value, values: Value) -> Value: ... + def batch_reduce_mean(self, in1: Value) -> Value: ... + def batch_reduce_mean_keepdim(self, in1: Value) -> Value: ... def batch_scatter(self, indices: Value, target: Value, source: Value) -> Value: ... def batch_size(self, args: collections.abc.Sequence[Value]) -> Value: ... def batch_split(self, in1: Value, counts: Value) -> list[Value]: ... @@ -713,6 +671,7 @@ class FunctionBuilder: def com_p_in( self, e_cm: Value ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... + def current_stream(self) -> int: ... def cut_all(self, obs: Value, min: Value, max: Value) -> Value: ... def cut_any(self, obs: Value, min: Value, max: Value) -> Value: ... def cut_one(self, obs: Value, min: Value, max: Value) -> Value: ... @@ -729,6 +688,19 @@ class FunctionBuilder: def discrete_histogram( self, input: Value, weights: Value, option_count: Value ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... + def div(self, in1: Value, in2: Value) -> Value: ... + def double_t_scattering( + self, + r_phi: Value, + pa: Value, + pb: Value, + t1_abs: Value, + t2_abs: Value, + m1: Value, + ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... + def double_t_scattering_inverse( + self, p1: Value, p2: Value, pa: Value, pb: Value + ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... def elu(self, in1: Value) -> Value: ... def fast_rambo_massive( self, r: Value, e_cm: Value, masses: Value, p0: Value @@ -780,6 +752,13 @@ class FunctionBuilder: ) -> Value: ... def invariants_from_momenta(self, p_ext: Value, factors: Value) -> Value: ... def leaky_relu(self, in1: Value) -> Value: ... + def madnis_abs_weight(self, f: Value, q: Value) -> Value: ... + def madnis_multi_channel_variance(self, vars: Value, abs_means: Value) -> Value: ... + def madnis_single_channel_variance(self, var: Value, abs_mean: Value) -> Value: ... + def madnis_softclip( + self, f: Value, q: Value, norm: Value, threshold: Value + ) -> Value: ... + def madnis_variance(self, f: Value, g: Value, q: Value, mean: Value) -> Value: ... def matmul(self, x: Value, weight: Value, bias: Value) -> Value: ... def matrix_element(self, args: collections.abc.Sequence[Value]) -> list[Value]: ... def max(self, in1: Value, in2: Value) -> Value: ... @@ -869,6 +848,7 @@ class FunctionBuilder: def select(self, input: Value, indices: Value) -> Value: ... def select_int(self, input: Value, indices: Value) -> Value: ... def select_vector(self, input: Value, indices: Value) -> Value: ... + def set_current_stream(self, arg0: typing.SupportsInt) -> None: ... def sigmoid(self, in1: Value) -> Value: ... def softmax(self, input: Value) -> Value: ... def softmax_prior(self, input: Value, prior: Value) -> Value: ... @@ -889,6 +869,7 @@ class FunctionBuilder: self, s: Value, mass: Value, nu: Value, s_min: Value, s_max: Value ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... def stack(self, args: collections.abc.Sequence[Value]) -> Value: ... + def stack_sizes(self, args: collections.abc.Sequence[Value]) -> Value: ... def sub(self, in1: Value, in2: Value) -> Value: ... def subchannel_weights( self, @@ -899,6 +880,18 @@ class FunctionBuilder: on_shell: Value, group_sizes: Value, ) -> Value: ... + def t1_inv_min_max_doublet( + self, pa: Value, pb: Value, m1: Value, mir_min: Value + ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... + def t1_inv_value_and_min_max_doublet( + self, pa: Value, pb: Value, p1: Value, m1: Value, mir_min: Value + ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... + def t2_inv_min_max_doublet( + self, pa: Value, pb: Value, m1: Value, mir_min: Value, t1_abs: Value + ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... + def t2_inv_value_and_min_max_doublet( + self, pa: Value, pb: Value, p1: Value, m1: Value, mir_min: Value, t1_abs: Value + ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... def t_inv_min_max( self, pa: Value, pb: Value, m1: Value, m2: Value ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... @@ -1013,16 +1006,18 @@ class FunctionBuilder: ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... class FunctionGenerator: - def __call__(self, *args): ... + def __call__(self, *args, **kwargs): ... def __init__( - self, - name: str, - arg_types: collections.abc.Sequence[Type], - return_types: collections.abc.Sequence[Type], + self, name: str, arg_types: NamedTypes, return_types: NamedTypes ) -> None: ... + @typing.overload def build_function( self, builder: FunctionBuilder, args: collections.abc.Sequence[Value] - ) -> list[Value]: ... + ) -> NamedValues: ... + @typing.overload + def build_function( + self, builder: FunctionBuilder, args: NamedValues + ) -> NamedValues: ... def function(self) -> Function: ... class FunctionRuntime: @@ -1047,6 +1042,117 @@ class FunctionRuntime: arg1: collections.abc.Sequence[bool], ) -> tuple[list[Tensor], list[madspace._madspace_py.Tensor | None], list[bool]]: ... +class GeneratorConfig: + verbosity: Verbosity + write_live_data: bool + def __init__(self) -> None: ... + @property + def combine_thread_count(self) -> int: ... + @combine_thread_count.setter + def combine_thread_count(self, arg0: typing.SupportsInt) -> None: ... + @property + def cpu_batch_size(self) -> int: ... + @cpu_batch_size.setter + def cpu_batch_size(self, arg0: typing.SupportsInt) -> None: ... + @property + def freeze_max_weight_after(self) -> int: ... + @freeze_max_weight_after.setter + def freeze_max_weight_after(self, arg0: typing.SupportsInt) -> None: ... + @property + def gpu_batch_size(self) -> int: ... + @gpu_batch_size.setter + def gpu_batch_size(self, arg0: typing.SupportsInt) -> None: ... + @property + def max_batch_size(self) -> int: ... + @max_batch_size.setter + def max_batch_size(self, arg0: typing.SupportsInt) -> None: ... + @property + def max_overweight_truncation(self) -> float: ... + @max_overweight_truncation.setter + def max_overweight_truncation(self, arg0: typing.SupportsFloat) -> None: ... + @property + def optimization_patience(self) -> int: ... + @optimization_patience.setter + def optimization_patience(self, arg0: typing.SupportsInt) -> None: ... + @property + def optimization_threshold(self) -> float: ... + @optimization_threshold.setter + def optimization_threshold(self, arg0: typing.SupportsFloat) -> None: ... + @property + def start_batch_size(self) -> int: ... + @start_batch_size.setter + def start_batch_size(self, arg0: typing.SupportsInt) -> None: ... + @property + def survey_max_iters(self) -> int: ... + @survey_max_iters.setter + def survey_max_iters(self, arg0: typing.SupportsInt) -> None: ... + @property + def survey_min_iters(self) -> int: ... + @survey_min_iters.setter + def survey_min_iters(self, arg0: typing.SupportsInt) -> None: ... + @property + def survey_target_precision(self) -> float: ... + @survey_target_precision.setter + def survey_target_precision(self, arg0: typing.SupportsFloat) -> None: ... + @property + def target_count(self) -> int: ... + @target_count.setter + def target_count(self, arg0: typing.SupportsInt) -> None: ... + @property + def vegas_damping(self) -> float: ... + @vegas_damping.setter + def vegas_damping(self, arg0: typing.SupportsFloat) -> None: ... + +class GeneratorStatus: + done: bool + name: str + optimized: bool + def __init__(self) -> None: ... + @property + def count(self) -> int: ... + @count.setter + def count(self, arg0: typing.SupportsInt) -> None: ... + @property + def count_after_cuts(self) -> int: ... + @count_after_cuts.setter + def count_after_cuts(self, arg0: typing.SupportsInt) -> None: ... + @property + def count_after_cuts_opt(self) -> int: ... + @count_after_cuts_opt.setter + def count_after_cuts_opt(self, arg0: typing.SupportsInt) -> None: ... + @property + def count_opt(self) -> int: ... + @count_opt.setter + def count_opt(self, arg0: typing.SupportsInt) -> None: ... + @property + def count_target(self) -> float: ... + @count_target.setter + def count_target(self, arg0: typing.SupportsFloat) -> None: ... + @property + def count_unweighted(self) -> float: ... + @count_unweighted.setter + def count_unweighted(self, arg0: typing.SupportsFloat) -> None: ... + @property + def error(self) -> float: ... + @error.setter + def error(self, arg0: typing.SupportsFloat) -> None: ... + @property + def iterations(self) -> int: ... + @iterations.setter + def iterations(self, arg0: typing.SupportsInt) -> None: ... + @property + def mean(self) -> float: ... + @mean.setter + def mean(self, arg0: typing.SupportsFloat) -> None: ... + @property + def rel_std_dev(self) -> float: ... + @rel_std_dev.setter + def rel_std_dev(self, arg0: typing.SupportsFloat) -> None: ... + @property + def subprocess(self) -> int: ... + @subprocess.setter + def subprocess(self, arg0: typing.SupportsInt) -> None: ... + class HistItem: def __init__( self, @@ -1064,6 +1170,18 @@ class HistItem: @property def observable(self) -> Observable: ... +class Histogram: + @property + def bin_errors(self) -> list[float]: ... + @property + def bin_values(self) -> list[float]: ... + @property + def max(self) -> float: ... + @property + def min(self) -> float: ... + @property + def name(self) -> str: ... + class Instruction: def __str__(self) -> str: ... @property @@ -1082,6 +1200,8 @@ class InstructionCall: def outputs(self) -> list[Value]: ... class Integrand(FunctionGenerator): + drop_cuts_and_rescale: typing.ClassVar[int] = 8192 + exclude_adaptive_and_chan_weight: typing.ClassVar[int] = 4096 matrix_element_inputs: typing.ClassVar[ list ] # value = [, , , , , ] @@ -1127,6 +1247,8 @@ class Integrand(FunctionGenerator): flags: typing.SupportsInt = 0, channel_indices: collections.abc.Sequence[typing.SupportsInt] = [], active_flavors: collections.abc.Sequence[typing.SupportsInt] = [], + flavor_remap: collections.abc.Sequence[typing.SupportsInt] = [], + flavor_factors: collections.abc.Sequence[typing.SupportsFloat] = [], ) -> None: ... def adaptive_map( self, @@ -1170,20 +1292,14 @@ class Invariant(Mapping): ) -> None: ... class LHECompleter: + @staticmethod + def load(file: str) -> LHECompleter: ... def __init__( self, subproc_args: collections.abc.Sequence[SubprocArgs], bw_cutoff: typing.SupportsFloat, ) -> None: ... - def complete_event_data( - self, - event: LHEEvent, - subprocess_index: typing.SupportsInt, - diagram_index: typing.SupportsInt, - color_index: typing.SupportsInt, - flavor_index: typing.SupportsInt, - helicity_index: typing.SupportsInt, - ) -> None: ... + def save(self, file: str) -> None: ... @property def max_particle_count(self) -> int: ... @@ -1539,30 +1655,154 @@ class MLP(FunctionGenerator): def input_dim(self) -> int: ... def output_dim(self) -> int: ... +class MadnisConfig: + lr_schedule: AdamOptimizer.LRSchedule + verbosity: Verbosity + def __init__(self) -> None: ... + @property + def adam_beta1(self) -> float: ... + @adam_beta1.setter + def adam_beta1(self, arg0: typing.SupportsFloat) -> None: ... + @property + def adam_beta2(self) -> float: ... + @adam_beta2.setter + def adam_beta2(self, arg0: typing.SupportsFloat) -> None: ... + @property + def adam_eps(self) -> float: ... + @adam_eps.setter + def adam_eps(self, arg0: typing.SupportsFloat) -> None: ... + @property + def batch_size_offset(self) -> int: ... + @batch_size_offset.setter + def batch_size_offset(self, arg0: typing.SupportsInt) -> None: ... + @property + def batch_size_per_channel(self) -> int: ... + @batch_size_per_channel.setter + def batch_size_per_channel(self, arg0: typing.SupportsInt) -> None: ... + @property + def batches(self) -> int: ... + @batches.setter + def batches(self, arg0: typing.SupportsInt) -> None: ... + @property + def buffer_capacity(self) -> int: ... + @buffer_capacity.setter + def buffer_capacity(self, arg0: typing.SupportsInt) -> None: ... + @property + def buffer_unweighting_quantile(self) -> float: ... + @buffer_unweighting_quantile.setter + def buffer_unweighting_quantile(self, arg0: typing.SupportsFloat) -> None: ... + @property + def buffered_steps(self) -> int: ... + @buffered_steps.setter + def buffered_steps(self, arg0: typing.SupportsInt) -> None: ... + @property + def channel_dropping_interval(self) -> int: ... + @channel_dropping_interval.setter + def channel_dropping_interval(self, arg0: typing.SupportsInt) -> None: ... + @property + def channel_dropping_threshold(self) -> float: ... + @channel_dropping_threshold.setter + def channel_dropping_threshold(self, arg0: typing.SupportsFloat) -> None: ... + @property + def cpu_generator_batch_size(self) -> int: ... + @cpu_generator_batch_size.setter + def cpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: ... + @property + def fixed_cwnet_fraction(self) -> float: ... + @fixed_cwnet_fraction.setter + def fixed_cwnet_fraction(self, arg0: typing.SupportsFloat) -> None: ... + @property + def generator_target_size_factor(self) -> int: ... + @generator_target_size_factor.setter + def generator_target_size_factor(self, arg0: typing.SupportsInt) -> None: ... + @property + def gpu_generator_batch_granularity(self) -> int: ... + @gpu_generator_batch_granularity.setter + def gpu_generator_batch_granularity(self, arg0: typing.SupportsInt) -> None: ... + @property + def gpu_generator_batch_size(self) -> int: ... + @gpu_generator_batch_size.setter + def gpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: ... + @property + def integration_history_length(self) -> int: ... + @integration_history_length.setter + def integration_history_length(self, arg0: typing.SupportsInt) -> None: ... + @property + def learning_rate(self) -> float: ... + @learning_rate.setter + def learning_rate(self, arg0: typing.SupportsFloat) -> None: ... + @property + def log_interval(self) -> int: ... + @log_interval.setter + def log_interval(self, arg0: typing.SupportsInt) -> None: ... + @property + def minimum_buffer_size(self) -> int: ... + @minimum_buffer_size.setter + def minimum_buffer_size(self, arg0: typing.SupportsInt) -> None: ... + @property + def softclip_threshold(self) -> float: ... + @softclip_threshold.setter + def softclip_threshold(self, arg0: typing.SupportsFloat) -> None: ... + @property + def uniform_channel_ratio(self) -> float: ... + @uniform_channel_ratio.setter + def uniform_channel_ratio(self, arg0: typing.SupportsFloat) -> None: ... + +class MadnisLoss(FunctionGenerator): + def __init__( + self, + functions: collections.abc.Sequence[FunctionGenerator], + cwnet: madspace._madspace_py.ChannelWeightNetwork | None, + softclip_threshold: typing.SupportsFloat = 0.0, + ) -> None: ... + +class MadnisTraining: + def __init__( + self, + generator_context: Context, + optimizer_context: Context, + config: MadnisConfig, + integrands: collections.abc.Sequence[Integrand], + cwnet: madspace._madspace_py.ChannelWeightNetwork | None, + ) -> None: ... + def active_channel_count(self) -> int: ... + def active_channels(self) -> list[int]: ... + def train_step(self, batch_index: typing.SupportsInt) -> None: ... + class Mapping: def __init__( self, name: str, - input_types: collections.abc.Sequence[Type], - output_types: collections.abc.Sequence[Type], - condition_types: collections.abc.Sequence[Type], + input_types: NamedTypes, + output_types: NamedTypes, + condition_types: NamedTypes, ) -> None: ... + @typing.overload def build_forward( self, builder: FunctionBuilder, inputs: collections.abc.Sequence[Value], conditions: collections.abc.Sequence[Value], - ) -> tuple[list[Value], Value]: ... + ) -> NamedValues: ... + @typing.overload + def build_forward( + self, builder: FunctionBuilder, inputs: NamedValues, conditions: NamedValues + ) -> NamedValues: ... + @typing.overload def build_inverse( self, builder: FunctionBuilder, inputs: collections.abc.Sequence[Value], conditions: collections.abc.Sequence[Value], - ) -> tuple[list[Value], Value]: ... + ) -> NamedValues: ... + @typing.overload + def build_inverse( + self, builder: FunctionBuilder, inputs: NamedValues, conditions: NamedValues + ) -> NamedValues: ... def forward_function(self) -> Function: ... def inverse_function(self) -> Function: ... - def map_forward(self, inputs, conditions=list()): ... - def map_inverse(self, inputs, conditions=list()): ... + def map_forward(self, inputs=list(), conditions=list(), **kwargs): ... + def map_inverse(self, inputs=list(), conditions=list(), **kwargs): ... class MatrixElement(FunctionGenerator): class MatrixElementInput: @@ -1744,10 +1984,6 @@ class MatrixElement(FunctionGenerator): def particle_count(self) -> int: ... class MatrixElementApi: - def __init__( - self, file: str, param_card: str, index: typing.SupportsInt = 0 - ) -> None: ... - def device(self) -> Device: ... def diagram_count(self) -> int: ... def helicity_count(self) -> int: ... def index(self) -> int: ... @@ -1763,11 +1999,71 @@ class MultiChannelFunction(FunctionGenerator): ) -> None: ... class MultiChannelIntegrand(FunctionGenerator): - def __init__(self, integrands: collections.abc.Sequence[Integrand]) -> None: ... + def __init__( + self, + integrands: collections.abc.Sequence[Integrand], + return_sizes: bool = False, + ) -> None: ... class MultiChannelMapping(Mapping): def __init__(self, mappings: collections.abc.Sequence[Mapping]) -> None: ... +class MultiMadnisTraining: + def __init__( + self, + generator_context: Context, + optimizer_context: Context, + config: MadnisConfig, + integrands: collections.abc.Sequence[collections.abc.Sequence[Integrand]], + cwnets: collections.abc.Sequence[ + madspace._madspace_py.ChannelWeightNetwork | None + ], + ) -> None: ... + def active_channels(self) -> list[list[int]]: ... + def train(self) -> None: ... + +class NamedTypes: + @typing.overload + def __getitem__(self, key: typing.SupportsInt) -> Type: ... + @typing.overload + def __getitem__(self, key: str) -> Type: ... + @typing.overload + def __init__(self) -> None: ... + @typing.overload + def __init__( + self, + keys: collections.abc.Sequence[str], + values: collections.abc.Sequence[Type], + ) -> None: ... + @typing.overload + def __init__(self, items: collections.abc.Sequence[tuple[str, Type]]) -> None: ... + def __len__(self) -> int: ... + def index_map(self) -> dict[str, int]: ... + def keys(self) -> list[str]: ... + def push_back(self, name: str, item: Type) -> None: ... + def values(self) -> list[Type]: ... + +class NamedValues: + @typing.overload + def __getitem__(self, key: typing.SupportsInt) -> Value: ... + @typing.overload + def __getitem__(self, key: str) -> Value: ... + @typing.overload + def __init__(self) -> None: ... + @typing.overload + def __init__( + self, + keys: collections.abc.Sequence[str], + values: collections.abc.Sequence[Value], + ) -> None: ... + @typing.overload + def __init__(self, items: collections.abc.Sequence[tuple[str, Value]]) -> None: ... + def __len__(self) -> int: ... + def index_map(self) -> dict[str, int]: ... + def keys(self) -> list[str]: ... + def push_back(self, name: str, item: Value) -> None: ... + def values(self) -> list[Value]: ... + class Observable(FunctionGenerator): class ObservableOption: """ @@ -2251,6 +2547,7 @@ class TPropagatorMapping(Mapping): integration_order: collections.abc.Sequence[typing.SupportsInt], invariant_power: typing.SupportsFloat = 0.0, ) -> None: ... + def random_dim(self) -> int: ... class Tensor: @staticmethod @@ -2340,7 +2637,7 @@ class Type: def shape(self) -> list[int]: ... class Unweighter(FunctionGenerator): - def __init__(self, types: collections.abc.Sequence[Type]) -> None: ... + def __init__(self, types: NamedTypes) -> None: ... class Value: @typing.overload @@ -2360,7 +2657,10 @@ class Value: class VegasGridOptimizer: def __init__( - self, context: Context, grid_name: str, damping: typing.SupportsFloat + self, + contexts: collections.abc.Sequence[Context], + grid_name: str, + damping: typing.SupportsFloat, ) -> None: ... def add_data(self, values: typing.Any, counts: typing.Any) -> None: ... def optimize(self) -> None: ... @@ -2380,13 +2680,48 @@ class VegasMapping(Mapping): def grid_name(self) -> str: ... def initialize_globals(self, context: Context) -> None: ... +class Verbosity: + """ + Members: + + silent + + log + + pretty + """ + + __members__: typing.ClassVar[ + dict[str, Verbosity] + ] # value = {'silent': , 'log': , 'pretty': } + log: typing.ClassVar[Verbosity] # value = + pretty: typing.ClassVar[Verbosity] # value = + silent: typing.ClassVar[Verbosity] # value = + def __eq__(self, other: typing.Any) -> bool: ... + def __getstate__(self) -> int: ... + def __hash__(self) -> int: ... + def __index__(self) -> int: ... + @typing.overload + def __init__(self, value: typing.SupportsInt) -> None: ... + @typing.overload + def __init__(self, name: str) -> None: ... + def __int__(self) -> int: ... + def __ne__(self, other: typing.Any) -> bool: ... + def __repr__(self) -> str: ... + def __setstate__(self, state: typing.SupportsInt) -> None: ... + def __str__(self) -> str: ... + @property + def name(self) -> str: ... + @property + def value(self) -> int: ... + def batch_float_array(count: typing.SupportsInt) -> Type: ... def batch_four_vec_array(count: typing.SupportsInt) -> Type: ... def cpu_device() -> Device: ... -def cuda_device() -> Device: ... +def cuda_device(index: typing.SupportsInt = 0) -> Device: ... def default_context() -> Context: ... -def default_cuda_context() -> Context: ... -def default_hip_context() -> Context: ... +def default_cuda_context(index: typing.SupportsInt = 0) -> Context: ... +def default_hip_context(index: typing.SupportsInt = 0) -> Context: ... def format_progress( progress: typing.SupportsFloat, width: typing.SupportsInt ) -> str: ... @@ -2394,12 +2729,11 @@ def format_si_prefix(value: typing.SupportsFloat) -> str: ... def format_with_error( value: typing.SupportsFloat, error: typing.SupportsFloat ) -> str: ... -def hip_device() -> Device: ... +def hip_device(index: typing.SupportsInt = 0) -> Device: ... def initialize_vegas_grid(context: Context, grid_name: str) -> None: ... def multichannel_batch_size(count: typing.SupportsInt) -> Type: ... def set_lib_path(lib_path: str) -> None: ... def set_simd_vector_size(vector_size: typing.SupportsInt) -> None: ... -def set_thread_count(new_count: typing.SupportsInt) -> None: ... batch_float: Type # value = float[batch_size] batch_four_vec: Type # value = float[batch_size, 4] @@ -2408,8 +2742,8 @@ batch_size: BatchSize # value = batch_size batch_sizes: DataType # value = float: DataType # value = int: DataType # value = -log: EventGeneratorVerbosity # value = -pretty: EventGeneratorVerbosity # value = -silent: EventGeneratorVerbosity # value = +log: Verbosity # value = +pretty: Verbosity # value = +silent: Verbosity # value = single_float: Type # value = float[1] single_int: Type # value = int[1] diff --git a/madspace/src/kernels/kernels.hpp b/madspace/src/kernels/kernels.hpp index c43636772..7c2d828ef 100644 --- a/madspace/src/kernels/kernels.hpp +++ b/madspace/src/kernels/kernels.hpp @@ -5,6 +5,7 @@ #include "discrete.hpp" #include "invariants.hpp" #include "kinematics.hpp" +#include "lup_det.hpp" #include "math.hpp" #include "multichannel.hpp" #include "nn.hpp" diff --git a/madspace/src/kernels/kinematics.hpp b/madspace/src/kernels/kinematics.hpp index a4f3a38f1..a0ba7916f 100644 --- a/madspace/src/kernels/kinematics.hpp +++ b/madspace/src/kernels/kinematics.hpp @@ -121,8 +121,6 @@ t2_inv_min_max_doublet(FVal s, FVal m1_2, FVal mir_min_2, FVal t1_ab // |t2| bounds given fixed s, m_i, mir_min, |t1|. // Derivation (abs-value convention): from the signed-t formulas in the // Fortran double_t, flip both signs. - // t2_min_abs = m_i^2 * (s - t1_abs - m_i^2) / (m_i^2 + t1_abs) - // t2_max_abs = s - t1_abs - m_i^2 - mir_min^2 auto denom = m1_2 + t1_abs + EPS; auto t2_min_raw = m1_2 * (s - t1_abs - m1_2) / denom; auto t2_max_raw = s - t1_abs - m1_2 - mir_min_2; @@ -228,7 +226,7 @@ rotate_two_ref(FourMom p, FourMom q_z, FourMom q_x) { // Forward rotation: take p from a canonical frame // (e_z along q_z, e_x in the plane spanned by q_z and q_x with positive // component along q_x's perpendicular part) - // into the world frame. Energy unchanged. + // into the lab frame. Energy unchanged. // // Used by kernel_two_to_three_particle_scattering: q_z = pa_com, // q_x = p3 boosted into the p_12 rest frame. Picks a unique azimuth @@ -265,7 +263,7 @@ rotate_two_ref(FourMom p, FourMom q_z, FourMom q_x) { template KERNELSPEC FourMom rotate_two_ref_inverse(FourMom p, FourMom q_z, FourMom q_x) { - // Inverse of rotate_two_ref: take a vector p from the world frame into the + // Inverse of rotate_two_ref: take a vector p from the lab frame into the // canonical frame defined by (q_z, q_x). Energy unchanged. auto qz_n2 = q_z[1]*q_z[1] + q_z[2]*q_z[2] + q_z[3]*q_z[3]; diff --git a/madspace/src/kernels/lup_det.hpp b/madspace/src/kernels/lup_det.hpp new file mode 100644 index 000000000..da0b8b704 --- /dev/null +++ b/madspace/src/kernels/lup_det.hpp @@ -0,0 +1,371 @@ +#pragma once + +#include "definitions.hpp" + +namespace madspace { +namespace kernels { + +// Local tolerance constant, intentionally tighter than the module-level EPS. +// Chosen so the relative tolerance scale * LU_TOL sits above eps_mach but +// below the noise floor at which the polynomial expansion blows up. +inline constexpr double LU_TOL = 1e-14; + +// Variadic max-of-|.| over an arbitrary list of FVal. +template +KERNELSPEC FVal max_abs(FVal first, Args... rest) { + if constexpr (sizeof...(rest) == 0) { + return fabs(first); + } else { + return max(fabs(first), max_abs(rest...)); + } +} + +// det of a symmetric 3x3 matrix via LU with partial pivoting. +// +// [ a11 a12 a13 ] +// [ a12 a22 a23 ] +// [ a13 a23 a33 ] +// +// Falls back lane-wise to `polynomial_fallback` if any pivot is below +// LU_TOL * max(|entries|). +template +KERNELSPEC FVal lup_det3( + FVal a11, FVal a12, FVal a13, + FVal a22, FVal a23, + FVal a33, + FVal polynomial_fallback +) { + // Row-major working copy. + auto r00 = a11, r01 = a12, r02 = a13; + auto r10 = a12, r11 = a22, r12 = a23; + auto r20 = a13, r21 = a23, r22 = a33; + + auto scale = max_abs(a11, a12, a13, a22, a23, a33); + auto tol = LU_TOL * scale; + auto sign = FVal(1.0); + auto degenerate = (scale == 0.); + + // ---- Step 0: pivot from column 0, rows {0,1,2} ---- + { + auto m0 = fabs(r00), m1 = fabs(r10), m2 = fabs(r20); + auto best = IVal(0); + auto best_mag = m0; + auto pick1 = (m1 > best_mag); + best = where(pick1, IVal(1), best); + best_mag = where(pick1, m1, best_mag); + auto pick2 = (m2 > best_mag); + best = where(pick2, IVal(2), best); + best_mag = where(pick2, m2, best_mag); + + auto sel1 = (best == IVal(1)); + auto sel2 = (best == IVal(2)); + auto any_swap = sel1 | sel2; + + auto n00 = where(sel1, r10, where(sel2, r20, r00)); + auto n01 = where(sel1, r11, where(sel2, r21, r01)); + auto n02 = where(sel1, r12, where(sel2, r22, r02)); + auto n10 = where(sel1, r00, r10); + auto n11 = where(sel1, r01, r11); + auto n12 = where(sel1, r02, r12); + auto n20 = where(sel2, r00, r20); + auto n21 = where(sel2, r01, r21); + auto n22 = where(sel2, r02, r22); + r00 = n00; r01 = n01; r02 = n02; + r10 = n10; r11 = n11; r12 = n12; + r20 = n20; r21 = n21; r22 = n22; + + sign = where(any_swap, -sign, sign); + degenerate = degenerate | (best_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r00); + auto f10 = r10 / pivot_safe; + auto f20 = r20 / pivot_safe; + r10 = f10; + r11 = r11 - f10 * r01; + r12 = r12 - f10 * r02; + r20 = f20; + r21 = r21 - f20 * r01; + r22 = r22 - f20 * r02; + } + + // ---- Step 1: pivot from column 1, rows {1, 2} ---- + { + auto m1 = fabs(r11), m2 = fabs(r21); + auto sel2 = (m2 > m1); + auto pivot_mag = where(sel2, m2, m1); + + auto n11 = where(sel2, r21, r11); + auto n12 = where(sel2, r22, r12); + auto n21 = where(sel2, r11, r21); + auto n22 = where(sel2, r12, r22); + r11 = n11; r12 = n12; r21 = n21; r22 = n22; + // Multipliers in column 0 follow the row swap (preserves L, not + // strictly needed for the determinant but kept for correctness). + auto n10 = where(sel2, r20, r10); + auto n20 = where(sel2, r10, r20); + r10 = n10; r20 = n20; + + sign = where(sel2, -sign, sign); + degenerate = degenerate | (pivot_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r11); + auto f21 = r21 / pivot_safe; + r21 = f21; + r22 = r22 - f21 * r12; + } + + // ---- Final pivot check: the trailing diagonal entry is also a pivot ---- + degenerate = degenerate | (fabs(r22) < tol); + + auto lu_det = sign * r00 * r11 * r22; + return where(degenerate, polynomial_fallback, lu_det); +} + +// det of a general (non-symmetric) 3x3 matrix via LU with partial pivoting. +// +// [ a11 a12 a13 ] +// [ a21 a22 a23 ] +// [ a31 a32 a33 ] +// +// Same fallback policy as lup_det3. Needed for bk_V, whose matrix has +// a21 = a12 but a13 != a31 and a23 != a32 (and a22 = a33 = 0). +template +KERNELSPEC FVal lup_det3_general( + FVal a11, FVal a12, FVal a13, + FVal a21, FVal a22, FVal a23, + FVal a31, FVal a32, FVal a33, + FVal polynomial_fallback +) { + auto r00 = a11, r01 = a12, r02 = a13; + auto r10 = a21, r11 = a22, r12 = a23; + auto r20 = a31, r21 = a32, r22 = a33; + + auto scale = max_abs(a11, a12, a13, a21, a22, a23, a31, a32, a33); + auto tol = LU_TOL * scale; + auto sign = FVal(1.0); + auto degenerate = (scale == 0.); + + // ---- Step 0: pivot from column 0, rows {0,1,2} ---- + { + auto m0 = fabs(r00), m1 = fabs(r10), m2 = fabs(r20); + auto best = IVal(0); + auto best_mag = m0; + auto pick1 = (m1 > best_mag); + best = where(pick1, IVal(1), best); + best_mag = where(pick1, m1, best_mag); + auto pick2 = (m2 > best_mag); + best = where(pick2, IVal(2), best); + best_mag = where(pick2, m2, best_mag); + + auto sel1 = (best == IVal(1)); + auto sel2 = (best == IVal(2)); + auto any_swap = sel1 | sel2; + + auto n00 = where(sel1, r10, where(sel2, r20, r00)); + auto n01 = where(sel1, r11, where(sel2, r21, r01)); + auto n02 = where(sel1, r12, where(sel2, r22, r02)); + auto n10 = where(sel1, r00, r10); + auto n11 = where(sel1, r01, r11); + auto n12 = where(sel1, r02, r12); + auto n20 = where(sel2, r00, r20); + auto n21 = where(sel2, r01, r21); + auto n22 = where(sel2, r02, r22); + r00 = n00; r01 = n01; r02 = n02; + r10 = n10; r11 = n11; r12 = n12; + r20 = n20; r21 = n21; r22 = n22; + + sign = where(any_swap, -sign, sign); + degenerate = degenerate | (best_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r00); + auto f10 = r10 / pivot_safe; + auto f20 = r20 / pivot_safe; + r10 = f10; + r11 = r11 - f10 * r01; + r12 = r12 - f10 * r02; + r20 = f20; + r21 = r21 - f20 * r01; + r22 = r22 - f20 * r02; + } + + // ---- Step 1: pivot from column 1, rows {1, 2} ---- + { + auto m1 = fabs(r11), m2 = fabs(r21); + auto sel2 = (m2 > m1); + auto pivot_mag = where(sel2, m2, m1); + + auto n11 = where(sel2, r21, r11); + auto n12 = where(sel2, r22, r12); + auto n21 = where(sel2, r11, r21); + auto n22 = where(sel2, r12, r22); + r11 = n11; r12 = n12; r21 = n21; r22 = n22; + auto n10 = where(sel2, r20, r10); + auto n20 = where(sel2, r10, r20); + r10 = n10; r20 = n20; + + sign = where(sel2, -sign, sign); + degenerate = degenerate | (pivot_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r11); + auto f21 = r21 / pivot_safe; + r21 = f21; + r22 = r22 - f21 * r12; + } + + degenerate = degenerate | (fabs(r22) < tol); + + auto lu_det = sign * r00 * r11 * r22; + return where(degenerate, polynomial_fallback, lu_det); +} + +// det of a symmetric 4x4 matrix via LU with partial pivoting. See lup_det3. +template +KERNELSPEC FVal lup_det4( + FVal a11, FVal a12, FVal a13, FVal a14, + FVal a22, FVal a23, FVal a24, + FVal a33, FVal a34, + FVal a44, + FVal polynomial_fallback +) { + auto r00 = a11, r01 = a12, r02 = a13, r03 = a14; + auto r10 = a12, r11 = a22, r12 = a23, r13 = a24; + auto r20 = a13, r21 = a23, r22 = a33, r23 = a34; + auto r30 = a14, r31 = a24, r32 = a34, r33 = a44; + + auto scale = max_abs(a11, a12, a13, a14, a22, a23, a24, a33, a34, a44); + auto tol = LU_TOL * scale; + auto sign = FVal(1.0); + auto degenerate = (scale == 0.); + + // ---- Step 0: pivot from column 0, rows {0,1,2,3} ---- + { + auto m0 = fabs(r00), m1 = fabs(r10), m2 = fabs(r20), m3 = fabs(r30); + auto best = IVal(0); + auto best_mag = m0; + auto pick1 = (m1 > best_mag); + best = where(pick1, IVal(1), best); + best_mag = where(pick1, m1, best_mag); + auto pick2 = (m2 > best_mag); + best = where(pick2, IVal(2), best); + best_mag = where(pick2, m2, best_mag); + auto pick3 = (m3 > best_mag); + best = where(pick3, IVal(3), best); + best_mag = where(pick3, m3, best_mag); + + auto sel1 = (best == IVal(1)); + auto sel2 = (best == IVal(2)); + auto sel3 = (best == IVal(3)); + auto any_swap = sel1 | sel2 | sel3; + + auto n00 = where(sel1, r10, where(sel2, r20, where(sel3, r30, r00))); + auto n01 = where(sel1, r11, where(sel2, r21, where(sel3, r31, r01))); + auto n02 = where(sel1, r12, where(sel2, r22, where(sel3, r32, r02))); + auto n03 = where(sel1, r13, where(sel2, r23, where(sel3, r33, r03))); + auto n10 = where(sel1, r00, r10); + auto n11 = where(sel1, r01, r11); + auto n12 = where(sel1, r02, r12); + auto n13 = where(sel1, r03, r13); + auto n20 = where(sel2, r00, r20); + auto n21 = where(sel2, r01, r21); + auto n22 = where(sel2, r02, r22); + auto n23 = where(sel2, r03, r23); + auto n30 = where(sel3, r00, r30); + auto n31 = where(sel3, r01, r31); + auto n32 = where(sel3, r02, r32); + auto n33 = where(sel3, r03, r33); + r00 = n00; r01 = n01; r02 = n02; r03 = n03; + r10 = n10; r11 = n11; r12 = n12; r13 = n13; + r20 = n20; r21 = n21; r22 = n22; r23 = n23; + r30 = n30; r31 = n31; r32 = n32; r33 = n33; + + sign = where(any_swap, -sign, sign); + degenerate = degenerate | (best_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r00); + auto f10 = r10 / pivot_safe; + auto f20 = r20 / pivot_safe; + auto f30 = r30 / pivot_safe; + r10 = f10; r20 = f20; r30 = f30; + r11 = r11 - f10 * r01; r12 = r12 - f10 * r02; r13 = r13 - f10 * r03; + r21 = r21 - f20 * r01; r22 = r22 - f20 * r02; r23 = r23 - f20 * r03; + r31 = r31 - f30 * r01; r32 = r32 - f30 * r02; r33 = r33 - f30 * r03; + } + + // ---- Step 1: pivot from column 1, rows {1, 2, 3} ---- + { + auto m1 = fabs(r11), m2 = fabs(r21), m3 = fabs(r31); + auto best = IVal(1); + auto best_mag = m1; + auto pick2 = (m2 > best_mag); + best = where(pick2, IVal(2), best); + best_mag = where(pick2, m2, best_mag); + auto pick3 = (m3 > best_mag); + best = where(pick3, IVal(3), best); + best_mag = where(pick3, m3, best_mag); + + auto sel2 = (best == IVal(2)); + auto sel3 = (best == IVal(3)); + auto any_swap = sel2 | sel3; + + auto n11 = where(sel2, r21, where(sel3, r31, r11)); + auto n12 = where(sel2, r22, where(sel3, r32, r12)); + auto n13 = where(sel2, r23, where(sel3, r33, r13)); + auto n21 = where(sel2, r11, r21); + auto n22 = where(sel2, r12, r22); + auto n23 = where(sel2, r13, r23); + auto n31 = where(sel3, r11, r31); + auto n32 = where(sel3, r12, r32); + auto n33 = where(sel3, r13, r33); + auto n10 = where(sel2, r20, where(sel3, r30, r10)); + auto n20 = where(sel2, r10, r20); + auto n30 = where(sel3, r10, r30); + r10 = n10; r11 = n11; r12 = n12; r13 = n13; + r20 = n20; r21 = n21; r22 = n22; r23 = n23; + r30 = n30; r31 = n31; r32 = n32; r33 = n33; + + sign = where(any_swap, -sign, sign); + degenerate = degenerate | (best_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r11); + auto f21 = r21 / pivot_safe; + auto f31 = r31 / pivot_safe; + r21 = f21; r31 = f31; + r22 = r22 - f21 * r12; r23 = r23 - f21 * r13; + r32 = r32 - f31 * r12; r33 = r33 - f31 * r13; + } + + // ---- Step 2: pivot from column 2, rows {2, 3} ---- + { + auto m2 = fabs(r22), m3 = fabs(r32); + auto sel3 = (m3 > m2); + auto pivot_mag = where(sel3, m3, m2); + + auto n22 = where(sel3, r32, r22); + auto n23 = where(sel3, r33, r23); + auto n32 = where(sel3, r22, r32); + auto n33 = where(sel3, r23, r33); + r22 = n22; r23 = n23; r32 = n32; r33 = n33; + auto n20 = where(sel3, r30, r20); + auto n21 = where(sel3, r31, r21); + auto n30 = where(sel3, r20, r30); + auto n31 = where(sel3, r21, r31); + r20 = n20; r21 = n21; r30 = n30; r31 = n31; + + sign = where(sel3, -sign, sign); + degenerate = degenerate | (pivot_mag < tol); + + auto pivot_safe = where(degenerate, FVal(1.0), r22); + auto f32 = r32 / pivot_safe; + r32 = f32; + r33 = r33 - f32 * r23; + } + + // ---- Final pivot check ---- + degenerate = degenerate | (fabs(r33) < tol); + + auto lu_det = sign * r00 * r11 * r22 * r33; + return where(degenerate, polynomial_fallback, lu_det); +} + +} // namespace kernels +} // namespace madspace diff --git a/madspace/src/kernels/threeparticle.hpp b/madspace/src/kernels/threeparticle.hpp index fc105f24e..e6f8af062 100644 --- a/madspace/src/kernels/threeparticle.hpp +++ b/madspace/src/kernels/threeparticle.hpp @@ -1,6 +1,7 @@ #pragma once #include "kinematics.hpp" +#include "lup_det.hpp" namespace madspace { namespace kernels { @@ -183,9 +184,8 @@ KERNELSPEC FVal bk_V( FVal t2, FVal s12 ) { - // Determinant of the 3x3 V-matrix, - // see Eq.(11) in 10.1103/PhysRev.187.2008. - // Note: expects the absolute value of t1 + // Matrix from Byckling-Kajantie eq.(11), 10.1103/PhysRev.187.2008. + // Asymmetric (a22 = a33 = 0; a12 = a21 but a13 != a31, a23 != a32). auto a11 = 2.0 * s12; auto a12 = ma_2 + s12 - t2; auto a13 = s12 + m1_2 - m2_2; @@ -194,8 +194,18 @@ KERNELSPEC FVal bk_V( auto a31 = m0_2 + s12 - m3_2; auto a32 = m0_2 + ma_2 - mb_2; - // Computes the determinant of the 3x3 V-matrix (hard-coded because easier) - auto det = a12 * a23 * a31 + a12 * a13 * a32 - a11 * a23 * a32 - a13 * a22 * a31; + // Polynomial fallback (the previous implementation), used by the LU + // helper if the matrix's leading pivots are below the LU tolerance. + // Expansion along row 3, using a22 = a33 = 0 and a21 = a12. + auto poly_det = a31 * (a12 * a23 - a13 * a22) + + a32 * (a12 * a13 - a11 * a23); + + auto det = lup_det3_general( + a11, a12, a13, + a12, a22, a23, + a31, a32, FVal(0.0), + poly_det + ); return -det / 8.0; } @@ -212,12 +222,10 @@ KERNELSPEC FVal bk_gram4( FVal s12, FVal s23 ) { - // omputes the 4x4 Gram determinant, - // see Eq.(B6) in 10.1103/PhysRev.187.2008. - // Note: expects the absolute value of t1 + // 4x4 Gram determinant, see Eq.(B6) in 10.1103/PhysRev.187.2008. + // Note: expects the absolute value of t1. - // Get upper triangular matrix components which are non-zero - // as the Gram matrix is symmetric, i.e. (a_{ij} = a_{ji}) + // Upper-triangular entries of the symmetric Gram matrix. auto a11 = 2.0 * ma_2; auto a12 = ma_2 - t1_abs - m1_2; auto a13 = ma_2 + t2 - s12; @@ -229,14 +237,28 @@ KERNELSPEC FVal bk_gram4( auto a34 = t2 + mb_2 - m3_2; auto a44 = 2.0 * mb_2; - // Computes the determinant of the 4x4 Gram matrix (hard-coded because easier) - auto det = a14 * a23 * a14 * a23 + a13 * a24 * a13 * a24 + a12 * a34 * a12 * a34 - - a14 * a14 * a22 * a33 - a13 * a13 * a22 * a44 - a12 * a12 * a33 * a44 - - a23 * a23 * a11 * a44 - a24 * a24 * a11 * a33 - a34 * a34 * a11 * a22 + - 2 * a11 * a23 * a24 * a34 + 2 * a12 * a13 * a23 * a44 + - 2 * a12 * a14 * a24 * a33 + 2 * a13 * a14 * a22 * a34 - - 2 * a12 * a13 * a24 * a34 - 2 * a12 * a14 * a23 * a34 - - 2 * a13 * a14 * a23 * a24 + a11 * a22 * a33 * a44; + // Polynomial fallback (the previous "hard-coded because easier" expansion), + // re-organized to keep the inner 2x2 minors as their own subexpressions. + // Used by the LU helper if any pivot is below tolerance. + auto poly_det = a14 * a14 * (a23 * a23 - a22 * a33) + + a13 * a13 * (a24 * a24 - a22 * a44) + + a12 * a12 * (a34 * a34 - a33 * a44) + - a23 * a23 * a11 * a44 + - a24 * a24 * a11 * a33 + - a34 * a34 * a11 * a22 + + a11 * a22 * a33 * a44 + + 2.0 * a11 * a23 * a24 * a34 + + 2.0 * a12 * a13 * (a23 * a44 - a24 * a34) + + 2.0 * a12 * a14 * (a24 * a33 - a23 * a34) + + 2.0 * a13 * a14 * (a22 * a34 - a23 * a24); + + auto det = lup_det4( + a11, a12, a13, a14, + a22, a23, a24, + a33, a34, + a44, + poly_det + ); return det / 16.0; } @@ -252,9 +274,6 @@ KERNELSPEC FVal bk_sqrt_g3i_g3im1( FVal t2, FVal s12 ) { - // This is the squaet root of the product of the two 3x3 Gram determinants g3i and - // g3im1, as in Eq.(11) in 10.1103/PhysRev.187.2008. Note: expects the absolute - // value of t1 auto a11 = 2 * s12; auto a12 = s12 + ma_2 - t2; auto a13 = s12 + m0_2 - m3_2; @@ -265,12 +284,28 @@ KERNELSPEC FVal bk_sqrt_g3i_g3im1( auto a33 = 2 * m0_2; auto b33 = 2 * m1_2; - // Calculate the two gramm determinants g3i and g3im1 - // (ard-coded because easier) - auto g3i = a11 * a22 * a33 + 2 * a12 * a23 * a13 - a11 * a23 * a23 - - a22 * a13 * a13 - a33 * a12 * a12; - auto g3im1 = a11 * a22 * b33 + 2 * a12 * b23 * b13 - a11 * b23 * b23 - - a22 * b13 * b13 - b33 * a12 * a12; + // Polynomial fallback for each 3x3 Gram (cofactor expansion). Used by the + // LU helper independently per matrix if its pivots are below tolerance. + auto poly_g3i = a11 * (a22 * a33 - a23 * a23) + - a12 * (a12 * a33 - a13 * a23) + + a13 * (a12 * a23 - a13 * a22); + auto poly_g3im1 = a11 * (a22 * b33 - b23 * b23) + - a12 * (a12 * b33 - b13 * b23) + + b13 * (a12 * b23 - b13 * a22); + + auto g3i = lup_det3( + a11, a12, a13, + a22, a23, + a33, + poly_g3i + ); + auto g3im1 = lup_det3( + a11, a12, b13, + a22, b23, + b33, + poly_g3im1 + ); + return sqrt(g3i * g3im1) / 8.0; } diff --git a/madspace/src/phasespace/t_propagator_mapping_23.cpp b/madspace/src/phasespace/color_ordered_mapping.cpp similarity index 67% rename from madspace/src/phasespace/t_propagator_mapping_23.cpp rename to madspace/src/phasespace/color_ordered_mapping.cpp index d539e3cc3..58d2c0579 100644 --- a/madspace/src/phasespace/t_propagator_mapping_23.cpp +++ b/madspace/src/phasespace/color_ordered_mapping.cpp @@ -1,4 +1,4 @@ -#include "madspace/phasespace/t_propagator_mapping_23.hpp" +#include "madspace/phasespace/color_ordered_mapping.hpp" #include #include @@ -12,19 +12,15 @@ namespace { // Cyclically rotate the colour order so that particle 0 is first, then // split the rest into two sets: particles strictly between 0 and 1 in the // cyclic order go into set1, particles after 1 (wrapping around) go into -// set2. Both sets are returned as 0-indexed outgoing-particle indices, -// i.e. particle p (0-indexed external) becomes (p - 2) in our 0-indexed -// convention for outgoing particles. (Particles 0 and 1 are the beams -// and not in either set.) +// set2. std::pair, std::vector> split_sets_from_colour_order(const std::vector& colour_order) { std::size_t n = colour_order.size(); if (n < 4) { throw std::invalid_argument( - "TPropagatorMapping23 requires at least 4 particles (2 beams + 2 outgoing)" + "ColorOrderedMapping requires at least 4 particles (2 beams + 2 outgoing)" ); } - // Find index of particle 0 and rotate auto it = std::find(colour_order.begin(), colour_order.end(), 0u); if (it == colour_order.end()) { throw std::invalid_argument("colour_order must contain particle 0"); @@ -35,18 +31,16 @@ split_sets_from_colour_order(const std::vector& colour_order) { for (std::size_t k = 0; k < n; ++k) { rotated.push_back(colour_order[(i0 + k) % n]); } - // rotated[0] is now particle 0. Find particle 1. auto it1 = std::find(rotated.begin(), rotated.end(), 1u); if (it1 == rotated.end()) { throw std::invalid_argument("colour_order must contain particle 1"); } std::size_t i1 = std::distance(rotated.begin(), it1); - // set1 = rotated[1..i1-1], set2 = rotated[i1+1..n-1] std::vector set1, set2; for (std::size_t k = 1; k < i1; ++k) { - std::size_t p = rotated[k]; // 0-indexed external + std::size_t p = rotated[k]; if (p <= 1) throw std::invalid_argument("invalid colour_order"); - set1.push_back(p - 2); // 0-indexed outgoing + set1.push_back(p - 2); } for (std::size_t k = i1 + 1; k < n; ++k) { std::size_t p = rotated[k]; @@ -55,21 +49,17 @@ split_sets_from_colour_order(const std::vector& colour_order) { } if (set1.empty() || set2.empty()) { throw std::invalid_argument( - "TPropagatorMapping23 requires both sets to be non-empty " + "ColorOrderedMapping requires both sets to be non-empty " "(particles 0 and 1 must not be adjacent in the colour order)" ); } return {set1, set2}; } -// Number of intermediate rest masses to sample for a walk of size k: -// max(0, k - 2) std::size_t n_intermediate_masses_for_set_size(std::size_t k) { return (k >= 2) ? (k - 2) : 0; } -// Number of randoms a walk of size k consumes for its t-channel-block calls -// (excluding mass samples). 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). @@ -78,21 +68,18 @@ std::size_t n_block_randoms_for_set_size(std::size_t k) { } // namespace -TPropagatorMapping23::TPropagatorMapping23( +ColorOrderedMapping::ColorOrderedMapping( const std::vector& colour_order, double t_invariant_power, double s_invariant_power ) : Mapping( - "TPropagatorMapping23", + "ColorOrderedMapping", [&] { auto [s1, s2] = split_sets_from_colour_order(colour_order); bool use_double_t = (s1.size() == 1) != (s2.size() == 1); - // When DoubleT is used: neither set-mass is sampled (the - // single-particle side has fixed mass, the multi-particle side's - // mass is derived by DoubleT), and the central block consumes - // 3 randoms (r_phi, r_t1, r_t2). Otherwise: 1 set-mass per - // multi-particle side, plus 2 randoms for the central 2->2. + // DoubleT branch: 0 set-masses + 3 central randoms. + // Standard branch: 1 set-mass per multi-particle side + 2 central randoms. std::size_t n_set_masses = use_double_t ? 0u : (s1.size() >= 2 ? 1u : 0u) + (s2.size() >= 2 ? 1u : 0u); @@ -150,23 +137,7 @@ TPropagatorMapping23::TPropagatorMapping23( _random_dim = n_set_masses + n_intermediate_masses + n_central + n_walk; } -// === Helpers used in both build_forward and build_inverse === -namespace { - -// Sum of masses of particles at indices `idxs` (0-indexed outgoing) in m_out. -Value sum_of_masses( - FunctionBuilder& fb, const ValueVec& m_out, const std::vector& idxs -) { - Value s = m_out.at(idxs[0]); - for (std::size_t k = 1; k < idxs.size(); ++k) { - s = fb.add(s, m_out.at(idxs[k])); - } - return s; -} - -} // namespace - -Mapping::Result TPropagatorMapping23::build_forward_impl( +Mapping::Result ColorOrderedMapping::build_forward_impl( FunctionBuilder& fb, const NamedVector& inputs, const NamedVector& conditions @@ -177,23 +148,19 @@ Mapping::Result TPropagatorMapping23::build_forward_impl( auto next_random = [&]() { return *(r++); }; ValueVec dets; - // ============================================================ // Phase 1a: pre-sample set composite masses. - // When _use_double_t: neither set mass is sampled here. The single- - // particle side has fixed mass; the multi-particle side's mass is - // derived later from the DoubleT central block. - // Otherwise: sample multi-particle set masses, fix single-particle ones. - // ============================================================ + // DoubleT branch defers the multi-particle side's mass until after Phase 2. Value m_set1, m_set2; - Value mass_sum_set1 = sum_of_masses(fb, m_out, _set1); - Value mass_sum_set2 = sum_of_masses(fb, m_out, _set2); + auto masses_of = [&](const std::vector& idxs) { + ValueVec v; v.reserve(idxs.size()); + for (auto i : idxs) v.push_back(m_out.at(i)); + return fb.sum(v); + }; + Value mass_sum_set1 = masses_of(_set1); + Value mass_sum_set2 = masses_of(_set2); if (_use_double_t) { - // Both set masses set later (multi-side from DoubleT output, - // single-side fixed). if (_set1.size() == 1) m_set1 = m_out.at(_set1[0]); if (_set2.size() == 1) m_set2 = m_out.at(_set2[0]); - // The multi-side m_set will be filled in after the DoubleT call, - // before sampling intermediate rest masses. } else { if (_set1.size() >= 2) { auto s_min = fb.square(mass_sum_set1); @@ -215,9 +182,7 @@ Mapping::Result TPropagatorMapping23::build_forward_impl( } } - // ============================================================ // Phase 2: central block producing (P_set1, P_set2). - // ============================================================ auto [pa, pb] = fb.com_p_in(e_cm); Value P_set1, P_set2; if (_use_double_t) { @@ -232,15 +197,13 @@ Mapping::Result TPropagatorMapping23::build_forward_impl( {next_random(), next_random(), next_random()}, {pa, pb, m_single, mir_min} ); - Value p_single = central.at(0); // mass = m_single - Value p_recoil = central.at(1); // recoil; mass derived from t1, t2 + Value p_single = central.at(0); + Value p_recoil = central.at(1); dets.push_back(central["det"]); if (single_is_set1) { P_set1 = p_single; P_set2 = p_recoil; - // m_set1 already set above. Derive m_set2 from p_recoil for - // downstream intermediate-mass sampling. invariants_from_momenta - // returns m^2 for the summed momenta indicated by the factor row. + // Derive multi-side m_set from the recoil momentum for Phase 1b. nested_vector2 factors_recoil{{1.0}}; auto m2 = fb.unstack( fb.invariants_from_momenta(fb.stack({p_recoil}), factors_recoil) @@ -266,44 +229,24 @@ Mapping::Result TPropagatorMapping23::build_forward_impl( dets.push_back(central["det"]); } - // The "other beam after central block" for each set's walk: - // by 4-momentum conservation, p_other_beam - P_other_set is what is - // (kinematically) left on that beam line after the central block emitted - // the other set. + // R_b for each walk: the beam minus what the central block emitted on the other side. Value R_b_for_set1 = fb.sub(pb, P_set2); Value R_b_for_set2 = fb.sub(pa, P_set1); - // ============================================================ // Phase 1b: pre-sample intermediate rest masses for each walk. - // For a walk peeling particles set[0], set[1], ..., set[k-1]: - // - mass of rest after peeling set[0] = m_rest_1 (sampled) - // - mass of rest after peeling set[0..1] = m_rest_2 (sampled) - // - ... - // - mass of rest after peeling set[0..k-3] = m_rest_{k-2} (sampled) - // - mass of rest after peeling set[0..k-2] = m of set[k-1] (fixed, residual) - // Bounds: - // m_min = sum of masses of [set[j+1], ..., set[k-1]] - // m_max = m_set - sum of masses of [set[0], ..., set[j]] - // Note: we sample these AFTER the central block, because for DoubleT - // we don't know m_set (multi-side) until then. - // ============================================================ + // For each step j in 0..k-3 (last is the residual mass(s[k-1])): + // m_min = sum masses [set[j+1]..set[k-1]] + // m_max = m_rest_j - mass[set[j]], m_rest_0 = m_set + // Sampling from [m_min, m_set - sum_peeled] independently would + // over-count by (k-2)! by ignoring the monotonic ordering. + // Sampled after the central block: DoubleT only fixes m_set then. auto sample_intermediate_masses = [&](const std::vector& s, Value m_set) -> ValueVec { std::size_t k = s.size(); ValueVec res; if (k <= 2) { - return res; // no intermediate masses needed + return res; } - // The walk peels particles s[0], s[1], ..., s[k-2] (s[k-1] is the - // final residual). After peeling s[0..j], the "rest" has mass - // m_rest_{j+1}. The chain is monotonic: - // m_set > m_rest_1 > m_rest_2 > ... > m_rest_{k-2} > mass(s[k-1]) - // We sample m_rest_{j+1} from - // m_min = sum of masses [s[j+1]..s[k-1]] (kinematic lower bound) - // m_max = m_rest_j - mass[s[j]] (monotonicity + peel-off room) - // where m_rest_0 = m_set. - // Sampling independently from [m_min, m_set - sum_of_already_peeled] - // would over-count by (k-2)! because it ignores the ordering. Value prev_mass = m_set; for (std::size_t j = 0; j < k - 2; ++j) { Value m_min = m_out.at(s[j + 1]); @@ -324,10 +267,8 @@ Mapping::Result TPropagatorMapping23::build_forward_impl( ValueVec rest_masses_set1 = sample_intermediate_masses(_set1, m_set1); ValueVec rest_masses_set2 = sample_intermediate_masses(_set2, m_set2); - // ============================================================ // Phase 3: peel-off walks - // ============================================================ - ValueVec p_out(_n_out); // 0-indexed by outgoing particle index + ValueVec p_out(_n_out); auto walk = [&](const std::vector& s, Value P_set, @@ -342,25 +283,16 @@ Mapping::Result TPropagatorMapping23::build_forward_impl( Value im1; bool first = true; for (std::size_t j = 0; j < k - 1; ++j) { - // Peel particle s[j]. The "new rest" after this peel has either - // mass = rest_masses[j] (intermediate) or mass = m_out[s[k-1]] - // (the final residual, when j == k-2). + // 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]); - // Block input: (R_b, R_a). Block output: p1_out + p2_out = R_b + R_a, - // where p1_out has mass m_rest (the chain carrier, sits on the pa=R_b - // side per the block's convention) and p2_out has mass m_peel (the - // newly peeled particle, sits on the pb=R_a side -- which is what we - // want since R_a is the "active" leg that emits). - // R_a is updated as R_a -= peeled. R_b stays constant. + // 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. // - // For the 2->3 block, the kernel internally computes - // p_12 = pa + pb - p_3, where p_3 = im1 (previous peel on this side). - // We want p_12 = remaining-system-left-to-produce = R_a + R_b - // (using R_a *after* the previous peel was subtracted from it). - // Since the kernel subtracts im1 internally, and im1 is already - // subtracted from R_a in our bookkeeping, we restore it: pass - // pb = R_a + im1, so p_12 = R_b + (R_a + im1) - im1 = R_b + R_a. ✓ + // 2->3 kernel internally subtracts p_3 = im1 from pa+pb. R_a already + // 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) { auto ks = _lab_scattering.build_forward( fb, @@ -406,7 +338,7 @@ Mapping::Result TPropagatorMapping23::build_forward_impl( return {{output_types().keys(), p_ext}, fb.product(dets)}; } -Mapping::Result TPropagatorMapping23::build_inverse_impl( +Mapping::Result ColorOrderedMapping::build_inverse_impl( FunctionBuilder& fb, const NamedVector& inputs, const NamedVector& conditions @@ -416,22 +348,13 @@ Mapping::Result TPropagatorMapping23::build_inverse_impl( ValueVec random_out; ValueVec dets; - // Inputs: pa, pb, p_out[0], p_out[1], ..., p_out[n_out-1] Value pa = inputs.at(0); Value pb = inputs.at(1); auto p_outgoing = [&](std::size_t i) { return inputs.at(2 + i); }; - // ============================================================ - // Build a factor matrix to extract all needed invariant masses in - // one shot. Rows (in order): - // row 0: m^2(P_set1) (if set1.size() >= 2) - // row 1: m^2(P_set2) (if set2.size() >= 2) - // then for set1 if size >= 3: - // k_1 - 2 rows, each summing momenta s1[j+1..k_1-1] for j=0..k_1-3 - // then for set2 if size >= 3: similarly - // Column = index into inputs (n_out + 2 columns). - // We record the layout so we can pull invariants out by name afterwards. - // ============================================================ + // Factor matrix: extract all needed invariant masses in one fb call. + // Layout: m^2(P_set1), m^2(P_set2), then per-set rest-system masses. + // idx_* below record where each band starts. nested_vector2 invariant_factors; auto n_inputs = _n_out + 2; auto factor_set = [&](const std::vector& idxs) { @@ -439,7 +362,6 @@ Mapping::Result TPropagatorMapping23::build_inverse_impl( for (auto idx : idxs) row.at(idx + 2) = 1.0; return row; }; - // Track positions in the stacked invariants output: int idx_m2_set1 = -1, idx_m2_set2 = -1; int idx_rest_set1_start = -1, idx_rest_set2_start = -1; if (_set1.size() >= 2) { @@ -471,9 +393,6 @@ Mapping::Result TPropagatorMapping23::build_inverse_impl( ); } - // ============================================================ - // Recover composites P_set1, P_set2 from the momenta. - // ============================================================ auto sum_momenta = [&](const std::vector& s) -> Value { Value p = p_outgoing(s[0]); for (std::size_t k = 1; k < s.size(); ++k) { @@ -484,17 +403,17 @@ Mapping::Result TPropagatorMapping23::build_inverse_impl( Value P_set1 = sum_momenta(_set1); Value P_set2 = sum_momenta(_set2); - // ============================================================ - // Phase 1a inverse: recover set-mass randoms (skipped in DoubleT case). - // Even when randoms are skipped, we still need m_set1/m_set2 for the - // Phase 1b intermediate-mass-bound calculation; those come from the - // invariants_from_momenta result for multi-particle sides. - // ============================================================ + // Phase 1a inverse: recover set-mass randoms (none for DoubleT). + // m_set values come from invariants_from_momenta for multi-particle sides. Value m_set1, m_set2; - Value mass_sum_set1 = sum_of_masses(fb, m_out, _set1); - Value mass_sum_set2 = sum_of_masses(fb, m_out, _set2); + auto masses_of = [&](const std::vector& idxs) { + ValueVec v; v.reserve(idxs.size()); + for (auto i : idxs) v.push_back(m_out.at(i)); + return fb.sum(v); + }; + Value mass_sum_set1 = masses_of(_set1); + Value mass_sum_set2 = masses_of(_set2); if (_use_double_t) { - // No set-mass randoms to recover; just set m_set1/m_set2. if (_set1.size() == 1) { m_set1 = m_out.at(_set1[0]); m_set2 = fb.sqrt(invariants.at(idx_m2_set2)); @@ -527,11 +446,8 @@ Mapping::Result TPropagatorMapping23::build_inverse_impl( } } - // ============================================================ - // Phase 2 inverse: central block from (pa, pb). Comes BEFORE Phase 1b - // because the forward emits central randoms before intermediate-mass - // randoms (so the multi-side m_set is known by Phase 1b). - // ============================================================ + // Phase 2 inverse: central block. Comes before Phase 1b because the + // forward emits central randoms first (so multi-side m_set is known by 1b). if (_use_double_t) { bool single_is_set1 = (_set1.size() == 1); Value p_single = single_is_set1 ? P_set1 : P_set2; @@ -543,8 +459,6 @@ Mapping::Result TPropagatorMapping23::build_inverse_impl( {p_single, p_recoil}, {pa, pb, m_single, mir_min} ); - // DoubleT inverse returns (r_phi, r_t1, r_t2, det). m1 is no longer - // an output (it's a condition now). random_out.push_back(central.at(0)); random_out.push_back(central.at(1)); random_out.push_back(central.at(2)); @@ -560,9 +474,7 @@ Mapping::Result TPropagatorMapping23::build_inverse_impl( dets.push_back(central["det"]); } - // ============================================================ // Phase 1b inverse: recover intermediate rest-mass randoms - // ============================================================ auto recover_intermediate_masses = [&](const std::vector& s, Value m_set, int idx_start) { std::size_t k = s.size(); @@ -591,9 +503,7 @@ Mapping::Result TPropagatorMapping23::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) { @@ -604,9 +514,7 @@ Mapping::Result TPropagatorMapping23::build_inverse_impl( bool first = true; for (std::size_t j = 0; j < k - 1; ++j) { Value peeled = p_outgoing(s[j]); - // The block at this rung had inputs (R_a, R_b) and produced - // (p1_out, p2_out) where p1_out has mass m_rest (carrier of - // remaining) and p2_out = peeled. By conservation, + // 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) { diff --git a/madspace/src/python/instruction_set.hpp b/madspace/src/python/instruction_set.hpp index 10dc81741..5f3509a19 100644 --- a/madspace/src/python/instruction_set.hpp +++ b/madspace/src/python/instruction_set.hpp @@ -34,8 +34,8 @@ void add_instructions(py::classh& fb) { fb.def("div", &FunctionBuilder::div, py::arg("in1"), py::arg("in2")); fb.def("reduce_sum", &FunctionBuilder::reduce_sum, py::arg("in1")); fb.def("reduce_sum_vector", &FunctionBuilder::reduce_sum_vector, py::arg("in1")); - fb.def("batch_reduce_mean", &FunctionBuilder::batch_reduce_mean, py::arg("in")); - fb.def("batch_reduce_mean_keepdim", &FunctionBuilder::batch_reduce_mean_keepdim, py::arg("in")); + fb.def("batch_reduce_mean", &FunctionBuilder::batch_reduce_mean, py::arg("in1")); + fb.def("batch_reduce_mean_keepdim", &FunctionBuilder::batch_reduce_mean_keepdim, py::arg("in1")); fb.def("reduce_product", &FunctionBuilder::reduce_product, py::arg("in1")); fb.def("sqrt", &FunctionBuilder::sqrt, py::arg("in1")); fb.def("square", &FunctionBuilder::square, py::arg("in1")); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index a58736f48..49323941c 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -488,14 +488,14 @@ PYBIND11_MODULE(_madspace_py, m) { ) .def("random_dim", &TPropagatorMapping::random_dim); - py::classh(m, "TPropagatorMapping23") + py::classh(m, "ColorOrderedMapping") .def( py::init, double, double>(), py::arg("colour_order"), - py::arg("t_invariant_power") = 0.8, - py::arg("s_invariant_power") = 0.8 + py::arg("t_invariant_power") = 0., + py::arg("s_invariant_power") = 0. ) - .def("random_dim", &TPropagatorMapping23::random_dim); + .def("random_dim", &ColorOrderedMapping::random_dim); py::classh(m, "VegasHistogram") .def( @@ -1432,81 +1432,6 @@ PYBIND11_MODULE(_madspace_py, m) { .def_readonly("bin_values", &Histogram::bin_values) .def_readonly("bin_errors", &Histogram::bin_errors); - py::classh(m, "ChannelEventGenerator") - .def_static( - "load", - &ChannelEventGenerator::load, - py::arg("channel_file"), - py::arg("contexts"), - py::arg("event_file"), - py::arg("weight_file"), - py::arg("config") - ) - .def( - py::init< - const std::vector&, - const Integrand&, - const std::string&, - const std::string&, - const GeneratorConfig&, - std::size_t, - const std::string&, - const std::optional&>(), - py::arg("contexts"), - py::arg("integrand"), - py::arg("event_file"), - py::arg("weight_file"), - py::arg("config"), - py::arg("subprocess_index"), - py::arg("name"), - py::arg("histograms") - ) - .def_readonly_static("integrand_flags", &ChannelEventGenerator::integrand_flags) - .def("status", &ChannelEventGenerator::status) - .def("save", &ChannelEventGenerator::save, py::arg("save")); - - py::classh(m, "EventGenerator") - .def_readonly_static("default_config", &EventGenerator::default_config) - .def( - py::init< - const std::vector&, - const std::vector>&, - const std::string&, - const GeneratorConfig&>(), - py::arg("contexts"), - py::arg("channels"), - py::arg("status_file") = "", - py::arg_v( - "config", - EventGenerator::default_config, - "EventGenerator.default_config" - ) - ) - .def("survey", &EventGenerator::survey) - .def("generate", &EventGenerator::generate) - .def( - "combine_to_compact_npy", - &EventGenerator::combine_to_compact_npy, - py::arg("file_name") - ) - .def( - "combine_to_lhe_npy", - &EventGenerator::combine_to_lhe_npy, - py::arg("file_name"), - py::arg("lhe_completer") - ) - .def( - "combine_to_lhe", - &EventGenerator::combine_to_lhe, - py::arg("file_name"), - py::arg("lhe_completer") - ) - .def("status", &EventGenerator::status) - .def("channel_status", &EventGenerator::channel_status) - .def("histograms", &EventGenerator::histograms) - .def("used_globals", &EventGenerator::used_globals) - .def("channels", &EventGenerator::channels); - py::classh(m, "LHEHeader") .def( py::init(), @@ -1726,6 +1651,81 @@ PYBIND11_MODULE(_madspace_py, m) { .def("print_update", &PrettyBox::print_update) .def_property_readonly("line_count", &PrettyBox::line_count); + py::classh(m, "ChannelEventGenerator") + .def_static( + "load", + &ChannelEventGenerator::load, + py::arg("channel_file"), + py::arg("contexts"), + py::arg("event_file"), + py::arg("weight_file"), + py::arg("config") + ) + .def( + py::init< + const std::vector&, + const Integrand&, + const std::string&, + const std::string&, + const GeneratorConfig&, + std::size_t, + const std::string&, + const std::optional&>(), + py::arg("contexts"), + py::arg("integrand"), + py::arg("event_file"), + py::arg("weight_file"), + py::arg("config"), + py::arg("subprocess_index"), + py::arg("name"), + py::arg("histograms") + ) + .def_readonly_static("integrand_flags", &ChannelEventGenerator::integrand_flags) + .def("status", &ChannelEventGenerator::status) + .def("save", &ChannelEventGenerator::save, py::arg("save")); + + py::classh(m, "EventGenerator") + .def_readonly_static("default_config", &EventGenerator::default_config) + .def( + py::init< + const std::vector&, + const std::vector>&, + const std::string&, + const GeneratorConfig&>(), + py::arg("contexts"), + py::arg("channels"), + py::arg("status_file") = "", + py::arg_v( + "config", + EventGenerator::default_config, + "EventGenerator.default_config" + ) + ) + .def("survey", &EventGenerator::survey) + .def("generate", &EventGenerator::generate) + .def( + "combine_to_compact_npy", + &EventGenerator::combine_to_compact_npy, + py::arg("file_name") + ) + .def( + "combine_to_lhe_npy", + &EventGenerator::combine_to_lhe_npy, + py::arg("file_name"), + py::arg("lhe_completer") + ) + .def( + "combine_to_lhe", + &EventGenerator::combine_to_lhe, + py::arg("file_name"), + py::arg("lhe_completer") + ) + .def("status", &EventGenerator::status) + .def("channel_status", &EventGenerator::channel_status) + .def("histograms", &EventGenerator::histograms) + .def("used_globals", &EventGenerator::used_globals) + .def("channels", &EventGenerator::channels); + py::classh logger(m, "Logger"); add_enum( logger, diff --git a/madspace/tests/test_2to3_scattering.py b/madspace/tests/test_2to3_scattering.py index ba96e8909..dfa52bed6 100644 --- a/madspace/tests/test_2to3_scattering.py +++ b/madspace/tests/test_2to3_scattering.py @@ -40,8 +40,7 @@ def boost(k: np.ndarray, p_boost: np.ndarray, inverse: bool = False) -> np.ndarr # Fixtures # ---------------------------- -# N = 50_000 # keep this moderate for CI speed; bump locally for tighter stats -N = 100_000 # keep this moderate for CI speed; bump locally for tighter stats +N = 10_000 # keep this moderate for CI speed; bump locally for tighter stats InputPoint = namedtuple( "InputPoint", diff --git a/madspace/tests/test_color_ordered_mapping.py b/madspace/tests/test_color_ordered_mapping.py new file mode 100644 index 000000000..4b86af1eb --- /dev/null +++ b/madspace/tests/test_color_ordered_mapping.py @@ -0,0 +1,209 @@ +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: (colour_order, int_order_for_old_class, topology_label). +TOPOLOGIES = [ + ([0, 2, 3, 1, 4], [0, 1], "n=5: set1={2,3}, set2={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, 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}"), +] + + +def _expand(topologies): + """Flat list of (colour_order, int_order, masses, label) pairs.""" + out = [] + for colour_order, int_order, topo_label in topologies: + n_out = len(colour_order) - 2 + for masses, mass_label in _mass_sets_for(n_out): + label = f"{mass_label}-{topo_label}" + out.append((colour_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("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_momentum_conservation(rng, colour_order, _int_order, masses, _label): + mapping = ms.ColorOrderedMapping(colour_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("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_on_shell_masses(rng, colour_order, _int_order, masses, _label): + mapping = ms.ColorOrderedMapping(colour_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("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) +def test_inverse(rng, colour_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(colour_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("colour_order,int_order,masses,_label", CASES, ids=CASE_IDS) +def test_phase_space_volume_matches_old(rng, colour_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(colour_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)" + ) \ No newline at end of file diff --git a/madspace/tests/test_t_propagator_23.py b/madspace/tests/test_t_propagator_23.py deleted file mode 100644 index 07d1f1934..000000000 --- a/madspace/tests/test_t_propagator_23.py +++ /dev/null @@ -1,264 +0,0 @@ -"""Tests for TPropagatorMapping23 (per-set, colour-ordered topology). - -Colour order convention: 0-indexed permutation of {0, ..., n-1}. -Particles 0 and 1 are the two incoming beams; 2, 3, ..., n-1 are outgoing. - -Note on non-physical events: For inner-rung 2->3 peels, the kinematic -constraint imposed by the 2->3 block (Gram-determinant condition, s23/t1 -range non-empty) can be violated for some sampled (rest-mass, t, s) combinations. -These events produce non-finite momenta/det and should be discarded in any -Monte Carlo integration -- they correspond to zero-weight points outside the -physical region. The tests therefore filter on `np.isfinite(det)` before -checking conservation, on-shell, and inverse round-trip. - -Note on numerical tails: events near kinematic boundaries (very small forward -det) have accumulated floating-point error from chained sqrts, boosts, and -rotations. A 10th-percentile det floor filters these without biasing the -integrand (they have near-zero weight in any MC integral anyway). -""" -import numpy as np -import pytest -from pytest import approx - -import madspace as ms - - -@pytest.fixture -def rng(): - return np.random.default_rng(1234) - - -N = 100_000 -CM_ENERGY = 13000.0 -MIN_PHYSICAL_FRACTION = 0.10 # at least 10% of events should be physical -LOW_DET_QUANTILE = 0.10 # drop bottom 10% by |det| in stability-sensitive tests - - -# Mass set families. The labels are templates; 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 defined for n_out >= 3. - return [173.0, 80.0] + [0.0] * (n_out - 2), f"{n_out} mixed" - - -def _mass_sets_for(n_out): - """Mass families that make sense at this n_out.""" - sets = [_massless(n_out), _w_like(n_out)] - if n_out >= 3: - sets.append(_mixed_top(n_out)) - return sets - - -# Each entry: (colour_order, int_order_for_old_class, topology_label). -# The topology determines n_out = len(colour_order) - 2. -TOPOLOGIES = [ - ([0, 2, 3, 1, 4], [0, 1], "n=5: set1={2,3}, set2={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, 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}"), -] - - -def _expand(topologies): - """Produce a flat list of (colour_order, int_order, masses, label) tuples - where each topology is paired only with its compatible mass families. - This replaces the cross-product parametrize that produced many skips.""" - out = [] - for colour_order, int_order, topo_label in topologies: - n_out = len(colour_order) - 2 - for masses, mass_label in _mass_sets_for(n_out): - label = f"{mass_label}-{topo_label}" - out.append((colour_order, int_order, masses, label)) - return out - - -CASES = _expand(TOPOLOGIES) -CASE_IDS = [c[3] for c in CASES] - - -def _physical_mask(det, p_ext): - """Events are physical iff det is finite AND every momentum is finite.""" - 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): - mask = _physical_mask(det, p_ext) - if mask.sum() < MIN_PHYSICAL_FRACTION * len(det): - pytest.fail( - f"Too few physical events: {mask.sum()}/{len(det)} " - f"({100*mask.sum()/len(det):.1f}%), expected >= " - f"{100*MIN_PHYSICAL_FRACTION:.0f}%" - ) - return mask, [p[mask] for p in p_ext], det[mask] - - -def _drop_low_det_tail(det, *arrays, q=LOW_DET_QUANTILE, median_relative=1e-6): - """Drop the bottom-q quantile AND any events below median_relative * median - of events by |det|. These boundary events accumulate floating-point error - from chained operations and have near-zero weight in MC anyway. The - median-relative floor matters for heavy-tailed DoubleT topologies where - the quantile floor alone is still many decades below the median.""" - median = np.median(np.abs(det)) - floor = max(np.quantile(np.abs(det), q), median_relative * median) - keep = np.abs(det) > floor - return (det[keep], *(a[keep] for a in arrays)) - - -@pytest.mark.parametrize("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_momentum_conservation(rng, colour_order, _int_order, masses, _label): - mapping = ms.TPropagatorMapping23(colour_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("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_on_shell_masses(rng, colour_order, _int_order, masses, _label): - mapping = ms.TPropagatorMapping23(colour_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) - - # Drop boundary events that accumulate FP error from chained kinematics. - det, *p_ext = _drop_low_det_tail(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("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_inverse(rng, colour_order, _int_order, masses, _label): - mapping = ms.TPropagatorMapping23(colour_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) - 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) - - # Drop the lowest-det events: boundary events whose forward det is many - # orders of magnitude below the median have inverse det near infinity, - # and the round-trip |det * inv_det - 1| can hit a few percent (or more) - # from accumulated FP error in atan2/sqrt/boosts. They have negligible - # weight in any MC integral. We use both a quantile floor (handles - # well-behaved distributions) and a median-relative floor (handles the - # heavy-tailed DoubleT-using topologies where the 10th percentile is - # still many decades below the median). - median = np.median(np.abs(det_phys)) - floor_quantile = np.quantile(np.abs(det_phys), LOW_DET_QUANTILE) - floor_median = 1e-6 * median - floor = max(floor_quantile, floor_median) - keep = ( - np.isfinite(det_inv) - & np.isfinite(det_phys) - & (np.abs(det_phys) > floor) - ) - for ri in r_inv: - keep &= np.isfinite(ri) - if keep.sum() < MIN_PHYSICAL_FRACTION * len(det): - pytest.fail( - f"Too few invertible events: {keep.sum()}/{len(det)} " - f"({100*keep.sum()/len(det):.1f}%)" - ) - - det_kept = det_phys[keep] - det_inv_kept = det_inv[keep] - # Inverse-roundtrip tolerance: 0.1% accounts for accumulated FP error - # in chained kinematic operations (boosts, rotations, sqrts) for the - # deeper-walk topologies. Correctness is established by - # 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); - # round-trip only recovers which-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("colour_order,int_order,masses,_label", CASES, ids=CASE_IDS) -def test_phase_space_volume_matches_old(rng, colour_order, int_order, masses, _label): - """The new class should integrate to the same phase-space volume as - the old TPropagatorMapping (with matrix element = 1). - - This is the actual correctness check: both classes are MC estimators of - the same n-body phase-space integral, and their averages must agree - within a few sigma. (The inverse-roundtrip test is a precision check; - this one is the physics check.) - """ - new = ms.TPropagatorMapping23(colour_order) - old = ms.TPropagatorMapping(int_order) - - r_new = rng.random((N, new.random_dim())) - r_old = rng.random((N, old.random_dim())) - e_cm = np.full(N, CM_ENERGY) - cond_new = [e_cm] + [np.full(N, m) for m in masses] - cond_old = cond_new - - *p_new, det_new = new.map_forward(list(r_new.T), cond_new) - *p_old, det_old = old.map_forward(list(r_old.T), cond_old) - - mask_new = np.isfinite(det_new) & (det_new > 0) - mask_old = np.isfinite(det_old) & (det_old > 0) - det_new = det_new[mask_new] - det_old = det_old[mask_old] - - assert len(det_new) > MIN_PHYSICAL_FRACTION * N, \ - f"new class: {len(det_new)}/{N} physical, too few" - assert len(det_old) > MIN_PHYSICAL_FRACTION * N, \ - f"old class: {len(det_old)}/{N} physical, too few" - - mean_new = np.mean(det_new) - mean_old = 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)" - ) \ No newline at end of file From 2a605bed8412056163e1ea98d057f95524ed3a7d Mon Sep 17 00:00:00 2001 From: Ramon Winterhalder Date: Tue, 19 May 2026 15:18:19 +0200 Subject: [PATCH 5/6] Go up with version --- madspace/pyproject.toml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/madspace/pyproject.toml b/madspace/pyproject.toml index 0c544e012..e271cdd3f 100644 --- a/madspace/pyproject.toml +++ b/madspace/pyproject.toml @@ -4,11 +4,12 @@ build-backend = "scikit_build_core.build" [project] name = "madspace" -version = "0.1.3" -description="Fast and flexible phase space mappings" +version = "0.2.0" +description="Fast and flexible phase-space mappings" readme = "README.md" authors = [ { name = "Theo Heimel", email = "theo.heimel@uclouvain.be" }, + { name = "Ramon Winterhalder", email = "ramon.winterhalder@unimi.it" }, ] requires-python = ">=3.11" classifiers = [ From 20b692dcfb95736aec6a2df2a89ade2f21a1c4f8 Mon Sep 17 00:00:00 2001 From: Ramon Winterhalder Date: Thu, 21 May 2026 15:43:10 +0200 Subject: [PATCH 6/6] Update spelling of color and add test dependencies to toml --- .../phasespace/color_ordered_mapping.hpp | 4 +- madspace/madspace/_madspace_py.pyi | 3978 ++++++++--------- madspace/pyproject.toml | 10 +- .../src/phasespace/color_ordered_mapping.cpp | 36 +- madspace/src/python/madspace.cpp | 2 +- madspace/tests/test_color_ordered_mapping.py | 34 +- 6 files changed, 1873 insertions(+), 2191 deletions(-) diff --git a/madspace/include/madspace/phasespace/color_ordered_mapping.hpp b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp index dedc5d6e6..cd1d70a9c 100644 --- a/madspace/include/madspace/phasespace/color_ordered_mapping.hpp +++ b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp @@ -12,10 +12,10 @@ namespace madspace { class ColorOrderedMapping : public Mapping { public: - // colour_order: 0-indexed permutation of {0, ..., n-1} (n = n_out + 2). + // 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& colour_order, + const std::vector& color_order, double t_invariant_power = 0.8, double s_invariant_power = 0.8 ); diff --git a/madspace/madspace/_madspace_py.pyi b/madspace/madspace/_madspace_py.pyi index 8b0251fe5..ce4e6ad39 100644 --- a/madspace/madspace/_madspace_py.pyi +++ b/madspace/madspace/_madspace_py.pyi @@ -1,1213 +1,1062 @@ from __future__ import annotations import collections.abc import typing - -__all__: list[str] = [ - "AdamOptimizer", - "AlphaSGrid", - "BatchSize", - "ChannelEventGenerator", - "ChannelWeightNetwork", - "ColorOrderedMapping", - "Context", - "CutItem", - "Cuts", - "DataType", - "Decay", - "Device", - "Diagram", - "DifferentialCrossSection", - "DiscreteFlow", - "DiscreteHistogram", - "DiscreteOptimizer", - "DiscreteSampler", - "DoubleT", - "EnergyScale", - "EventGenerator", - "FastRamboMapping", - "Flow", - "Function", - "FunctionBuilder", - "FunctionGenerator", - "FunctionRuntime", - "GeneratorConfig", - "GeneratorStatus", - "HistItem", - "Histogram", - "Instruction", - "InstructionCall", - "Integrand", - "IntegrandProbability", - "Invariant", - "LHECompleter", - "LHEEvent", - "LHEFileWriter", - "LHEHeader", - "LHEMeta", - "LHEParticle", - "LHEProcess", - "LineRef", - "Logger", - "Luminosity", - "MLP", - "MadnisConfig", - "MadnisLoss", - "MadnisTraining", - "Mapping", - "MatrixElement", - "MatrixElementApi", - "MomentumPreprocessing", - "MultiChannelFunction", - "MultiChannelIntegrand", - "MultiChannelMapping", - "MultiMadnisTraining", - "NamedTypes", - "NamedValues", - "Observable", - "ObservableHistograms", - "PartonDensity", - "PdfGrid", - "PhaseSpaceMapping", - "PrettyBox", - "Propagator", - "PropagatorChannelWeights", - "RunningCoupling", - "SubchannelWeights", - "SubprocArgs", - "TPropagatorMapping", - "Tensor", - "ThreeBodyDecay", - "Topology", - "TwoBodyDecay", - "TwoToThreeParticleScattering", - "TwoToTwoParticleScattering", - "Type", - "Unweighter", - "Value", - "VegasGridOptimizer", - "VegasHistogram", - "VegasMapping", - "Verbosity", - "batch_float", - "batch_float_array", - "batch_four_vec", - "batch_four_vec_array", - "batch_int", - "batch_size", - "batch_sizes", - "cpu_device", - "cuda_device", - "default_context", - "default_cuda_context", - "default_hip_context", - "float", - "format_progress", - "format_si_prefix", - "format_with_error", - "hip_device", - "initialize_vegas_grid", - "int", - "log", - "multichannel_batch_size", - "pretty", - "set_lib_path", - "set_simd_vector_size", - "silent", - "single_float", - "single_int", -] - +__all__: list[str] = ['AdamOptimizer', 'AlphaSGrid', 'BatchSize', 'ChannelEventGenerator', 'ChannelWeightNetwork', 'ColorOrderedMapping', 'Context', 'CutItem', 'Cuts', 'DataType', 'Decay', 'Device', 'Diagram', 'DifferentialCrossSection', 'DiscreteFlow', 'DiscreteHistogram', 'DiscreteOptimizer', 'DiscreteSampler', 'DoubleT', 'EnergyScale', 'EventGenerator', 'FastRamboMapping', 'Flow', 'Function', 'FunctionBuilder', 'FunctionGenerator', 'FunctionRuntime', 'GeneratorConfig', 'GeneratorStatus', 'HistItem', 'Histogram', 'Instruction', 'InstructionCall', 'Integrand', 'IntegrandProbability', 'Invariant', 'LHECompleter', 'LHEEvent', 'LHEFileWriter', 'LHEHeader', 'LHEMeta', 'LHEParticle', 'LHEProcess', 'LineRef', 'Logger', 'Luminosity', 'MLP', 'MadnisConfig', 'MadnisLoss', 'MadnisTraining', 'Mapping', 'MatrixElement', 'MatrixElementApi', 'MomentumPreprocessing', 'MultiChannelFunction', 'MultiChannelIntegrand', 'MultiChannelMapping', 'MultiMadnisTraining', 'NamedTypes', 'NamedValues', 'Observable', 'ObservableHistograms', 'PartonDensity', 'PdfGrid', 'PhaseSpaceMapping', 'PrettyBox', 'Propagator', 'PropagatorChannelWeights', 'RunningCoupling', 'SubchannelWeights', 'SubprocArgs', 'TPropagatorMapping', 'Tensor', 'ThreeBodyDecay', 'Topology', 'TwoBodyDecay', 'TwoToThreeParticleScattering', 'TwoToTwoParticleScattering', 'Type', 'Unweighter', 'Value', 'VegasGridOptimizer', 'VegasHistogram', 'VegasMapping', 'Verbosity', 'batch_float', 'batch_float_array', 'batch_four_vec', 'batch_four_vec_array', 'batch_int', 'batch_size', 'batch_sizes', 'cpu_device', 'cuda_device', 'default_context', 'default_cuda_context', 'default_hip_context', 'float', 'format_progress', 'format_si_prefix', 'format_with_error', 'hip_device', 'initialize_vegas_grid', 'int', 'log', 'multichannel_batch_size', 'pretty', 'set_lib_path', 'set_simd_vector_size', 'silent', 'single_float', 'single_int'] class AdamOptimizer: class LRSchedule: """ Members: - + none - + cosine """ - - __members__: typing.ClassVar[ - dict[str, AdamOptimizer.LRSchedule] - ] # value = {'none': , 'cosine': } - cosine: typing.ClassVar[ - AdamOptimizer.LRSchedule - ] # value = + __members__: typing.ClassVar[dict[str, AdamOptimizer.LRSchedule]] # value = {'none': , 'cosine': } + cosine: typing.ClassVar[AdamOptimizer.LRSchedule] # value = none: typing.ClassVar[AdamOptimizer.LRSchedule] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... cosine: typing.ClassVar[AdamOptimizer.LRSchedule] # value = none: typing.ClassVar[AdamOptimizer.LRSchedule] # value = - def __init__( - self, - function: Function, - context: Context, - learning_rate: typing.SupportsFloat, - schedule: AdamOptimizer.LRSchedule = AdamOptimizer.LRSchedule.LRSchedule.none, - step_count: typing.SupportsInt = 0, - beta1: typing.SupportsFloat = 0.9, - beta2: typing.SupportsFloat = 0.999, - eps: typing.SupportsFloat = 1e-08, - ) -> None: ... - def context(self) -> Context: ... - def input_types(self) -> list[Type]: ... - def learning_rate(self) -> float: ... - def step(self, inputs: collections.abc.Sequence[typing.Any]) -> list[Tensor]: ... - + def __init__(self, function: Function, context: Context, learning_rate: typing.SupportsFloat, schedule: AdamOptimizer.LRSchedule = AdamOptimizer.LRSchedule.LRSchedule.none, step_count: typing.SupportsInt = 0, beta1: typing.SupportsFloat = 0.9, beta2: typing.SupportsFloat = 0.999, eps: typing.SupportsFloat = 1e-08) -> None: + ... + def context(self) -> Context: + ... + def input_types(self) -> list[Type]: + ... + def learning_rate(self) -> float: + ... + def step(self, inputs: collections.abc.Sequence[typing.Any]) -> list[Tensor]: + ... class AlphaSGrid: - def __init__(self, file: str) -> None: ... - def coefficients_shape(self, batch_dim: bool = False) -> list[int]: ... - def initialize_globals(self, context: Context, prefix: str = "") -> None: ... - def logq2_shape(self, batch_dim: bool = False) -> list[int]: ... + def __init__(self, file: str) -> None: + ... + def coefficients_shape(self, batch_dim: bool = False) -> list[int]: + ... + def initialize_globals(self, context: Context, prefix: str = '') -> None: + ... + def logq2_shape(self, batch_dim: bool = False) -> list[int]: + ... @property - def logq2(self) -> list[float]: ... + def logq2(self) -> list[float]: + ... @property - def q(self) -> list[float]: ... + def q(self) -> list[float]: + ... @property - def q_count(self) -> int: ... + def q_count(self) -> int: + ... @property - def region_sizes(self) -> list[int]: ... + def region_sizes(self) -> list[int]: + ... @property - def values(self) -> list[float]: ... - + def values(self) -> list[float]: + ... class BatchSize: one: typing.ClassVar[BatchSize] # value = 1 @typing.overload - def __init__(self) -> None: ... + def __init__(self) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... - + def __init__(self, name: str) -> None: + ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... class ChannelEventGenerator: integrand_flags: typing.ClassVar[int] = 1077 @staticmethod - def load( - channel_file: str, - contexts: collections.abc.Sequence[Context], - event_file: str, - weight_file: str, - config: GeneratorConfig, - ) -> ChannelEventGenerator: ... - def __init__( - self, - contexts: collections.abc.Sequence[Context], - integrand: Integrand, - event_file: str, - weight_file: str, - config: GeneratorConfig, - subprocess_index: typing.SupportsInt, - name: str, - histograms: madspace._madspace_py.ObservableHistograms | None, - ) -> None: ... - def save(self, save: str) -> None: ... - def status(self) -> GeneratorStatus: ... - + def load(channel_file: str, contexts: collections.abc.Sequence[Context], event_file: str, weight_file: str, config: GeneratorConfig) -> ChannelEventGenerator: + ... + def __init__(self, contexts: collections.abc.Sequence[Context], integrand: Integrand, event_file: str, weight_file: str, config: GeneratorConfig, subprocess_index: typing.SupportsInt, name: str, histograms: madspace._madspace_py.ObservableHistograms | None) -> None: + ... + def save(self, save: str) -> None: + ... + def status(self) -> GeneratorStatus: + ... class ChannelWeightNetwork(FunctionGenerator): - def __init__( - self, - channel_count: typing.SupportsInt, - particle_count: typing.SupportsInt, - hidden_dim: typing.SupportsInt = 32, - layers: typing.SupportsInt = 3, - activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, - prefix: str = "", - include_preprocessing: bool = True, - ) -> None: ... - def initialize_globals(self, context: Context) -> None: ... - def mask_name(self) -> str: ... - def mlp(self) -> MLP: ... - def preprocessing(self) -> MomentumPreprocessing: ... - + def __init__(self, channel_count: typing.SupportsInt, particle_count: typing.SupportsInt, hidden_dim: typing.SupportsInt = 32, layers: typing.SupportsInt = 3, activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, prefix: str = '', include_preprocessing: bool = True) -> None: + ... + def initialize_globals(self, context: Context) -> None: + ... + def mask_name(self) -> str: + ... + def mlp(self) -> MLP: + ... + def preprocessing(self) -> MomentumPreprocessing: + ... class ColorOrderedMapping(Mapping): - def __init__( - self, - colour_order: collections.abc.Sequence[typing.SupportsInt], - t_invariant_power: typing.SupportsFloat = 0.0, - s_invariant_power: typing.SupportsFloat = 0.0, - ) -> None: ... - def random_dim(self) -> int: ... - + def __init__(self, colour_order: collections.abc.Sequence[typing.SupportsInt], t_invariant_power: typing.SupportsFloat = 0.0, s_invariant_power: typing.SupportsFloat = 0.0) -> None: + ... + def random_dim(self) -> int: + ... class Context: @typing.overload - def __init__(self, thread_count: typing.SupportsInt = -1) -> None: ... + def __init__(self, thread_count: typing.SupportsInt = -1) -> None: + ... @typing.overload - def __init__( - self, device: Device, thread_count: typing.SupportsInt = -1 - ) -> None: ... - def copy_globals_from(self, context: Context) -> None: ... - def define_global( - self, - name: str, - dtype: DataType, - shape: collections.abc.Sequence[typing.SupportsInt], - requires_grad: bool = False, - ) -> Tensor: ... - def delete_global(self, name: str) -> None: ... - def device(self) -> Device: ... - def get_global(self, name: str) -> Tensor: ... - def global_exists(self, name: str) -> bool: ... - def global_names(self) -> list[str]: ... - def global_requires_grad(self, name: str) -> bool: ... - def load_globals(self, dir: str) -> None: ... - def load_matrix_element(self, file: str, param_card: str) -> MatrixElementApi: ... - def matrix_element(self, index: typing.SupportsInt) -> MatrixElementApi: ... - def save_globals(self, dir: str) -> None: ... - + def __init__(self, device: Device, thread_count: typing.SupportsInt = -1) -> None: + ... + def copy_globals_from(self, context: Context) -> None: + ... + def define_global(self, name: str, dtype: DataType, shape: collections.abc.Sequence[typing.SupportsInt], requires_grad: bool = False) -> Tensor: + ... + def delete_global(self, name: str) -> None: + ... + def device(self) -> Device: + ... + def get_global(self, name: str) -> Tensor: + ... + def global_exists(self, name: str) -> bool: + ... + def global_names(self) -> list[str]: + ... + def global_requires_grad(self, name: str) -> bool: + ... + def load_globals(self, dir: str) -> None: + ... + def load_matrix_element(self, file: str, param_card: str) -> MatrixElementApi: + ... + def matrix_element(self, index: typing.SupportsInt) -> MatrixElementApi: + ... + def save_globals(self, dir: str) -> None: + ... class CutItem: - def __init__( - self, - observable: Observable, - min: typing.SupportsFloat = ..., - max: typing.SupportsFloat = ..., - mode: Cuts.CutMode = Cuts.CutMode.CutMode.all, - ) -> None: ... + def __init__(self, observable: Observable, min: typing.SupportsFloat = ..., max: typing.SupportsFloat = ..., mode: Cuts.CutMode = Cuts.CutMode.CutMode.all) -> None: + ... @property - def max(self) -> float: ... + def max(self) -> float: + ... @property - def min(self) -> float: ... + def min(self) -> float: + ... @property - def mode(self) -> Cuts.CutMode: ... + def mode(self) -> Cuts.CutMode: + ... @property - def observable(self) -> Observable: ... - + def observable(self) -> Observable: + ... class Cuts(FunctionGenerator): class CutMode: """ Members: - + any - + all """ - - __members__: typing.ClassVar[ - dict[str, Cuts.CutMode] - ] # value = {'any': , 'all': } + __members__: typing.ClassVar[dict[str, Cuts.CutMode]] # value = {'any': , 'all': } all: typing.ClassVar[Cuts.CutMode] # value = any: typing.ClassVar[Cuts.CutMode] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... all: typing.ClassVar[Cuts.CutMode] # value = any: typing.ClassVar[Cuts.CutMode] # value = @typing.overload - def __init__(self, cut_data: collections.abc.Sequence[CutItem]) -> None: ... + def __init__(self, cut_data: collections.abc.Sequence[CutItem]) -> None: + ... @typing.overload - def __init__(self, particle_count: typing.SupportsInt) -> None: ... - def eta_max(self) -> list[float]: ... - def pt_min(self) -> list[float]: ... - def sqrt_s_min(self) -> float: ... - + def __init__(self, particle_count: typing.SupportsInt) -> None: + ... + def eta_max(self) -> list[float]: + ... + def pt_min(self) -> list[float]: + ... + def sqrt_s_min(self) -> float: + ... class DataType: """ Members: - + int - + float - + batch_sizes """ - - __members__: typing.ClassVar[ - dict[str, DataType] - ] # value = {'int': , 'float': , 'batch_sizes': } + __members__: typing.ClassVar[dict[str, DataType]] # value = {'int': , 'float': , 'batch_sizes': } batch_sizes: typing.ClassVar[DataType] # value = float: typing.ClassVar[DataType] # value = int: typing.ClassVar[DataType] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... - @property - def name(self) -> str: ... - @property - def value(self) -> int: ... - + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... + @property + def name(self) -> str: + ... + @property + def value(self) -> int: + ... class Decay: @property - def child_indices(self) -> list[int]: ... + def child_indices(self) -> list[int]: + ... @property - def e_max(self) -> float: ... + def e_max(self) -> float: + ... @property - def e_min(self) -> float: ... + def e_min(self) -> float: + ... @property - def index(self) -> int: ... + def index(self) -> int: + ... @property - def mass(self) -> float: ... + def mass(self) -> float: + ... @property - def on_shell(self) -> bool: ... + def on_shell(self) -> bool: + ... @property - def parent_index(self) -> int: ... + def parent_index(self) -> int: + ... @property - def pdg_id(self) -> int: ... + def pdg_id(self) -> int: + ... @property - def width(self) -> float: ... - + def width(self) -> float: + ... class Device: pass - class Diagram: - def __init__( - self, - incoming_masses: collections.abc.Sequence[typing.SupportsFloat], - outgoing_masses: collections.abc.Sequence[typing.SupportsFloat], - propagators: collections.abc.Sequence[Propagator], - vertices: collections.abc.Sequence[collections.abc.Sequence[LineRef]], - ) -> None: ... + def __init__(self, incoming_masses: collections.abc.Sequence[typing.SupportsFloat], outgoing_masses: collections.abc.Sequence[typing.SupportsFloat], propagators: collections.abc.Sequence[Propagator], vertices: collections.abc.Sequence[collections.abc.Sequence[LineRef]]) -> None: + ... @property - def incoming_masses(self) -> list[float]: ... + def incoming_masses(self) -> list[float]: + ... @property - def incoming_vertices(self) -> typing.Annotated[list[int], "FixedSize(2)"]: ... + def incoming_vertices(self) -> typing.Annotated[list[int], "FixedSize(2)"]: + ... @property - def outgoing_masses(self) -> list[float]: ... + def outgoing_masses(self) -> list[float]: + ... @property - def outgoing_vertices(self) -> list[int]: ... + def outgoing_vertices(self) -> list[int]: + ... @property - def propagator_vertices(self) -> list[list[int]]: ... + def propagator_vertices(self) -> list[list[int]]: + ... @property - def propagators(self) -> list[Propagator]: ... + def propagators(self) -> list[Propagator]: + ... @property - def vertices(self) -> list[list[LineRef]]: ... - + def vertices(self) -> list[list[LineRef]]: + ... class DifferentialCrossSection(FunctionGenerator): - def __init__( - self, - matrix_element: MatrixElement, - cm_energy: typing.SupportsFloat, - running_coupling: RunningCoupling, - energy_scale: EnergyScale, - pid_options: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ] = [], - has_pdf1: bool = False, - has_pdf2: bool = False, - pdf_grid1: madspace._madspace_py.PdfGrid | None = None, - pdf_grid2: madspace._madspace_py.PdfGrid | None = None, - has_mirror: bool = False, - input_momentum_fraction: bool = True, - ) -> None: ... - def has_mirror(self) -> bool: ... - def matrix_element(self) -> MatrixElement: ... - def pid_options(self) -> list[list[int]]: ... - + def __init__(self, matrix_element: MatrixElement, cm_energy: typing.SupportsFloat, running_coupling: RunningCoupling, energy_scale: EnergyScale, pid_options: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] = [], has_pdf1: bool = False, has_pdf2: bool = False, pdf_grid1: madspace._madspace_py.PdfGrid | None = None, pdf_grid2: madspace._madspace_py.PdfGrid | None = None, has_mirror: bool = False, input_momentum_fraction: bool = True) -> None: + ... + def has_mirror(self) -> bool: + ... + def matrix_element(self) -> MatrixElement: + ... + def pid_options(self) -> list[list[int]]: + ... class DiscreteFlow(Mapping): - def __init__( - self, - option_counts: collections.abc.Sequence[typing.SupportsInt], - prefix: str = "", - dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = [], - condition_dim: typing.SupportsInt = 0, - subnet_hidden_dim: typing.SupportsInt = 32, - subnet_layers: typing.SupportsInt = 3, - subnet_activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, - ) -> None: ... - def condition_dim(self) -> int: ... - def initialize_globals(self, context: Context) -> None: ... - def option_counts(self) -> list[int]: ... - + def __init__(self, option_counts: collections.abc.Sequence[typing.SupportsInt], prefix: str = '', dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = [], condition_dim: typing.SupportsInt = 0, subnet_hidden_dim: typing.SupportsInt = 32, subnet_layers: typing.SupportsInt = 3, subnet_activation: MLP.Activation = MLP.Activation.Activation.leaky_relu) -> None: + ... + def condition_dim(self) -> int: + ... + def initialize_globals(self, context: Context) -> None: + ... + def option_counts(self) -> list[int]: + ... class DiscreteHistogram(FunctionGenerator): - def __init__( - self, option_counts: collections.abc.Sequence[typing.SupportsInt] - ) -> None: ... - + def __init__(self, option_counts: collections.abc.Sequence[typing.SupportsInt]) -> None: + ... class DiscreteOptimizer: - def __init__( - self, - contexts: collections.abc.Sequence[Context], - prob_names: collections.abc.Sequence[str], - ) -> None: ... - def add_data( - self, values_and_counts: collections.abc.Sequence[typing.Any] - ) -> None: ... - def optimize(self) -> None: ... - + def __init__(self, contexts: collections.abc.Sequence[Context], prob_names: collections.abc.Sequence[str]) -> None: + ... + def add_data(self, values_and_counts: collections.abc.Sequence[typing.Any]) -> None: + ... + def optimize(self) -> None: + ... class DiscreteSampler(Mapping): - def __init__( - self, - option_counts: collections.abc.Sequence[typing.SupportsInt], - prefix: str = "", - dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = [], - ) -> None: ... - def initialize_globals(self, context: Context) -> None: ... - def option_counts(self) -> list[int]: ... - def prob_names(self) -> list[str]: ... - + def __init__(self, option_counts: collections.abc.Sequence[typing.SupportsInt], prefix: str = '', dims_with_prior: collections.abc.Sequence[typing.SupportsInt] = []) -> None: + ... + def initialize_globals(self, context: Context) -> None: + ... + def option_counts(self) -> list[int]: + ... + def prob_names(self) -> list[str]: + ... class DoubleT(Mapping): - def __init__( - self, - t1_invariant_power: typing.SupportsFloat = 0.0, - t1_mass: typing.SupportsFloat = 0.0, - t1_width: typing.SupportsFloat = 0.0, - t2_invariant_power: typing.SupportsFloat = 0.0, - t2_mass: typing.SupportsFloat = 0.0, - t2_width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, t1_invariant_power: typing.SupportsFloat = 0.0, t1_mass: typing.SupportsFloat = 0.0, t1_width: typing.SupportsFloat = 0.0, t2_invariant_power: typing.SupportsFloat = 0.0, t2_mass: typing.SupportsFloat = 0.0, t2_width: typing.SupportsFloat = 0.0) -> None: + ... class EnergyScale(FunctionGenerator): class DynamicalScaleType: """ Members: - + transverse_energy - + transverse_mass - + half_transverse_mass - + partonic_energy """ - - __members__: typing.ClassVar[ - dict[str, EnergyScale.DynamicalScaleType] - ] # value = {'transverse_energy': , 'transverse_mass': , 'half_transverse_mass': , 'partonic_energy': } - half_transverse_mass: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - partonic_energy: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - transverse_energy: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - transverse_mass: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, EnergyScale.DynamicalScaleType]] # value = {'transverse_energy': , 'transverse_mass': , 'half_transverse_mass': , 'partonic_energy': } + half_transverse_mass: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + partonic_energy: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + transverse_energy: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + transverse_mass: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - - half_transverse_mass: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - partonic_energy: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - transverse_energy: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = - transverse_mass: typing.ClassVar[ - EnergyScale.DynamicalScaleType - ] # value = + def value(self) -> int: + ... + half_transverse_mass: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + partonic_energy: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + transverse_energy: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = + transverse_mass: typing.ClassVar[EnergyScale.DynamicalScaleType] # value = @typing.overload - def __init__(self, particle_count: typing.SupportsInt) -> None: ... + def __init__(self, particle_count: typing.SupportsInt) -> None: + ... @typing.overload - def __init__( - self, particle_count: typing.SupportsInt, type: EnergyScale.DynamicalScaleType - ) -> None: ... + def __init__(self, particle_count: typing.SupportsInt, type: EnergyScale.DynamicalScaleType) -> None: + ... @typing.overload - def __init__( - self, particle_count: typing.SupportsInt, fixed_scale: typing.SupportsFloat - ) -> None: ... + def __init__(self, particle_count: typing.SupportsInt, fixed_scale: typing.SupportsFloat) -> None: + ... @typing.overload - def __init__( - self, - particle_count: typing.SupportsInt, - dynamical_scale_type: EnergyScale.DynamicalScaleType, - ren_scale_fixed: bool, - fact_scale_fixed: bool, - ren_scale: typing.SupportsFloat, - fact_scale1: typing.SupportsFloat, - fact_scale2: typing.SupportsFloat, - ) -> None: ... - + def __init__(self, particle_count: typing.SupportsInt, dynamical_scale_type: EnergyScale.DynamicalScaleType, ren_scale_fixed: bool, fact_scale_fixed: bool, ren_scale: typing.SupportsFloat, fact_scale1: typing.SupportsFloat, fact_scale2: typing.SupportsFloat) -> None: + ... class EventGenerator: - default_config: typing.ClassVar[ - GeneratorConfig - ] # value = - def __init__( - self, - contexts: collections.abc.Sequence[Context], - channels: collections.abc.Sequence[ChannelEventGenerator], - status_file: str = "", - config: GeneratorConfig = ..., - ) -> None: ... - def channel_status(self) -> list[GeneratorStatus]: ... - def channels(self) -> list[ChannelEventGenerator]: ... - def combine_to_compact_npy(self, file_name: str) -> None: ... - def combine_to_lhe(self, file_name: str, lhe_completer: LHECompleter) -> None: ... - def combine_to_lhe_npy( - self, file_name: str, lhe_completer: LHECompleter - ) -> None: ... - def generate(self) -> None: ... - def histograms(self) -> list[Histogram]: ... - def status(self) -> GeneratorStatus: ... - def survey(self) -> None: ... - def used_globals(self) -> set[str]: ... - + default_config: typing.ClassVar[GeneratorConfig] # value = + def __init__(self, contexts: collections.abc.Sequence[Context], channels: collections.abc.Sequence[ChannelEventGenerator], status_file: str = '', config: GeneratorConfig = ...) -> None: + ... + def channel_status(self) -> list[GeneratorStatus]: + ... + def channels(self) -> list[ChannelEventGenerator]: + ... + def combine_to_compact_npy(self, file_name: str) -> None: + ... + def combine_to_lhe(self, file_name: str, lhe_completer: LHECompleter) -> None: + ... + def combine_to_lhe_npy(self, file_name: str, lhe_completer: LHECompleter) -> None: + ... + def generate(self) -> None: + ... + def histograms(self) -> list[Histogram]: + ... + def status(self) -> GeneratorStatus: + ... + def survey(self) -> None: + ... + def used_globals(self) -> set[str]: + ... class FastRamboMapping(Mapping): - def __init__(self, n_particles: typing.SupportsInt, massless: bool) -> None: ... - + def __init__(self, n_particles: typing.SupportsInt, massless: bool) -> None: + ... class Flow(Mapping): - def __init__( - self, - input_dim: typing.SupportsInt, - condition_dim: typing.SupportsInt = 0, - prefix: str = "", - bin_count: typing.SupportsInt = 10, - subnet_hidden_dim: typing.SupportsInt = 32, - subnet_layers: typing.SupportsInt = 3, - subnet_activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, - invert_spline: bool = True, - ) -> None: ... - def condition_dim(self) -> int: ... - def initialize_from_vegas(self, context: Context, grid_name: str) -> None: ... - def initialize_globals(self, context: Context) -> None: ... - def input_dim(self) -> int: ... - + def __init__(self, input_dim: typing.SupportsInt, condition_dim: typing.SupportsInt = 0, prefix: str = '', bin_count: typing.SupportsInt = 10, subnet_hidden_dim: typing.SupportsInt = 32, subnet_layers: typing.SupportsInt = 3, subnet_activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, invert_spline: bool = True) -> None: + ... + def condition_dim(self) -> int: + ... + def initialize_from_vegas(self, context: Context, grid_name: str) -> None: + ... + def initialize_globals(self, context: Context) -> None: + ... + def input_dim(self) -> int: + ... class Function: @staticmethod - def load(file: str) -> Function: ... - def __call__(self, *args, **kwargs): ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... - def save(self, file: str) -> None: ... - @property - def globals(self) -> list[tuple[str, Value]]: ... - @property - def inputs(self) -> NamedValues: ... - @property - def instructions(self) -> list[InstructionCall]: ... - @property - def locals(self) -> list[Value]: ... - @property - def outputs(self) -> NamedValues: ... - + def load(file: str) -> Function: + ... + def __call__(self, *args, **kwargs): + ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... + def save(self, file: str) -> None: + ... + @property + def globals(self) -> list[tuple[str, Value]]: + ... + @property + def inputs(self) -> NamedValues: + ... + @property + def instructions(self) -> list[InstructionCall]: + ... + @property + def locals(self) -> list[Value]: + ... + @property + def outputs(self) -> NamedValues: + ... class FunctionBuilder: - def __init__(self, input_types: NamedTypes, output_types: NamedTypes) -> None: ... - def accept_norm(self, accepted_batch: Value, full_batch: Value) -> Value: ... - def add(self, in1: Value, in2: Value) -> Value: ... - def add_int(self, in1: Value, in2: Value) -> Value: ... - def apply_subchannel_weights( - self, - channel_weights_in: Value, - subchannel_weights: Value, - channel_indices: Value, - subchannel_indices: Value, - ) -> Value: ... - def argsort(self, input: Value) -> Value: ... - def batch_cat( - self, args: collections.abc.Sequence[Value] - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def batch_gather(self, indices: Value, values: Value) -> Value: ... - def batch_reduce_mean(self, in1: Value) -> Value: ... - def batch_reduce_mean_keepdim(self, in1: Value) -> Value: ... - def batch_scatter(self, indices: Value, target: Value, source: Value) -> Value: ... - def batch_size(self, args: collections.abc.Sequence[Value]) -> Value: ... - def batch_split(self, in1: Value, counts: Value) -> list[Value]: ... - def boost_beam(self, p1: Value, x1: Value, x2: Value) -> Value: ... - def boost_beam_inverse(self, p1: Value, x1: Value, x2: Value) -> Value: ... - def breit_wigner_invariant( - self, r: Value, mass: Value, width: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def breit_wigner_invariant_inverse( - self, s: Value, mass: Value, width: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def cat(self, args: collections.abc.Sequence[Value]) -> Value: ... - def chili_forward( - self, r: Value, e_cm: Value, m_out: Value, pt_min: Value, y_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def chili_inverse( - self, p_ext: Value, e_cm: Value, m_out: Value, pt_min: Value, y_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def collect_channel_weights( - self, amp2: Value, channel_indices: Value, channel_count: Value - ) -> Value: ... - def com_p_in( - self, e_cm: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def current_stream(self) -> int: ... - def cut_all(self, obs: Value, min: Value, max: Value) -> Value: ... - def cut_any(self, obs: Value, min: Value, max: Value) -> Value: ... - def cut_one(self, obs: Value, min: Value, max: Value) -> Value: ... - def cut_unphysical(self, w_in: Value, p: Value, x1: Value, x2: Value) -> Value: ... - def diff_cross_section( - self, - x1: Value, - x2: Value, - pdf1: Value, - pdf2: Value, - matrix_element: Value, - e_cm2: Value, - ) -> Value: ... - def discrete_histogram( - self, input: Value, weights: Value, option_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def div(self, in1: Value, in2: Value) -> Value: ... - def double_t_scattering( - self, - r_phi: Value, - pa: Value, - pb: Value, - t1_abs: Value, - t2_abs: Value, - m1: Value, - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def double_t_scattering_inverse( - self, p1: Value, p2: Value, pa: Value, pb: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def elu(self, in1: Value) -> Value: ... - def fast_rambo_massive( - self, r: Value, e_cm: Value, masses: Value, p0: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def fast_rambo_massive_com( - self, r: Value, e_cm: Value, masses: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def fast_rambo_massive_inverse( - self, p_out: Value, e_cm: Value, masses: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def fast_rambo_massless( - self, r: Value, e_cm: Value, p0: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def fast_rambo_massless_com( - self, r: Value, e_cm: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def fast_rambo_massless_inverse( - self, p_out: Value, e_cm: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def full(self, args: collections.abc.Sequence[Value]) -> Value: ... - def function(self) -> Function: ... - def gather(self, index: Value, choices: Value) -> Value: ... - def gather_int(self, index: Value, choices: Value) -> Value: ... - def gelu(self, in1: Value) -> Value: ... - def get_global( - self, - name: str, - dtype: DataType, - shape: collections.abc.Sequence[typing.SupportsInt], - ) -> Value: ... - def histogram( - self, input: Value, weights: Value, min: Value, max: Value, bin_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def input(self, index: typing.SupportsInt) -> Value: ... - def input_range( - self, start_index: typing.SupportsInt, end_index: typing.SupportsInt - ) -> list[Value]: ... - def interpolate_alpha_s( - self, q2: Value, grid_logq2: Value, grid_coeffs: Value - ) -> Value: ... - def interpolate_pdf( - self, - x: Value, - q2: Value, - pid_indices: Value, - grid_logx: Value, - grid_logq2: Value, - grid_coeffs: Value, - ) -> Value: ... - def invariants_from_momenta(self, p_ext: Value, factors: Value) -> Value: ... - def leaky_relu(self, in1: Value) -> Value: ... - def madnis_abs_weight(self, f: Value, q: Value) -> Value: ... - def madnis_multi_channel_variance(self, vars: Value, abs_means: Value) -> Value: ... - def madnis_single_channel_variance(self, var: Value, abs_mean: Value) -> Value: ... - def madnis_softclip( - self, f: Value, q: Value, norm: Value, threshold: Value - ) -> Value: ... - def madnis_variance(self, f: Value, g: Value, q: Value, mean: Value) -> Value: ... - def matmul(self, x: Value, weight: Value, bias: Value) -> Value: ... - def matrix_element(self, args: collections.abc.Sequence[Value]) -> list[Value]: ... - def max(self, in1: Value, in2: Value) -> Value: ... - def min(self, in1: Value, in2: Value) -> Value: ... - def mirror_momenta(self, p_ext: Value, mirror: Value) -> Value: ... - def momenta_to_x1x2( - self, p_ext: Value, e_cm: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def mul(self, in1: Value, in2: Value) -> Value: ... - def nonzero(self, input: Value) -> Value: ... - def obs_delta_eta(self, p1: Value, p2: Value) -> Value: ... - def obs_delta_phi(self, p1: Value, p2: Value) -> Value: ... - def obs_delta_r(self, p1: Value, p2: Value) -> Value: ... - def obs_e(self, p: Value) -> Value: ... - def obs_eta(self, p: Value) -> Value: ... - def obs_eta_abs(self, p: Value) -> Value: ... - def obs_mass(self, p: Value) -> Value: ... - def obs_p_mag(self, p: Value) -> Value: ... - def obs_phi(self, p: Value) -> Value: ... - def obs_pt(self, p: Value) -> Value: ... - def obs_px(self, p: Value) -> Value: ... - def obs_py(self, p: Value) -> Value: ... - def obs_pz(self, p: Value) -> Value: ... - def obs_sqrt_s(self, p_ext: Value) -> Value: ... - def obs_theta(self, p: Value) -> Value: ... - def obs_y(self, p: Value) -> Value: ... - def obs_y_abs(self, p: Value) -> Value: ... - def offset_indices( - self, batch_sizes_offset: Value, batch_sizes_out: Value - ) -> Value: ... - def one_hot(self, index: Value, option_count: Value) -> Value: ... - def output(self, index: typing.SupportsInt, value: Value) -> None: ... - def output_range( - self, start_index: typing.SupportsInt, values: collections.abc.Sequence[Value] - ) -> None: ... - def permute_momenta( - self, momenta: Value, permutations: Value, index: Value - ) -> Value: ... - def pop(self, in1: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def product(self, values: collections.abc.Sequence[Value]) -> Value: ... - def pt_eta_phi_x(self, p_ext: Value, x1: Value, x2: Value) -> Value: ... - def r_to_x1x2( - self, r: Value, s_hat: Value, s_lab: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def random(self, batch_size: Value, count: Value) -> Value: ... - def reduce_product(self, in1: Value) -> Value: ... - def reduce_sum(self, in1: Value) -> Value: ... - def reduce_sum_vector(self, in1: Value) -> Value: ... - def relu(self, in1: Value) -> Value: ... - def rqs_find_bin( - self, input: Value, in_sizes: Value, out_sizes: Value, derivatives: Value - ) -> Value: ... - def rqs_forward( - self, input: Value, condition: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def rqs_inverse( - self, input: Value, condition: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def rqs_reshape( - self, input: Value, bin_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def s23_min_max( - self, pa: Value, pb: Value, p3: Value, t1_abs: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def s23_value_and_min_max( - self, pa: Value, pb: Value, p3: Value, t1_abs: Value, p1: Value, p2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def sample_discrete( - self, r: Value, option_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def sample_discrete_inverse( - self, index: Value, option_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def sample_discrete_probs( - self, r: Value, probs: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def sample_discrete_probs_inverse( - self, index: Value, probs: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def scale_half_transverse_mass(self, momenta: Value) -> Value: ... - def scale_partonic_energy(self, momenta: Value) -> Value: ... - def scale_transverse_energy(self, momenta: Value) -> Value: ... - def scale_transverse_mass(self, momenta: Value) -> Value: ... - def sde2_channel_weights( - self, invariants: Value, masses: Value, widths: Value, indices: Value - ) -> Value: ... - def select(self, input: Value, indices: Value) -> Value: ... - def select_int(self, input: Value, indices: Value) -> Value: ... - def select_vector(self, input: Value, indices: Value) -> Value: ... - def set_current_stream(self, arg0: typing.SupportsInt) -> None: ... - def sigmoid(self, in1: Value) -> Value: ... - def softmax(self, input: Value) -> Value: ... - def softmax_prior(self, input: Value, prior: Value) -> Value: ... - def softplus(self, in1: Value) -> Value: ... - def sqrt(self, in1: Value) -> Value: ... - def square(self, in1: Value) -> Value: ... - def squeeze(self, input: Value) -> Value: ... - def stable_invariant( - self, r: Value, mass: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def stable_invariant_inverse( - self, s: Value, mass: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def stable_invariant_nu( - self, r: Value, mass: Value, nu: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def stable_invariant_nu_inverse( - self, s: Value, mass: Value, nu: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def stack(self, args: collections.abc.Sequence[Value]) -> Value: ... - def stack_sizes(self, args: collections.abc.Sequence[Value]) -> Value: ... - def sub(self, in1: Value, in2: Value) -> Value: ... - def subchannel_weights( - self, - invariants: Value, - masses: Value, - widths: Value, - indices: Value, - on_shell: Value, - group_sizes: Value, - ) -> Value: ... - def t1_inv_min_max_doublet( - self, pa: Value, pb: Value, m1: Value, mir_min: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def t1_inv_value_and_min_max_doublet( - self, pa: Value, pb: Value, p1: Value, m1: Value, mir_min: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def t2_inv_min_max_doublet( - self, pa: Value, pb: Value, m1: Value, mir_min: Value, t1_abs: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def t2_inv_value_and_min_max_doublet( - self, pa: Value, pb: Value, p1: Value, m1: Value, mir_min: Value, t1_abs: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def t_inv_min_max( - self, pa: Value, pb: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def t_inv_value_and_min_max( - self, pa: Value, pb: Value, p1: Value, p2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def three_body_decay( - self, - r_e1: Value, - r_e2: Value, - r_phi: Value, - r_cos_theta: Value, - r_beta: Value, - m0: Value, - m1: Value, - m2: Value, - m3: Value, - p0: Value, - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def three_body_decay_com( - self, - r_e1: Value, - r_e2: Value, - r_phi: Value, - r_cos_theta: Value, - r_beta: Value, - m0: Value, - m1: Value, - m2: Value, - m3: Value, - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def three_body_decay_com_inverse( - self, p1: Value, p2: Value, p3: Value - ) -> typing.Annotated[list[Value], "FixedSize(10)"]: ... - def three_body_decay_inverse( - self, p1: Value, p2: Value, p3: Value - ) -> typing.Annotated[list[Value], "FixedSize(11)"]: ... - def two_body_decay( - self, - r_phi: Value, - r_cos_theta: Value, - m0: Value, - m1: Value, - m2: Value, - p0: Value, - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_body_decay_com( - self, r_phi: Value, r_cos_theta: Value, m0: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_body_decay_com_inverse( - self, p1: Value, p2: Value - ) -> typing.Annotated[list[Value], "FixedSize(6)"]: ... - def two_body_decay_inverse( - self, p1: Value, p2: Value - ) -> typing.Annotated[list[Value], "FixedSize(7)"]: ... - def two_to_three_particle_scattering( - self, - phi_choice: Value, - pa: Value, - pb: Value, - p3: Value, - s23: Value, - t1_abs: Value, - m1: Value, - m2: Value, - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_to_three_particle_scattering_inverse( - self, - p1: Value, - p2: Value, - p3: Value, - pa: Value, - pb: Value, - t1_abs: Value, - s23: Value, - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def two_to_two_particle_scattering( - self, r_phi: Value, pa: Value, pb: Value, t_abs: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_to_two_particle_scattering_com( - self, r_phi: Value, pa: Value, pb: Value, t_abs: Value, m1: Value, m2: Value - ) -> typing.Annotated[list[Value], "FixedSize(3)"]: ... - def two_to_two_particle_scattering_com_inverse( - self, p1: Value, p2: Value, pa: Value, pb: Value - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def two_to_two_particle_scattering_inverse( - self, p1: Value, p2: Value, pa: Value, pb: Value - ) -> typing.Annotated[list[Value], "FixedSize(4)"]: ... - def uniform_invariant( - self, r: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def uniform_invariant_inverse( - self, s: Value, s_min: Value, s_max: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def unsqueeze(self, input: Value) -> Value: ... - def unstack(self, in1: Value) -> list[Value]: ... - def unstack_sizes(self, in1: Value) -> list[Value]: ... - def unweight( - self, weights: Value, max_weight: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def vegas_forward( - self, input: Value, grid: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def vegas_histogram( - self, input: Value, weights: Value, bin_count: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def vegas_inverse( - self, input: Value, grid: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - def x1x2_to_r( - self, x1: Value, x2: Value, s_lab: Value - ) -> typing.Annotated[list[Value], "FixedSize(2)"]: ... - + def __init__(self, input_types: NamedTypes, output_types: NamedTypes) -> None: + ... + def accept_norm(self, accepted_batch: Value, full_batch: Value) -> Value: + ... + def add(self, in1: Value, in2: Value) -> Value: + ... + def add_int(self, in1: Value, in2: Value) -> Value: + ... + def apply_subchannel_weights(self, channel_weights_in: Value, subchannel_weights: Value, channel_indices: Value, subchannel_indices: Value) -> Value: + ... + def argsort(self, input: Value) -> Value: + ... + def batch_cat(self, args: collections.abc.Sequence[Value]) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def batch_gather(self, indices: Value, values: Value) -> Value: + ... + def batch_reduce_mean(self, in1: Value) -> Value: + ... + def batch_reduce_mean_keepdim(self, in1: Value) -> Value: + ... + def batch_scatter(self, indices: Value, target: Value, source: Value) -> Value: + ... + def batch_size(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def batch_split(self, in1: Value, counts: Value) -> list[Value]: + ... + def boost_beam(self, p1: Value, x1: Value, x2: Value) -> Value: + ... + def boost_beam_inverse(self, p1: Value, x1: Value, x2: Value) -> Value: + ... + def breit_wigner_invariant(self, r: Value, mass: Value, width: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def breit_wigner_invariant_inverse(self, s: Value, mass: Value, width: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def cat(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def chili_forward(self, r: Value, e_cm: Value, m_out: Value, pt_min: Value, y_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def chili_inverse(self, p_ext: Value, e_cm: Value, m_out: Value, pt_min: Value, y_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def collect_channel_weights(self, amp2: Value, channel_indices: Value, channel_count: Value) -> Value: + ... + def com_p_in(self, e_cm: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def current_stream(self) -> int: + ... + def cut_all(self, obs: Value, min: Value, max: Value) -> Value: + ... + def cut_any(self, obs: Value, min: Value, max: Value) -> Value: + ... + def cut_one(self, obs: Value, min: Value, max: Value) -> Value: + ... + def cut_unphysical(self, w_in: Value, p: Value, x1: Value, x2: Value) -> Value: + ... + def diff_cross_section(self, x1: Value, x2: Value, pdf1: Value, pdf2: Value, matrix_element: Value, e_cm2: Value) -> Value: + ... + def discrete_histogram(self, input: Value, weights: Value, option_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def div(self, in1: Value, in2: Value) -> Value: + ... + def double_t_scattering(self, r_phi: Value, pa: Value, pb: Value, t1_abs: Value, t2_abs: Value, m1: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def double_t_scattering_inverse(self, p1: Value, p2: Value, pa: Value, pb: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def elu(self, in1: Value) -> Value: + ... + def fast_rambo_massive(self, r: Value, e_cm: Value, masses: Value, p0: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def fast_rambo_massive_com(self, r: Value, e_cm: Value, masses: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def fast_rambo_massive_inverse(self, p_out: Value, e_cm: Value, masses: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def fast_rambo_massless(self, r: Value, e_cm: Value, p0: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def fast_rambo_massless_com(self, r: Value, e_cm: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def fast_rambo_massless_inverse(self, p_out: Value, e_cm: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def full(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def function(self) -> Function: + ... + def gather(self, index: Value, choices: Value) -> Value: + ... + def gather_int(self, index: Value, choices: Value) -> Value: + ... + def gelu(self, in1: Value) -> Value: + ... + def get_global(self, name: str, dtype: DataType, shape: collections.abc.Sequence[typing.SupportsInt]) -> Value: + ... + def histogram(self, input: Value, weights: Value, min: Value, max: Value, bin_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def input(self, index: typing.SupportsInt) -> Value: + ... + def input_range(self, start_index: typing.SupportsInt, end_index: typing.SupportsInt) -> list[Value]: + ... + def interpolate_alpha_s(self, q2: Value, grid_logq2: Value, grid_coeffs: Value) -> Value: + ... + def interpolate_pdf(self, x: Value, q2: Value, pid_indices: Value, grid_logx: Value, grid_logq2: Value, grid_coeffs: Value) -> Value: + ... + def invariants_from_momenta(self, p_ext: Value, factors: Value) -> Value: + ... + def leaky_relu(self, in1: Value) -> Value: + ... + def madnis_abs_weight(self, f: Value, q: Value) -> Value: + ... + def madnis_multi_channel_variance(self, vars: Value, abs_means: Value) -> Value: + ... + def madnis_single_channel_variance(self, var: Value, abs_mean: Value) -> Value: + ... + def madnis_softclip(self, f: Value, q: Value, norm: Value, threshold: Value) -> Value: + ... + def madnis_variance(self, f: Value, g: Value, q: Value, mean: Value) -> Value: + ... + def matmul(self, x: Value, weight: Value, bias: Value) -> Value: + ... + def matrix_element(self, args: collections.abc.Sequence[Value]) -> list[Value]: + ... + def max(self, in1: Value, in2: Value) -> Value: + ... + def min(self, in1: Value, in2: Value) -> Value: + ... + def mirror_momenta(self, p_ext: Value, mirror: Value) -> Value: + ... + def momenta_to_x1x2(self, p_ext: Value, e_cm: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def mul(self, in1: Value, in2: Value) -> Value: + ... + def nonzero(self, input: Value) -> Value: + ... + def obs_delta_eta(self, p1: Value, p2: Value) -> Value: + ... + def obs_delta_phi(self, p1: Value, p2: Value) -> Value: + ... + def obs_delta_r(self, p1: Value, p2: Value) -> Value: + ... + def obs_e(self, p: Value) -> Value: + ... + def obs_eta(self, p: Value) -> Value: + ... + def obs_eta_abs(self, p: Value) -> Value: + ... + def obs_mass(self, p: Value) -> Value: + ... + def obs_p_mag(self, p: Value) -> Value: + ... + def obs_phi(self, p: Value) -> Value: + ... + def obs_pt(self, p: Value) -> Value: + ... + def obs_px(self, p: Value) -> Value: + ... + def obs_py(self, p: Value) -> Value: + ... + def obs_pz(self, p: Value) -> Value: + ... + def obs_sqrt_s(self, p_ext: Value) -> Value: + ... + def obs_theta(self, p: Value) -> Value: + ... + def obs_y(self, p: Value) -> Value: + ... + def obs_y_abs(self, p: Value) -> Value: + ... + def offset_indices(self, batch_sizes_offset: Value, batch_sizes_out: Value) -> Value: + ... + def one_hot(self, index: Value, option_count: Value) -> Value: + ... + def output(self, index: typing.SupportsInt, value: Value) -> None: + ... + def output_range(self, start_index: typing.SupportsInt, values: collections.abc.Sequence[Value]) -> None: + ... + def permute_momenta(self, momenta: Value, permutations: Value, index: Value) -> Value: + ... + def pop(self, in1: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def product(self, values: collections.abc.Sequence[Value]) -> Value: + ... + def pt_eta_phi_x(self, p_ext: Value, x1: Value, x2: Value) -> Value: + ... + def r_to_x1x2(self, r: Value, s_hat: Value, s_lab: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def random(self, batch_size: Value, count: Value) -> Value: + ... + def reduce_product(self, in1: Value) -> Value: + ... + def reduce_sum(self, in1: Value) -> Value: + ... + def reduce_sum_vector(self, in1: Value) -> Value: + ... + def relu(self, in1: Value) -> Value: + ... + def rqs_find_bin(self, input: Value, in_sizes: Value, out_sizes: Value, derivatives: Value) -> Value: + ... + def rqs_forward(self, input: Value, condition: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def rqs_inverse(self, input: Value, condition: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def rqs_reshape(self, input: Value, bin_count: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def s23_min_max(self, pa: Value, pb: Value, p3: Value, t1_abs: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def s23_value_and_min_max(self, pa: Value, pb: Value, p3: Value, t1_abs: Value, p1: Value, p2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def sample_discrete(self, r: Value, option_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def sample_discrete_inverse(self, index: Value, option_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def sample_discrete_probs(self, r: Value, probs: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def sample_discrete_probs_inverse(self, index: Value, probs: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def scale_half_transverse_mass(self, momenta: Value) -> Value: + ... + def scale_partonic_energy(self, momenta: Value) -> Value: + ... + def scale_transverse_energy(self, momenta: Value) -> Value: + ... + def scale_transverse_mass(self, momenta: Value) -> Value: + ... + def sde2_channel_weights(self, invariants: Value, masses: Value, widths: Value, indices: Value) -> Value: + ... + def select(self, input: Value, indices: Value) -> Value: + ... + def select_int(self, input: Value, indices: Value) -> Value: + ... + def select_vector(self, input: Value, indices: Value) -> Value: + ... + def set_current_stream(self, arg0: typing.SupportsInt) -> None: + ... + def sigmoid(self, in1: Value) -> Value: + ... + def softmax(self, input: Value) -> Value: + ... + def softmax_prior(self, input: Value, prior: Value) -> Value: + ... + def softplus(self, in1: Value) -> Value: + ... + def sqrt(self, in1: Value) -> Value: + ... + def square(self, in1: Value) -> Value: + ... + def squeeze(self, input: Value) -> Value: + ... + def stable_invariant(self, r: Value, mass: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def stable_invariant_inverse(self, s: Value, mass: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def stable_invariant_nu(self, r: Value, mass: Value, nu: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def stable_invariant_nu_inverse(self, s: Value, mass: Value, nu: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def stack(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def stack_sizes(self, args: collections.abc.Sequence[Value]) -> Value: + ... + def sub(self, in1: Value, in2: Value) -> Value: + ... + def subchannel_weights(self, invariants: Value, masses: Value, widths: Value, indices: Value, on_shell: Value, group_sizes: Value) -> Value: + ... + def t1_inv_min_max_doublet(self, pa: Value, pb: Value, m1: Value, mir_min: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def t1_inv_value_and_min_max_doublet(self, pa: Value, pb: Value, p1: Value, m1: Value, mir_min: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def t2_inv_min_max_doublet(self, pa: Value, pb: Value, m1: Value, mir_min: Value, t1_abs: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def t2_inv_value_and_min_max_doublet(self, pa: Value, pb: Value, p1: Value, m1: Value, mir_min: Value, t1_abs: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def t_inv_min_max(self, pa: Value, pb: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def t_inv_value_and_min_max(self, pa: Value, pb: Value, p1: Value, p2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def three_body_decay(self, r_e1: Value, r_e2: Value, r_phi: Value, r_cos_theta: Value, r_beta: Value, m0: Value, m1: Value, m2: Value, m3: Value, p0: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def three_body_decay_com(self, r_e1: Value, r_e2: Value, r_phi: Value, r_cos_theta: Value, r_beta: Value, m0: Value, m1: Value, m2: Value, m3: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def three_body_decay_com_inverse(self, p1: Value, p2: Value, p3: Value) -> typing.Annotated[list[Value], "FixedSize(10)"]: + ... + def three_body_decay_inverse(self, p1: Value, p2: Value, p3: Value) -> typing.Annotated[list[Value], "FixedSize(11)"]: + ... + def two_body_decay(self, r_phi: Value, r_cos_theta: Value, m0: Value, m1: Value, m2: Value, p0: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_body_decay_com(self, r_phi: Value, r_cos_theta: Value, m0: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_body_decay_com_inverse(self, p1: Value, p2: Value) -> typing.Annotated[list[Value], "FixedSize(6)"]: + ... + def two_body_decay_inverse(self, p1: Value, p2: Value) -> typing.Annotated[list[Value], "FixedSize(7)"]: + ... + def two_to_three_particle_scattering(self, phi_choice: Value, pa: Value, pb: Value, p3: Value, s23: Value, t1_abs: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_to_three_particle_scattering_inverse(self, p1: Value, p2: Value, p3: Value, pa: Value, pb: Value, t1_abs: Value, s23: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def two_to_two_particle_scattering(self, r_phi: Value, pa: Value, pb: Value, t_abs: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_to_two_particle_scattering_com(self, r_phi: Value, pa: Value, pb: Value, t_abs: Value, m1: Value, m2: Value) -> typing.Annotated[list[Value], "FixedSize(3)"]: + ... + def two_to_two_particle_scattering_com_inverse(self, p1: Value, p2: Value, pa: Value, pb: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def two_to_two_particle_scattering_inverse(self, p1: Value, p2: Value, pa: Value, pb: Value) -> typing.Annotated[list[Value], "FixedSize(4)"]: + ... + def uniform_invariant(self, r: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def uniform_invariant_inverse(self, s: Value, s_min: Value, s_max: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def unsqueeze(self, input: Value) -> Value: + ... + def unstack(self, in1: Value) -> list[Value]: + ... + def unstack_sizes(self, in1: Value) -> list[Value]: + ... + def unweight(self, weights: Value, max_weight: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def vegas_forward(self, input: Value, grid: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def vegas_histogram(self, input: Value, weights: Value, bin_count: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def vegas_inverse(self, input: Value, grid: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... + def x1x2_to_r(self, x1: Value, x2: Value, s_lab: Value) -> typing.Annotated[list[Value], "FixedSize(2)"]: + ... class FunctionGenerator: - def __call__(self, *args, **kwargs): ... - def __init__( - self, name: str, arg_types: NamedTypes, return_types: NamedTypes - ) -> None: ... + def __call__(self, *args, **kwargs): + ... + def __init__(self, name: str, arg_types: NamedTypes, return_types: NamedTypes) -> None: + ... @typing.overload - def build_function( - self, builder: FunctionBuilder, args: collections.abc.Sequence[Value] - ) -> NamedValues: ... + def build_function(self, builder: FunctionBuilder, args: collections.abc.Sequence[Value]) -> NamedValues: + ... @typing.overload - def build_function( - self, builder: FunctionBuilder, args: NamedValues - ) -> NamedValues: ... - def function(self) -> Function: ... - + def build_function(self, builder: FunctionBuilder, args: NamedValues) -> NamedValues: + ... + def function(self) -> Function: + ... class FunctionRuntime: - def __call__(self, *args): ... + def __call__(self, *args): + ... @typing.overload - def __init__(self, function: Function) -> None: ... + def __init__(self, function: Function) -> None: + ... @typing.overload - def __init__(self, function: Function, context: Context) -> None: ... - def call(self, arg0: collections.abc.Sequence[typing.Any]) -> list[Tensor]: ... - def call_backward( - self, - arg0: collections.abc.Sequence[typing.Any], - arg1: collections.abc.Sequence[typing.Any], - arg2: collections.abc.Sequence[bool], - ) -> tuple[ - list[madspace._madspace_py.Tensor | None], - list[tuple[str, madspace._madspace_py.Tensor | None]], - ]: ... - def call_with_grad( - self, - arg0: collections.abc.Sequence[typing.Any], - arg1: collections.abc.Sequence[bool], - ) -> tuple[list[Tensor], list[madspace._madspace_py.Tensor | None], list[bool]]: ... - + def __init__(self, function: Function, context: Context) -> None: + ... + def call(self, arg0: collections.abc.Sequence[typing.Any]) -> list[Tensor]: + ... + def call_backward(self, arg0: collections.abc.Sequence[typing.Any], arg1: collections.abc.Sequence[typing.Any], arg2: collections.abc.Sequence[bool]) -> tuple[list[madspace._madspace_py.Tensor | None], list[tuple[str, madspace._madspace_py.Tensor | None]]]: + ... + def call_with_grad(self, arg0: collections.abc.Sequence[typing.Any], arg1: collections.abc.Sequence[bool]) -> tuple[list[Tensor], list[madspace._madspace_py.Tensor | None], list[bool]]: + ... class GeneratorConfig: verbosity: Verbosity write_live_data: bool - def __init__(self) -> None: ... + def __init__(self) -> None: + ... @property - def combine_thread_count(self) -> int: ... + def combine_thread_count(self) -> int: + ... @combine_thread_count.setter - def combine_thread_count(self, arg0: typing.SupportsInt) -> None: ... + def combine_thread_count(self, arg0: typing.SupportsInt) -> None: + ... @property - def cpu_batch_size(self) -> int: ... + def cpu_batch_size(self) -> int: + ... @cpu_batch_size.setter - def cpu_batch_size(self, arg0: typing.SupportsInt) -> None: ... + def cpu_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def freeze_max_weight_after(self) -> int: ... + def freeze_max_weight_after(self) -> int: + ... @freeze_max_weight_after.setter - def freeze_max_weight_after(self, arg0: typing.SupportsInt) -> None: ... + def freeze_max_weight_after(self, arg0: typing.SupportsInt) -> None: + ... @property - def gpu_batch_size(self) -> int: ... + def gpu_batch_size(self) -> int: + ... @gpu_batch_size.setter - def gpu_batch_size(self, arg0: typing.SupportsInt) -> None: ... + def gpu_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def max_batch_size(self) -> int: ... + def max_batch_size(self) -> int: + ... @max_batch_size.setter - def max_batch_size(self, arg0: typing.SupportsInt) -> None: ... + def max_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def max_overweight_truncation(self) -> float: ... + def max_overweight_truncation(self) -> float: + ... @max_overweight_truncation.setter - def max_overweight_truncation(self, arg0: typing.SupportsFloat) -> None: ... + def max_overweight_truncation(self, arg0: typing.SupportsFloat) -> None: + ... @property - def optimization_patience(self) -> int: ... + def optimization_patience(self) -> int: + ... @optimization_patience.setter - def optimization_patience(self, arg0: typing.SupportsInt) -> None: ... + def optimization_patience(self, arg0: typing.SupportsInt) -> None: + ... @property - def optimization_threshold(self) -> float: ... + def optimization_threshold(self) -> float: + ... @optimization_threshold.setter - def optimization_threshold(self, arg0: typing.SupportsFloat) -> None: ... + def optimization_threshold(self, arg0: typing.SupportsFloat) -> None: + ... @property - def start_batch_size(self) -> int: ... + def start_batch_size(self) -> int: + ... @start_batch_size.setter - def start_batch_size(self, arg0: typing.SupportsInt) -> None: ... + def start_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def survey_max_iters(self) -> int: ... + def survey_max_iters(self) -> int: + ... @survey_max_iters.setter - def survey_max_iters(self, arg0: typing.SupportsInt) -> None: ... + def survey_max_iters(self, arg0: typing.SupportsInt) -> None: + ... @property - def survey_min_iters(self) -> int: ... + def survey_min_iters(self) -> int: + ... @survey_min_iters.setter - def survey_min_iters(self, arg0: typing.SupportsInt) -> None: ... + def survey_min_iters(self, arg0: typing.SupportsInt) -> None: + ... @property - def survey_target_precision(self) -> float: ... + def survey_target_precision(self) -> float: + ... @survey_target_precision.setter - def survey_target_precision(self, arg0: typing.SupportsFloat) -> None: ... + def survey_target_precision(self, arg0: typing.SupportsFloat) -> None: + ... @property - def target_count(self) -> int: ... + def target_count(self) -> int: + ... @target_count.setter - def target_count(self, arg0: typing.SupportsInt) -> None: ... + def target_count(self, arg0: typing.SupportsInt) -> None: + ... @property - def vegas_damping(self) -> float: ... + def vegas_damping(self) -> float: + ... @vegas_damping.setter - def vegas_damping(self, arg0: typing.SupportsFloat) -> None: ... - + def vegas_damping(self, arg0: typing.SupportsFloat) -> None: + ... class GeneratorStatus: done: bool name: str optimized: bool - def __init__(self) -> None: ... + def __init__(self) -> None: + ... @property - def count(self) -> int: ... + def count(self) -> int: + ... @count.setter - def count(self, arg0: typing.SupportsInt) -> None: ... + def count(self, arg0: typing.SupportsInt) -> None: + ... @property - def count_after_cuts(self) -> int: ... + def count_after_cuts(self) -> int: + ... @count_after_cuts.setter - def count_after_cuts(self, arg0: typing.SupportsInt) -> None: ... + def count_after_cuts(self, arg0: typing.SupportsInt) -> None: + ... @property - def count_after_cuts_opt(self) -> int: ... + def count_after_cuts_opt(self) -> int: + ... @count_after_cuts_opt.setter - def count_after_cuts_opt(self, arg0: typing.SupportsInt) -> None: ... + def count_after_cuts_opt(self, arg0: typing.SupportsInt) -> None: + ... @property - def count_opt(self) -> int: ... + def count_opt(self) -> int: + ... @count_opt.setter - def count_opt(self, arg0: typing.SupportsInt) -> None: ... + def count_opt(self, arg0: typing.SupportsInt) -> None: + ... @property - def count_target(self) -> float: ... + def count_target(self) -> float: + ... @count_target.setter - def count_target(self, arg0: typing.SupportsFloat) -> None: ... + def count_target(self, arg0: typing.SupportsFloat) -> None: + ... @property - def count_unweighted(self) -> float: ... + def count_unweighted(self) -> float: + ... @count_unweighted.setter - def count_unweighted(self, arg0: typing.SupportsFloat) -> None: ... + def count_unweighted(self, arg0: typing.SupportsFloat) -> None: + ... @property - def error(self) -> float: ... + def error(self) -> float: + ... @error.setter - def error(self, arg0: typing.SupportsFloat) -> None: ... + def error(self, arg0: typing.SupportsFloat) -> None: + ... @property - def iterations(self) -> int: ... + def iterations(self) -> int: + ... @iterations.setter - def iterations(self, arg0: typing.SupportsInt) -> None: ... + def iterations(self, arg0: typing.SupportsInt) -> None: + ... @property - def mean(self) -> float: ... + def mean(self) -> float: + ... @mean.setter - def mean(self, arg0: typing.SupportsFloat) -> None: ... + def mean(self, arg0: typing.SupportsFloat) -> None: + ... @property - def rel_std_dev(self) -> float: ... + def rel_std_dev(self) -> float: + ... @rel_std_dev.setter - def rel_std_dev(self, arg0: typing.SupportsFloat) -> None: ... + def rel_std_dev(self, arg0: typing.SupportsFloat) -> None: + ... @property - def subprocess(self) -> int: ... + def subprocess(self) -> int: + ... @subprocess.setter - def subprocess(self, arg0: typing.SupportsInt) -> None: ... - + def subprocess(self, arg0: typing.SupportsInt) -> None: + ... class HistItem: - def __init__( - self, - observable: Observable, - min: typing.SupportsFloat, - max: typing.SupportsFloat, - bin_count: typing.SupportsInt, - ) -> None: ... + def __init__(self, observable: Observable, min: typing.SupportsFloat, max: typing.SupportsFloat, bin_count: typing.SupportsInt) -> None: + ... @property - def bin_count(self) -> int: ... + def bin_count(self) -> int: + ... @property - def max(self) -> float: ... + def max(self) -> float: + ... @property - def min(self) -> float: ... + def min(self) -> float: + ... @property - def observable(self) -> Observable: ... - + def observable(self) -> Observable: + ... class Histogram: @property - def bin_errors(self) -> list[float]: ... + def bin_errors(self) -> list[float]: + ... @property - def bin_values(self) -> list[float]: ... + def bin_values(self) -> list[float]: + ... @property - def max(self) -> float: ... + def max(self) -> float: + ... @property - def min(self) -> float: ... + def min(self) -> float: + ... @property - def name(self) -> str: ... - + def name(self) -> str: + ... class Instruction: - def __str__(self) -> str: ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def opcode(self) -> int: ... - + def opcode(self) -> int: + ... class InstructionCall: - def __repr__(self) -> str: ... - def __str__(self) -> str: ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... @property - def inputs(self) -> list[Value]: ... + def inputs(self) -> list[Value]: + ... @property - def instruction(self) -> Instruction: ... + def instruction(self) -> Instruction: + ... @property - def outputs(self) -> list[Value]: ... - + def outputs(self) -> list[Value]: + ... class Integrand(FunctionGenerator): drop_cuts_and_rescale: typing.ClassVar[int] = 8192 exclude_adaptive_and_chan_weight: typing.ClassVar[int] = 4096 - matrix_element_inputs: typing.ClassVar[ - list - ] # value = [, , , , , ] - matrix_element_outputs: typing.ClassVar[ - list - ] # value = [, , , , ] + matrix_element_inputs: typing.ClassVar[list] # value = [, , , , , ] + matrix_element_outputs: typing.ClassVar[list] # value = [, , , , ] return_chan_weights: typing.ClassVar[int] = 256 return_channel: typing.ClassVar[int] = 128 return_cwnet_input: typing.ClassVar[int] = 512 @@ -1220,421 +1069,419 @@ class Integrand(FunctionGenerator): return_x1_x2: typing.ClassVar[int] = 8 sample: typing.ClassVar[int] = 1 unweight: typing.ClassVar[int] = 2 - def __init__( - self, - mapping: PhaseSpaceMapping, - diff_xs: DifferentialCrossSection, - adaptive_map: ( - None | madspace._madspace_py.VegasMapping | madspace._madspace_py.Flow - ) = None, - discrete_before: ( - None - | madspace._madspace_py.DiscreteSampler - | madspace._madspace_py.DiscreteFlow - ) = None, - discrete_after: ( - None - | madspace._madspace_py.DiscreteSampler - | madspace._madspace_py.DiscreteFlow - ) = None, - pdf_grid: madspace._madspace_py.PdfGrid | None = None, - energy_scale: madspace._madspace_py.EnergyScale | None = None, - prop_chan_weights: madspace._madspace_py.PropagatorChannelWeights | None = None, - subchan_weights: madspace._madspace_py.SubchannelWeights | None = None, - chan_weight_net: madspace._madspace_py.ChannelWeightNetwork | None = None, - chan_weight_remap: collections.abc.Sequence[typing.SupportsInt] = [], - remapped_chan_count: typing.SupportsInt = 0, - flags: typing.SupportsInt = 0, - channel_indices: collections.abc.Sequence[typing.SupportsInt] = [], - active_flavors: collections.abc.Sequence[typing.SupportsInt] = [], - flavor_remap: collections.abc.Sequence[typing.SupportsInt] = [], - flavor_factors: collections.abc.Sequence[typing.SupportsFloat] = [], - ) -> None: ... - def adaptive_map( - self, - ) -> None | madspace._madspace_py.VegasMapping | madspace._madspace_py.Flow: ... - def chan_weight_net(self) -> madspace._madspace_py.ChannelWeightNetwork | None: ... - def diff_xs(self) -> DifferentialCrossSection: ... - def discrete_after( - self, - ) -> ( - None - | madspace._madspace_py.DiscreteSampler - | madspace._madspace_py.DiscreteFlow - ): ... - def discrete_before( - self, - ) -> ( - None - | madspace._madspace_py.DiscreteSampler - | madspace._madspace_py.DiscreteFlow - ): ... - def energy_scale(self) -> madspace._madspace_py.EnergyScale | None: ... - def flags(self) -> int: ... - def latent_dims(self) -> tuple[list[int], list[bool]]: ... - def mapping(self) -> PhaseSpaceMapping: ... - def particle_count(self) -> int: ... - def prop_chan_weights( - self, - ) -> madspace._madspace_py.PropagatorChannelWeights | None: ... - def random_dim(self) -> int: ... - def vegas_grid_name(self) -> str | None: ... - + def __init__(self, mapping: PhaseSpaceMapping, diff_xs: DifferentialCrossSection, adaptive_map: None | madspace._madspace_py.VegasMapping | madspace._madspace_py.Flow = None, discrete_before: None | madspace._madspace_py.DiscreteSampler | madspace._madspace_py.DiscreteFlow = None, discrete_after: None | madspace._madspace_py.DiscreteSampler | madspace._madspace_py.DiscreteFlow = None, pdf_grid: madspace._madspace_py.PdfGrid | None = None, energy_scale: madspace._madspace_py.EnergyScale | None = None, prop_chan_weights: madspace._madspace_py.PropagatorChannelWeights | None = None, subchan_weights: madspace._madspace_py.SubchannelWeights | None = None, chan_weight_net: madspace._madspace_py.ChannelWeightNetwork | None = None, chan_weight_remap: collections.abc.Sequence[typing.SupportsInt] = [], remapped_chan_count: typing.SupportsInt = 0, flags: typing.SupportsInt = 0, channel_indices: collections.abc.Sequence[typing.SupportsInt] = [], active_flavors: collections.abc.Sequence[typing.SupportsInt] = [], flavor_remap: collections.abc.Sequence[typing.SupportsInt] = [], flavor_factors: collections.abc.Sequence[typing.SupportsFloat] = []) -> None: + ... + def adaptive_map(self) -> None | madspace._madspace_py.VegasMapping | madspace._madspace_py.Flow: + ... + def chan_weight_net(self) -> madspace._madspace_py.ChannelWeightNetwork | None: + ... + def diff_xs(self) -> DifferentialCrossSection: + ... + def discrete_after(self) -> None | madspace._madspace_py.DiscreteSampler | madspace._madspace_py.DiscreteFlow: + ... + def discrete_before(self) -> None | madspace._madspace_py.DiscreteSampler | madspace._madspace_py.DiscreteFlow: + ... + def energy_scale(self) -> madspace._madspace_py.EnergyScale | None: + ... + def flags(self) -> int: + ... + def latent_dims(self) -> tuple[list[int], list[bool]]: + ... + def mapping(self) -> PhaseSpaceMapping: + ... + def particle_count(self) -> int: + ... + def prop_chan_weights(self) -> madspace._madspace_py.PropagatorChannelWeights | None: + ... + def random_dim(self) -> int: + ... + def vegas_grid_name(self) -> str | None: + ... class IntegrandProbability(FunctionGenerator): - def __init__(self, integrand: Integrand) -> None: ... - + def __init__(self, integrand: Integrand) -> None: + ... class Invariant(Mapping): - def __init__( - self, - power: typing.SupportsFloat = 0.0, - mass: typing.SupportsFloat = 0.0, - width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, power: typing.SupportsFloat = 0.0, mass: typing.SupportsFloat = 0.0, width: typing.SupportsFloat = 0.0) -> None: + ... class LHECompleter: @staticmethod - def load(file: str) -> LHECompleter: ... - def __init__( - self, - subproc_args: collections.abc.Sequence[SubprocArgs], - bw_cutoff: typing.SupportsFloat, - ) -> None: ... - def save(self, file: str) -> None: ... - @property - def max_particle_count(self) -> int: ... - + def load(file: str) -> LHECompleter: + ... + def __init__(self, subproc_args: collections.abc.Sequence[SubprocArgs], bw_cutoff: typing.SupportsFloat) -> None: + ... + def save(self, file: str) -> None: + ... + @property + def max_particle_count(self) -> int: + ... class LHEEvent: - def __init__( - self, - process_id: typing.SupportsInt = 0, - weight: typing.SupportsFloat = 0.0, - scale: typing.SupportsFloat = 0.0, - alpha_qed: typing.SupportsFloat = 0.0, - alpha_qcd: typing.SupportsFloat = 0.0, - particles: collections.abc.Sequence[LHEParticle] = [], - ) -> None: ... - @property - def alpha_qcd(self) -> int: ... + def __init__(self, process_id: typing.SupportsInt = 0, weight: typing.SupportsFloat = 0.0, scale: typing.SupportsFloat = 0.0, alpha_qed: typing.SupportsFloat = 0.0, alpha_qcd: typing.SupportsFloat = 0.0, particles: collections.abc.Sequence[LHEParticle] = []) -> None: + ... + @property + def alpha_qcd(self) -> int: + ... @alpha_qcd.setter - def alpha_qcd(self, arg0: typing.SupportsInt) -> None: ... + def alpha_qcd(self, arg0: typing.SupportsInt) -> None: + ... @property - def alpha_qed(self) -> float: ... + def alpha_qed(self) -> float: + ... @alpha_qed.setter - def alpha_qed(self, arg0: typing.SupportsFloat) -> None: ... + def alpha_qed(self, arg0: typing.SupportsFloat) -> None: + ... @property - def particles(self) -> list[LHEParticle]: ... + def particles(self) -> list[LHEParticle]: + ... @particles.setter - def particles(self, arg0: collections.abc.Sequence[LHEParticle]) -> None: ... + def particles(self, arg0: collections.abc.Sequence[LHEParticle]) -> None: + ... @property - def process_id(self) -> int: ... + def process_id(self) -> int: + ... @process_id.setter - def process_id(self, arg0: typing.SupportsInt) -> None: ... + def process_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def scale(self) -> float: ... + def scale(self) -> float: + ... @scale.setter - def scale(self, arg0: typing.SupportsFloat) -> None: ... + def scale(self, arg0: typing.SupportsFloat) -> None: + ... @property - def weight(self) -> float: ... + def weight(self) -> float: + ... @weight.setter - def weight(self, arg0: typing.SupportsFloat) -> None: ... - + def weight(self, arg0: typing.SupportsFloat) -> None: + ... class LHEFileWriter: - def __init__(self, file_name: str, meta: LHEMeta) -> None: ... - def write(self, event: LHEEvent) -> None: ... - def write_string(self, str: str) -> None: ... - + def __init__(self, file_name: str, meta: LHEMeta) -> None: + ... + def write(self, event: LHEEvent) -> None: + ... + def write_string(self, str: str) -> None: + ... class LHEHeader: content: str escape_content: bool name: str - def __init__( - self, name: str = "", content: str = "", escape_content: bool = False - ) -> None: ... - + def __init__(self, name: str = '', content: str = '', escape_content: bool = False) -> None: + ... class LHEMeta: - def __init__( - self, - beam1_pdg_id: typing.SupportsInt = 0, - beam2_pdg_id: typing.SupportsInt = 0, - beam1_energy: typing.SupportsFloat = 0.0, - beam2_energy: typing.SupportsFloat = 0.0, - beam1_pdf_authors: typing.SupportsInt = 0, - beam2_pdf_authors: typing.SupportsInt = 0, - beam1_pdf_id: typing.SupportsInt = 0, - beam2_pdf_id: typing.SupportsInt = 0, - weight_mode: typing.SupportsInt = 0, - processes: collections.abc.Sequence[LHEProcess] = [], - headers: collections.abc.Sequence[LHEHeader] = [], - ) -> None: ... - @property - def beam1_energy(self) -> float: ... + def __init__(self, beam1_pdg_id: typing.SupportsInt = 0, beam2_pdg_id: typing.SupportsInt = 0, beam1_energy: typing.SupportsFloat = 0.0, beam2_energy: typing.SupportsFloat = 0.0, beam1_pdf_authors: typing.SupportsInt = 0, beam2_pdf_authors: typing.SupportsInt = 0, beam1_pdf_id: typing.SupportsInt = 0, beam2_pdf_id: typing.SupportsInt = 0, weight_mode: typing.SupportsInt = 0, processes: collections.abc.Sequence[LHEProcess] = [], headers: collections.abc.Sequence[LHEHeader] = []) -> None: + ... + @property + def beam1_energy(self) -> float: + ... @beam1_energy.setter - def beam1_energy(self, arg0: typing.SupportsFloat) -> None: ... + def beam1_energy(self, arg0: typing.SupportsFloat) -> None: + ... @property - def beam1_pdf_authors(self) -> int: ... + def beam1_pdf_authors(self) -> int: + ... @beam1_pdf_authors.setter - def beam1_pdf_authors(self, arg0: typing.SupportsInt) -> None: ... + def beam1_pdf_authors(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam1_pdf_id(self) -> int: ... + def beam1_pdf_id(self) -> int: + ... @beam1_pdf_id.setter - def beam1_pdf_id(self, arg0: typing.SupportsInt) -> None: ... + def beam1_pdf_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam1_pdg_id(self) -> int: ... + def beam1_pdg_id(self) -> int: + ... @beam1_pdg_id.setter - def beam1_pdg_id(self, arg0: typing.SupportsInt) -> None: ... + def beam1_pdg_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam2_energy(self) -> float: ... + def beam2_energy(self) -> float: + ... @beam2_energy.setter - def beam2_energy(self, arg0: typing.SupportsFloat) -> None: ... + def beam2_energy(self, arg0: typing.SupportsFloat) -> None: + ... @property - def beam2_pdf_authors(self) -> int: ... + def beam2_pdf_authors(self) -> int: + ... @beam2_pdf_authors.setter - def beam2_pdf_authors(self, arg0: typing.SupportsInt) -> None: ... + def beam2_pdf_authors(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam2_pdf_id(self) -> int: ... + def beam2_pdf_id(self) -> int: + ... @beam2_pdf_id.setter - def beam2_pdf_id(self, arg0: typing.SupportsInt) -> None: ... + def beam2_pdf_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def beam2_pdg_id(self) -> int: ... + def beam2_pdg_id(self) -> int: + ... @beam2_pdg_id.setter - def beam2_pdg_id(self, arg0: typing.SupportsInt) -> None: ... + def beam2_pdg_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def headers(self) -> list[LHEHeader]: ... + def headers(self) -> list[LHEHeader]: + ... @headers.setter - def headers(self, arg0: collections.abc.Sequence[LHEHeader]) -> None: ... + def headers(self, arg0: collections.abc.Sequence[LHEHeader]) -> None: + ... @property - def processes(self) -> list[LHEProcess]: ... + def processes(self) -> list[LHEProcess]: + ... @processes.setter - def processes(self, arg0: collections.abc.Sequence[LHEProcess]) -> None: ... + def processes(self, arg0: collections.abc.Sequence[LHEProcess]) -> None: + ... @property - def weight_mode(self) -> int: ... + def weight_mode(self) -> int: + ... @weight_mode.setter - def weight_mode(self, arg0: typing.SupportsInt) -> None: ... - + def weight_mode(self, arg0: typing.SupportsInt) -> None: + ... class LHEParticle: status_incoming: typing.ClassVar[int] = -1 status_intermediate_resonance: typing.ClassVar[int] = 2 status_outgoing: typing.ClassVar[int] = 1 - def __init__( - self, - pdg_id: typing.SupportsInt = 0, - status_code: typing.SupportsInt = 0, - mother1: typing.SupportsInt = 0, - mother2: typing.SupportsInt = 0, - color: typing.SupportsInt = 0, - anti_color: typing.SupportsInt = 0, - p_x: typing.SupportsFloat = 0.0, - p_y: typing.SupportsFloat = 0.0, - p_z: typing.SupportsFloat = 0.0, - energy: typing.SupportsFloat = 0.0, - mass: typing.SupportsFloat = 0.0, - lifetime: typing.SupportsFloat = 0.0, - spin: typing.SupportsFloat = 0.0, - ) -> None: ... - @property - def anti_color(self) -> int: ... + def __init__(self, pdg_id: typing.SupportsInt = 0, status_code: typing.SupportsInt = 0, mother1: typing.SupportsInt = 0, mother2: typing.SupportsInt = 0, color: typing.SupportsInt = 0, anti_color: typing.SupportsInt = 0, p_x: typing.SupportsFloat = 0.0, p_y: typing.SupportsFloat = 0.0, p_z: typing.SupportsFloat = 0.0, energy: typing.SupportsFloat = 0.0, mass: typing.SupportsFloat = 0.0, lifetime: typing.SupportsFloat = 0.0, spin: typing.SupportsFloat = 0.0) -> None: + ... + @property + def anti_color(self) -> int: + ... @anti_color.setter - def anti_color(self, arg0: typing.SupportsInt) -> None: ... + def anti_color(self, arg0: typing.SupportsInt) -> None: + ... @property - def color(self) -> int: ... + def color(self) -> int: + ... @color.setter - def color(self, arg0: typing.SupportsInt) -> None: ... + def color(self, arg0: typing.SupportsInt) -> None: + ... @property - def energy(self) -> float: ... + def energy(self) -> float: + ... @energy.setter - def energy(self, arg0: typing.SupportsFloat) -> None: ... + def energy(self, arg0: typing.SupportsFloat) -> None: + ... @property - def lifetime(self) -> float: ... + def lifetime(self) -> float: + ... @lifetime.setter - def lifetime(self, arg0: typing.SupportsFloat) -> None: ... + def lifetime(self, arg0: typing.SupportsFloat) -> None: + ... @property - def mass(self) -> float: ... + def mass(self) -> float: + ... @mass.setter - def mass(self, arg0: typing.SupportsFloat) -> None: ... + def mass(self, arg0: typing.SupportsFloat) -> None: + ... @property - def mother1(self) -> int: ... + def mother1(self) -> int: + ... @mother1.setter - def mother1(self, arg0: typing.SupportsInt) -> None: ... + def mother1(self, arg0: typing.SupportsInt) -> None: + ... @property - def mother2(self) -> int: ... + def mother2(self) -> int: + ... @mother2.setter - def mother2(self, arg0: typing.SupportsInt) -> None: ... + def mother2(self, arg0: typing.SupportsInt) -> None: + ... @property - def pdg_id(self) -> int: ... + def pdg_id(self) -> int: + ... @pdg_id.setter - def pdg_id(self, arg0: typing.SupportsInt) -> None: ... + def pdg_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def px(self) -> float: ... + def px(self) -> float: + ... @px.setter - def px(self, arg0: typing.SupportsFloat) -> None: ... + def px(self, arg0: typing.SupportsFloat) -> None: + ... @property - def py(self) -> float: ... + def py(self) -> float: + ... @py.setter - def py(self, arg0: typing.SupportsFloat) -> None: ... + def py(self, arg0: typing.SupportsFloat) -> None: + ... @property - def pz(self) -> float: ... + def pz(self) -> float: + ... @pz.setter - def pz(self, arg0: typing.SupportsFloat) -> None: ... + def pz(self, arg0: typing.SupportsFloat) -> None: + ... @property - def spin(self) -> float: ... + def spin(self) -> float: + ... @spin.setter - def spin(self, arg0: typing.SupportsFloat) -> None: ... + def spin(self, arg0: typing.SupportsFloat) -> None: + ... @property - def status_code(self) -> int: ... + def status_code(self) -> int: + ... @status_code.setter - def status_code(self, arg0: typing.SupportsInt) -> None: ... - + def status_code(self, arg0: typing.SupportsInt) -> None: + ... class LHEProcess: @staticmethod - def __init__(*args, **kwargs) -> None: ... + def __init__(*args, **kwargs) -> None: + ... @property - def cross_section(self) -> float: ... + def cross_section(self) -> float: + ... @cross_section.setter - def cross_section(self, arg0: typing.SupportsFloat) -> None: ... + def cross_section(self, arg0: typing.SupportsFloat) -> None: + ... @property - def cross_section_error(self) -> float: ... + def cross_section_error(self) -> float: + ... @cross_section_error.setter - def cross_section_error(self, arg0: typing.SupportsFloat) -> None: ... + def cross_section_error(self, arg0: typing.SupportsFloat) -> None: + ... @property - def max_weight(self) -> float: ... + def max_weight(self) -> float: + ... @max_weight.setter - def max_weight(self, arg0: typing.SupportsFloat) -> None: ... + def max_weight(self, arg0: typing.SupportsFloat) -> None: + ... @property - def process_id(self) -> int: ... + def process_id(self) -> int: + ... @process_id.setter - def process_id(self, arg0: typing.SupportsInt) -> None: ... - + def process_id(self, arg0: typing.SupportsInt) -> None: + ... class LineRef: - def __init__(self, str: str) -> None: ... - def __repr__(self) -> str: ... - + def __init__(self, str: str) -> None: + ... + def __repr__(self) -> str: + ... class Logger: class LogLevel: """ Members: - + level_debug - + level_info - + level_warning - + level_error """ - - __members__: typing.ClassVar[ - dict[str, Logger.LogLevel] - ] # value = {'level_debug': , 'level_info': , 'level_warning': , 'level_error': } - level_debug: typing.ClassVar[ - Logger.LogLevel - ] # value = - level_error: typing.ClassVar[ - Logger.LogLevel - ] # value = + __members__: typing.ClassVar[dict[str, Logger.LogLevel]] # value = {'level_debug': , 'level_info': , 'level_warning': , 'level_error': } + level_debug: typing.ClassVar[Logger.LogLevel] # value = + level_error: typing.ClassVar[Logger.LogLevel] # value = level_info: typing.ClassVar[Logger.LogLevel] # value = - level_warning: typing.ClassVar[ - Logger.LogLevel - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + level_warning: typing.ClassVar[Logger.LogLevel] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... level_debug: typing.ClassVar[Logger.LogLevel] # value = level_error: typing.ClassVar[Logger.LogLevel] # value = level_info: typing.ClassVar[Logger.LogLevel] # value = - level_warning: typing.ClassVar[ - Logger.LogLevel - ] # value = + level_warning: typing.ClassVar[Logger.LogLevel] # value = @staticmethod - def debug(message: str) -> None: ... + def debug(message: str) -> None: + ... @staticmethod - def error(message: str) -> None: ... + def error(message: str) -> None: + ... @staticmethod - def info(message: str) -> None: ... + def info(message: str) -> None: + ... @staticmethod - def log(level: Logger.LogLevel, message: str) -> None: ... + def log(level: Logger.LogLevel, message: str) -> None: + ... @staticmethod - def set_log_handler( - func: collections.abc.Callable[[Logger.LogLevel, str], None], - ) -> None: ... + def set_log_handler(func: collections.abc.Callable[[Logger.LogLevel, str], None]) -> None: + ... @staticmethod - def warning(message: str) -> None: ... - + def warning(message: str) -> None: + ... class Luminosity(Mapping): - def __init__( - self, - s_lab: typing.SupportsFloat, - s_hat_min: typing.SupportsFloat, - s_hat_max: typing.SupportsFloat = 0.0, - invariant_power: typing.SupportsFloat = 0.0, - mass: typing.SupportsFloat = 0.0, - width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, s_lab: typing.SupportsFloat, s_hat_min: typing.SupportsFloat, s_hat_max: typing.SupportsFloat = 0.0, invariant_power: typing.SupportsFloat = 0.0, mass: typing.SupportsFloat = 0.0, width: typing.SupportsFloat = 0.0) -> None: + ... class MLP(FunctionGenerator): class Activation: """ Members: - + relu - + leaky_relu - + elu - + gelu - + sigmoid - + softplus - + linear """ - - __members__: typing.ClassVar[ - dict[str, MLP.Activation] - ] # value = {'relu': , 'leaky_relu': , 'elu': , 'gelu': , 'sigmoid': , 'softplus': , 'linear': } + __members__: typing.ClassVar[dict[str, MLP.Activation]] # value = {'relu': , 'leaky_relu': , 'elu': , 'gelu': , 'sigmoid': , 'softplus': , 'linear': } elu: typing.ClassVar[MLP.Activation] # value = gelu: typing.ClassVar[MLP.Activation] # value = - leaky_relu: typing.ClassVar[ - MLP.Activation - ] # value = + leaky_relu: typing.ClassVar[MLP.Activation] # value = linear: typing.ClassVar[MLP.Activation] # value = relu: typing.ClassVar[MLP.Activation] # value = sigmoid: typing.ClassVar[MLP.Activation] # value = softplus: typing.ClassVar[MLP.Activation] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... elu: typing.ClassVar[MLP.Activation] # value = gelu: typing.ClassVar[MLP.Activation] # value = leaky_relu: typing.ClassVar[MLP.Activation] # value = @@ -1642,1099 +1489,926 @@ class MLP(FunctionGenerator): relu: typing.ClassVar[MLP.Activation] # value = sigmoid: typing.ClassVar[MLP.Activation] # value = softplus: typing.ClassVar[MLP.Activation] # value = - def __init__( - self, - input_dim: typing.SupportsInt, - output_dim: typing.SupportsInt, - hidden_dim: typing.SupportsInt = 32, - layers: typing.SupportsInt = 3, - activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, - prefix: str = "", - ) -> None: ... - def initialize_globals(self, context: Context) -> None: ... - def input_dim(self) -> int: ... - def output_dim(self) -> int: ... - + def __init__(self, input_dim: typing.SupportsInt, output_dim: typing.SupportsInt, hidden_dim: typing.SupportsInt = 32, layers: typing.SupportsInt = 3, activation: MLP.Activation = MLP.Activation.Activation.leaky_relu, prefix: str = '') -> None: + ... + def initialize_globals(self, context: Context) -> None: + ... + def input_dim(self) -> int: + ... + def output_dim(self) -> int: + ... class MadnisConfig: lr_schedule: AdamOptimizer.LRSchedule verbosity: Verbosity - def __init__(self) -> None: ... + def __init__(self) -> None: + ... @property - def adam_beta1(self) -> float: ... + def adam_beta1(self) -> float: + ... @adam_beta1.setter - def adam_beta1(self, arg0: typing.SupportsFloat) -> None: ... + def adam_beta1(self, arg0: typing.SupportsFloat) -> None: + ... @property - def adam_beta2(self) -> float: ... + def adam_beta2(self) -> float: + ... @adam_beta2.setter - def adam_beta2(self, arg0: typing.SupportsFloat) -> None: ... + def adam_beta2(self, arg0: typing.SupportsFloat) -> None: + ... @property - def adam_eps(self) -> float: ... + def adam_eps(self) -> float: + ... @adam_eps.setter - def adam_eps(self, arg0: typing.SupportsFloat) -> None: ... + def adam_eps(self, arg0: typing.SupportsFloat) -> None: + ... @property - def batch_size_offset(self) -> int: ... + def batch_size_offset(self) -> int: + ... @batch_size_offset.setter - def batch_size_offset(self, arg0: typing.SupportsInt) -> None: ... + def batch_size_offset(self, arg0: typing.SupportsInt) -> None: + ... @property - def batch_size_per_channel(self) -> int: ... + def batch_size_per_channel(self) -> int: + ... @batch_size_per_channel.setter - def batch_size_per_channel(self, arg0: typing.SupportsInt) -> None: ... + def batch_size_per_channel(self, arg0: typing.SupportsInt) -> None: + ... @property - def batches(self) -> int: ... + def batches(self) -> int: + ... @batches.setter - def batches(self, arg0: typing.SupportsInt) -> None: ... + def batches(self, arg0: typing.SupportsInt) -> None: + ... @property - def buffer_capacity(self) -> int: ... + def buffer_capacity(self) -> int: + ... @buffer_capacity.setter - def buffer_capacity(self, arg0: typing.SupportsInt) -> None: ... + def buffer_capacity(self, arg0: typing.SupportsInt) -> None: + ... @property - def buffer_unweighting_quantile(self) -> float: ... + def buffer_unweighting_quantile(self) -> float: + ... @buffer_unweighting_quantile.setter - def buffer_unweighting_quantile(self, arg0: typing.SupportsFloat) -> None: ... + def buffer_unweighting_quantile(self, arg0: typing.SupportsFloat) -> None: + ... @property - def buffered_steps(self) -> int: ... + def buffered_steps(self) -> int: + ... @buffered_steps.setter - def buffered_steps(self, arg0: typing.SupportsInt) -> None: ... + def buffered_steps(self, arg0: typing.SupportsInt) -> None: + ... @property - def channel_dropping_interval(self) -> int: ... + def channel_dropping_interval(self) -> int: + ... @channel_dropping_interval.setter - def channel_dropping_interval(self, arg0: typing.SupportsInt) -> None: ... + def channel_dropping_interval(self, arg0: typing.SupportsInt) -> None: + ... @property - def channel_dropping_threshold(self) -> float: ... + def channel_dropping_threshold(self) -> float: + ... @channel_dropping_threshold.setter - def channel_dropping_threshold(self, arg0: typing.SupportsFloat) -> None: ... + def channel_dropping_threshold(self, arg0: typing.SupportsFloat) -> None: + ... @property - def cpu_generator_batch_size(self) -> int: ... + def cpu_generator_batch_size(self) -> int: + ... @cpu_generator_batch_size.setter - def cpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: ... + def cpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def fixed_cwnet_fraction(self) -> float: ... + def fixed_cwnet_fraction(self) -> float: + ... @fixed_cwnet_fraction.setter - def fixed_cwnet_fraction(self, arg0: typing.SupportsFloat) -> None: ... + def fixed_cwnet_fraction(self, arg0: typing.SupportsFloat) -> None: + ... @property - def generator_target_size_factor(self) -> int: ... + def generator_target_size_factor(self) -> int: + ... @generator_target_size_factor.setter - def generator_target_size_factor(self, arg0: typing.SupportsInt) -> None: ... + def generator_target_size_factor(self, arg0: typing.SupportsInt) -> None: + ... @property - def gpu_generator_batch_granularity(self) -> int: ... + def gpu_generator_batch_granularity(self) -> int: + ... @gpu_generator_batch_granularity.setter - def gpu_generator_batch_granularity(self, arg0: typing.SupportsInt) -> None: ... + def gpu_generator_batch_granularity(self, arg0: typing.SupportsInt) -> None: + ... @property - def gpu_generator_batch_size(self) -> int: ... + def gpu_generator_batch_size(self) -> int: + ... @gpu_generator_batch_size.setter - def gpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: ... + def gpu_generator_batch_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def integration_history_length(self) -> int: ... + def integration_history_length(self) -> int: + ... @integration_history_length.setter - def integration_history_length(self, arg0: typing.SupportsInt) -> None: ... + def integration_history_length(self, arg0: typing.SupportsInt) -> None: + ... @property - def learning_rate(self) -> float: ... + def learning_rate(self) -> float: + ... @learning_rate.setter - def learning_rate(self, arg0: typing.SupportsFloat) -> None: ... + def learning_rate(self, arg0: typing.SupportsFloat) -> None: + ... @property - def log_interval(self) -> int: ... + def log_interval(self) -> int: + ... @log_interval.setter - def log_interval(self, arg0: typing.SupportsInt) -> None: ... + def log_interval(self, arg0: typing.SupportsInt) -> None: + ... @property - def minimum_buffer_size(self) -> int: ... + def minimum_buffer_size(self) -> int: + ... @minimum_buffer_size.setter - def minimum_buffer_size(self, arg0: typing.SupportsInt) -> None: ... + def minimum_buffer_size(self, arg0: typing.SupportsInt) -> None: + ... @property - def softclip_threshold(self) -> float: ... + def softclip_threshold(self) -> float: + ... @softclip_threshold.setter - def softclip_threshold(self, arg0: typing.SupportsFloat) -> None: ... + def softclip_threshold(self, arg0: typing.SupportsFloat) -> None: + ... @property - def uniform_channel_ratio(self) -> float: ... + def uniform_channel_ratio(self) -> float: + ... @uniform_channel_ratio.setter - def uniform_channel_ratio(self, arg0: typing.SupportsFloat) -> None: ... - + def uniform_channel_ratio(self, arg0: typing.SupportsFloat) -> None: + ... class MadnisLoss(FunctionGenerator): - def __init__( - self, - functions: collections.abc.Sequence[FunctionGenerator], - cwnet: madspace._madspace_py.ChannelWeightNetwork | None, - softclip_threshold: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, functions: collections.abc.Sequence[FunctionGenerator], cwnet: madspace._madspace_py.ChannelWeightNetwork | None, softclip_threshold: typing.SupportsFloat = 0.0) -> None: + ... class MadnisTraining: - def __init__( - self, - generator_context: Context, - optimizer_context: Context, - config: MadnisConfig, - integrands: collections.abc.Sequence[Integrand], - cwnet: madspace._madspace_py.ChannelWeightNetwork | None, - ) -> None: ... - def active_channel_count(self) -> int: ... - def active_channels(self) -> list[int]: ... - def train_step(self, batch_index: typing.SupportsInt) -> None: ... - + def __init__(self, generator_context: Context, optimizer_context: Context, config: MadnisConfig, integrands: collections.abc.Sequence[Integrand], cwnet: madspace._madspace_py.ChannelWeightNetwork | None) -> None: + ... + def active_channel_count(self) -> int: + ... + def active_channels(self) -> list[int]: + ... + def train_step(self, batch_index: typing.SupportsInt) -> None: + ... class Mapping: - def __init__( - self, - name: str, - input_types: NamedTypes, - output_types: NamedTypes, - condition_types: NamedTypes, - ) -> None: ... + def __init__(self, name: str, input_types: NamedTypes, output_types: NamedTypes, condition_types: NamedTypes) -> None: + ... @typing.overload - def build_forward( - self, - builder: FunctionBuilder, - inputs: collections.abc.Sequence[Value], - conditions: collections.abc.Sequence[Value], - ) -> NamedValues: ... + def build_forward(self, builder: FunctionBuilder, inputs: collections.abc.Sequence[Value], conditions: collections.abc.Sequence[Value]) -> NamedValues: + ... @typing.overload - def build_forward( - self, builder: FunctionBuilder, inputs: NamedValues, conditions: NamedValues - ) -> NamedValues: ... + def build_forward(self, builder: FunctionBuilder, inputs: NamedValues, conditions: NamedValues) -> NamedValues: + ... @typing.overload - def build_inverse( - self, - builder: FunctionBuilder, - inputs: collections.abc.Sequence[Value], - conditions: collections.abc.Sequence[Value], - ) -> NamedValues: ... + def build_inverse(self, builder: FunctionBuilder, inputs: collections.abc.Sequence[Value], conditions: collections.abc.Sequence[Value]) -> NamedValues: + ... @typing.overload - def build_inverse( - self, builder: FunctionBuilder, inputs: NamedValues, conditions: NamedValues - ) -> NamedValues: ... - def forward_function(self) -> Function: ... - def inverse_function(self) -> Function: ... - def map_forward(self, inputs=list(), conditions=list(), **kwargs): ... - def map_inverse(self, inputs=list(), conditions=list(), **kwargs): ... - + def build_inverse(self, builder: FunctionBuilder, inputs: NamedValues, conditions: NamedValues) -> NamedValues: + ... + def forward_function(self) -> Function: + ... + def inverse_function(self) -> Function: + ... + def map_forward(self, inputs = list(), conditions = list(), **kwargs): + ... + def map_inverse(self, inputs = list(), conditions = list(), **kwargs): + ... class MatrixElement(FunctionGenerator): class MatrixElementInput: """ Members: - + momenta_in - + alpha_s_in - + flavor_in - + random_color_in - + random_helicity_in - + random_diagram_in - + helicity_in - + diagram_in """ - - __members__: typing.ClassVar[ - dict[str, MatrixElement.MatrixElementInput] - ] # value = {'momenta_in': , 'alpha_s_in': , 'flavor_in': , 'random_color_in': , 'random_helicity_in': , 'random_diagram_in': , 'helicity_in': , 'diagram_in': } - alpha_s_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - diagram_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - flavor_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - helicity_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - momenta_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_color_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_diagram_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_helicity_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, MatrixElement.MatrixElementInput]] # value = {'momenta_in': , 'alpha_s_in': , 'flavor_in': , 'random_color_in': , 'random_helicity_in': , 'random_diagram_in': , 'helicity_in': , 'diagram_in': } + alpha_s_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + diagram_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + flavor_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + helicity_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + momenta_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_color_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_diagram_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_helicity_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... class MatrixElementOutput: """ Members: - + matrix_element_out - + diagram_amp2_out - + color_index_out - + helicity_index_out - + diagram_index_out """ - - __members__: typing.ClassVar[ - dict[str, MatrixElement.MatrixElementOutput] - ] # value = {'matrix_element_out': , 'diagram_amp2_out': , 'color_index_out': , 'helicity_index_out': , 'diagram_index_out': } - color_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - diagram_amp2_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - diagram_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - helicity_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - matrix_element_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, MatrixElement.MatrixElementOutput]] # value = {'matrix_element_out': , 'diagram_amp2_out': , 'color_index_out': , 'helicity_index_out': , 'diagram_index_out': } + color_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + diagram_amp2_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + diagram_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + helicity_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + matrix_element_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - - alpha_s_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - color_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - diagram_amp2_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - diagram_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - diagram_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - flavor_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - helicity_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - helicity_index_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - matrix_element_out: typing.ClassVar[ - MatrixElement.MatrixElementOutput - ] # value = - momenta_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_color_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_diagram_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = - random_helicity_in: typing.ClassVar[ - MatrixElement.MatrixElementInput - ] # value = + def value(self) -> int: + ... + alpha_s_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + color_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + diagram_amp2_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + diagram_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + diagram_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + flavor_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + helicity_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + helicity_index_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + matrix_element_out: typing.ClassVar[MatrixElement.MatrixElementOutput] # value = + momenta_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_color_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_diagram_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = + random_helicity_in: typing.ClassVar[MatrixElement.MatrixElementInput] # value = @typing.overload - def __init__( - self, - matrix_element_index: typing.SupportsInt, - particle_count: typing.SupportsInt, - inputs: collections.abc.Sequence[MatrixElement.MatrixElementInput], - outputs: collections.abc.Sequence[MatrixElement.MatrixElementOutput], - diagram_count: typing.SupportsInt = 1, - sample_random_inputs: bool = False, - ) -> None: ... + def __init__(self, matrix_element_index: typing.SupportsInt, particle_count: typing.SupportsInt, inputs: collections.abc.Sequence[MatrixElement.MatrixElementInput], outputs: collections.abc.Sequence[MatrixElement.MatrixElementOutput], diagram_count: typing.SupportsInt = 1, sample_random_inputs: bool = False) -> None: + ... @typing.overload - def __init__( - self, - matrix_element_api: MatrixElementApi, - inputs: collections.abc.Sequence[MatrixElement.MatrixElementInput], - outputs: collections.abc.Sequence[MatrixElement.MatrixElementOutput], - sample_random_inputs: bool = False, - ) -> None: ... - def diagram_count(self) -> int: ... - def matrix_element_index(self) -> int: ... - def particle_count(self) -> int: ... - + def __init__(self, matrix_element_api: MatrixElementApi, inputs: collections.abc.Sequence[MatrixElement.MatrixElementInput], outputs: collections.abc.Sequence[MatrixElement.MatrixElementOutput], sample_random_inputs: bool = False) -> None: + ... + def diagram_count(self) -> int: + ... + def matrix_element_index(self) -> int: + ... + def particle_count(self) -> int: + ... class MatrixElementApi: - def diagram_count(self) -> int: ... - def helicity_count(self) -> int: ... - def index(self) -> int: ... - def particle_count(self) -> int: ... - + def diagram_count(self) -> int: + ... + def helicity_count(self) -> int: + ... + def index(self) -> int: + ... + def particle_count(self) -> int: + ... class MomentumPreprocessing(FunctionGenerator): - def __init__(self, particle_count: typing.SupportsInt) -> None: ... - def output_dim(self) -> int: ... - + def __init__(self, particle_count: typing.SupportsInt) -> None: + ... + def output_dim(self) -> int: + ... class MultiChannelFunction(FunctionGenerator): - def __init__( - self, functions: collections.abc.Sequence[FunctionGenerator] - ) -> None: ... - + def __init__(self, functions: collections.abc.Sequence[FunctionGenerator]) -> None: + ... class MultiChannelIntegrand(FunctionGenerator): - def __init__( - self, - integrands: collections.abc.Sequence[Integrand], - return_sizes: bool = False, - ) -> None: ... - + def __init__(self, integrands: collections.abc.Sequence[Integrand], return_sizes: bool = False) -> None: + ... class MultiChannelMapping(Mapping): - def __init__(self, mappings: collections.abc.Sequence[Mapping]) -> None: ... - + def __init__(self, mappings: collections.abc.Sequence[Mapping]) -> None: + ... class MultiMadnisTraining: - def __init__( - self, - generator_context: Context, - optimizer_context: Context, - config: MadnisConfig, - integrands: collections.abc.Sequence[collections.abc.Sequence[Integrand]], - cwnets: collections.abc.Sequence[ - madspace._madspace_py.ChannelWeightNetwork | None - ], - ) -> None: ... - def active_channels(self) -> list[list[int]]: ... - def train(self) -> None: ... - + def __init__(self, generator_context: Context, optimizer_context: Context, config: MadnisConfig, integrands: collections.abc.Sequence[collections.abc.Sequence[Integrand]], cwnets: collections.abc.Sequence[madspace._madspace_py.ChannelWeightNetwork | None]) -> None: + ... + def active_channels(self) -> list[list[int]]: + ... + def train(self) -> None: + ... class NamedTypes: @typing.overload - def __getitem__(self, key: typing.SupportsInt) -> Type: ... + def __getitem__(self, key: typing.SupportsInt) -> Type: + ... @typing.overload - def __getitem__(self, key: str) -> Type: ... + def __getitem__(self, key: str) -> Type: + ... @typing.overload - def __init__(self) -> None: ... + def __init__(self) -> None: + ... @typing.overload - def __init__( - self, - keys: collections.abc.Sequence[str], - values: collections.abc.Sequence[Type], - ) -> None: ... + def __init__(self, keys: collections.abc.Sequence[str], values: collections.abc.Sequence[Type]) -> None: + ... @typing.overload - def __init__(self, items: collections.abc.Sequence[tuple[str, Type]]) -> None: ... - def __len__(self) -> int: ... - def index_map(self) -> dict[str, int]: ... - def keys(self) -> list[str]: ... - def push_back(self, name: str, item: Type) -> None: ... - def values(self) -> list[Type]: ... - + def __init__(self, items: collections.abc.Sequence[tuple[str, Type]]) -> None: + ... + def __len__(self) -> int: + ... + def index_map(self) -> dict[str, int]: + ... + def keys(self) -> list[str]: + ... + def push_back(self, name: str, item: Type) -> None: + ... + def values(self) -> list[Type]: + ... class NamedValues: @typing.overload - def __getitem__(self, key: typing.SupportsInt) -> Value: ... + def __getitem__(self, key: typing.SupportsInt) -> Value: + ... @typing.overload - def __getitem__(self, key: str) -> Value: ... + def __getitem__(self, key: str) -> Value: + ... @typing.overload - def __init__(self) -> None: ... + def __init__(self) -> None: + ... @typing.overload - def __init__( - self, - keys: collections.abc.Sequence[str], - values: collections.abc.Sequence[Value], - ) -> None: ... + def __init__(self, keys: collections.abc.Sequence[str], values: collections.abc.Sequence[Value]) -> None: + ... @typing.overload - def __init__(self, items: collections.abc.Sequence[tuple[str, Value]]) -> None: ... - def __len__(self) -> int: ... - def index_map(self) -> dict[str, int]: ... - def keys(self) -> list[str]: ... - def push_back(self, name: str, item: Value) -> None: ... - def values(self) -> list[Value]: ... - + def __init__(self, items: collections.abc.Sequence[tuple[str, Value]]) -> None: + ... + def __len__(self) -> int: + ... + def index_map(self) -> dict[str, int]: + ... + def keys(self) -> list[str]: + ... + def push_back(self, name: str, item: Value) -> None: + ... + def values(self) -> list[Value]: + ... class Observable(FunctionGenerator): class ObservableOption: """ Members: - + obs_e - + obs_px - + obs_py - + obs_pz - + obs_mass - + obs_pt - + obs_p_mag - + obs_phi - + obs_theta - + obs_y - + obs_y_abs - + obs_eta - + obs_eta_abs - + obs_delta_eta - + obs_delta_phi - + obs_delta_r - + obs_sqrt_s """ - - __members__: typing.ClassVar[ - dict[str, Observable.ObservableOption] - ] # value = {'obs_e': , 'obs_px': , 'obs_py': , 'obs_pz': , 'obs_mass': , 'obs_pt': , 'obs_p_mag': , 'obs_phi': , 'obs_theta': , 'obs_y': , 'obs_y_abs': , 'obs_eta': , 'obs_eta_abs': , 'obs_delta_eta': , 'obs_delta_phi': , 'obs_delta_r': , 'obs_sqrt_s': } - obs_delta_eta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_delta_phi: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_delta_r: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_e: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_eta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_eta_abs: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_mass: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_p_mag: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_phi: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_pt: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_px: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_py: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_pz: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_sqrt_s: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_theta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_y: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_y_abs: typing.ClassVar[ - Observable.ObservableOption - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, Observable.ObservableOption]] # value = {'obs_e': , 'obs_px': , 'obs_py': , 'obs_pz': , 'obs_mass': , 'obs_pt': , 'obs_p_mag': , 'obs_phi': , 'obs_theta': , 'obs_y': , 'obs_y_abs': , 'obs_eta': , 'obs_eta_abs': , 'obs_delta_eta': , 'obs_delta_phi': , 'obs_delta_r': , 'obs_sqrt_s': } + obs_delta_eta: typing.ClassVar[Observable.ObservableOption] # value = + obs_delta_phi: typing.ClassVar[Observable.ObservableOption] # value = + obs_delta_r: typing.ClassVar[Observable.ObservableOption] # value = + obs_e: typing.ClassVar[Observable.ObservableOption] # value = + obs_eta: typing.ClassVar[Observable.ObservableOption] # value = + obs_eta_abs: typing.ClassVar[Observable.ObservableOption] # value = + obs_mass: typing.ClassVar[Observable.ObservableOption] # value = + obs_p_mag: typing.ClassVar[Observable.ObservableOption] # value = + obs_phi: typing.ClassVar[Observable.ObservableOption] # value = + obs_pt: typing.ClassVar[Observable.ObservableOption] # value = + obs_px: typing.ClassVar[Observable.ObservableOption] # value = + obs_py: typing.ClassVar[Observable.ObservableOption] # value = + obs_pz: typing.ClassVar[Observable.ObservableOption] # value = + obs_sqrt_s: typing.ClassVar[Observable.ObservableOption] # value = + obs_theta: typing.ClassVar[Observable.ObservableOption] # value = + obs_y: typing.ClassVar[Observable.ObservableOption] # value = + obs_y_abs: typing.ClassVar[Observable.ObservableOption] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - + def value(self) -> int: + ... bottom_pids: typing.ClassVar[list] = [-5, 5] jet_pids: typing.ClassVar[list] = [1, 2, 3, 4, -1, -2, -3, -4, 21] lepton_pids: typing.ClassVar[list] = [11, 13, 15, -11, -13, -15] missing_pids: typing.ClassVar[list] = [12, 14, 16, -12, -14, -16] - obs_delta_eta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_delta_phi: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_delta_r: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_e: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_eta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_eta_abs: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_mass: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_p_mag: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_phi: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_pt: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_px: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_py: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_pz: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_sqrt_s: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_theta: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_y: typing.ClassVar[ - Observable.ObservableOption - ] # value = - obs_y_abs: typing.ClassVar[ - Observable.ObservableOption - ] # value = + obs_delta_eta: typing.ClassVar[Observable.ObservableOption] # value = + obs_delta_phi: typing.ClassVar[Observable.ObservableOption] # value = + obs_delta_r: typing.ClassVar[Observable.ObservableOption] # value = + obs_e: typing.ClassVar[Observable.ObservableOption] # value = + obs_eta: typing.ClassVar[Observable.ObservableOption] # value = + obs_eta_abs: typing.ClassVar[Observable.ObservableOption] # value = + obs_mass: typing.ClassVar[Observable.ObservableOption] # value = + obs_p_mag: typing.ClassVar[Observable.ObservableOption] # value = + obs_phi: typing.ClassVar[Observable.ObservableOption] # value = + obs_pt: typing.ClassVar[Observable.ObservableOption] # value = + obs_px: typing.ClassVar[Observable.ObservableOption] # value = + obs_py: typing.ClassVar[Observable.ObservableOption] # value = + obs_pz: typing.ClassVar[Observable.ObservableOption] # value = + obs_sqrt_s: typing.ClassVar[Observable.ObservableOption] # value = + obs_theta: typing.ClassVar[Observable.ObservableOption] # value = + obs_y: typing.ClassVar[Observable.ObservableOption] # value = + obs_y_abs: typing.ClassVar[Observable.ObservableOption] # value = photon_pids: typing.ClassVar[list] = [22] - def __init__( - self, - pids: collections.abc.Sequence[typing.SupportsInt], - observable: Observable.ObservableOption, - select_pids: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ], - sum_momenta: bool = False, - sum_observable: bool = False, - order_observable: ( - madspace._madspace_py.Observable.ObservableOption | None - ) = None, - order_indices: collections.abc.Sequence[typing.SupportsInt] = [], - ignore_incoming: bool = True, - name: str = "", - ) -> None: ... - + def __init__(self, pids: collections.abc.Sequence[typing.SupportsInt], observable: Observable.ObservableOption, select_pids: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]], sum_momenta: bool = False, sum_observable: bool = False, order_observable: madspace._madspace_py.Observable.ObservableOption | None = None, order_indices: collections.abc.Sequence[typing.SupportsInt] = [], ignore_incoming: bool = True, name: str = '') -> None: + ... class ObservableHistograms(FunctionGenerator): - def __init__(self, observables: collections.abc.Sequence[HistItem]) -> None: ... - + def __init__(self, observables: collections.abc.Sequence[HistItem]) -> None: + ... class PartonDensity(FunctionGenerator): - def __init__( - self, - grid: PdfGrid, - pids: collections.abc.Sequence[typing.SupportsInt], - dynamic_pid: bool = False, - prefix: str = "", - ) -> None: ... - + def __init__(self, grid: PdfGrid, pids: collections.abc.Sequence[typing.SupportsInt], dynamic_pid: bool = False, prefix: str = '') -> None: + ... class PdfGrid: - def __init__(self, file: str) -> None: ... - def coefficients_shape(self, batch_dim: bool = False) -> list[int]: ... - def initialize_globals(self, context: Context, prefix: str = "") -> None: ... - def logq2_shape(self, batch_dim: bool = False) -> list[int]: ... - def logx_shape(self, batch_dim: bool = False) -> list[int]: ... + def __init__(self, file: str) -> None: + ... + def coefficients_shape(self, batch_dim: bool = False) -> list[int]: + ... + def initialize_globals(self, context: Context, prefix: str = '') -> None: + ... + def logq2_shape(self, batch_dim: bool = False) -> list[int]: + ... + def logx_shape(self, batch_dim: bool = False) -> list[int]: + ... @property - def grid_point_count(self) -> int: ... + def grid_point_count(self) -> int: + ... @property - def logq2(self) -> list[float]: ... + def logq2(self) -> list[float]: + ... @property - def logx(self) -> list[float]: ... + def logx(self) -> list[float]: + ... @property - def pids(self) -> list[int]: ... + def pids(self) -> list[int]: + ... @property - def q(self) -> list[float]: ... + def q(self) -> list[float]: + ... @property - def q_count(self) -> int: ... + def q_count(self) -> int: + ... @property - def region_sizes(self) -> list[int]: ... + def region_sizes(self) -> list[int]: + ... @property - def values(self) -> list[list[float]]: ... + def values(self) -> list[list[float]]: + ... @property - def x(self) -> list[float]: ... - + def x(self) -> list[float]: + ... class PhaseSpaceMapping(Mapping): class TChannelMode: """ Members: - + propagator - + rambo - + chili """ - - __members__: typing.ClassVar[ - dict[str, PhaseSpaceMapping.TChannelMode] - ] # value = {'propagator': , 'rambo': , 'chili': } - chili: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - propagator: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - rambo: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + __members__: typing.ClassVar[dict[str, PhaseSpaceMapping.TChannelMode]] # value = {'propagator': , 'rambo': , 'chili': } + chili: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + propagator: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + rambo: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... @property - def name(self) -> str: ... + def name(self) -> str: + ... @property - def value(self) -> int: ... - - chili: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - propagator: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = - rambo: typing.ClassVar[ - PhaseSpaceMapping.TChannelMode - ] # value = + def value(self) -> int: + ... + chili: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + propagator: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = + rambo: typing.ClassVar[PhaseSpaceMapping.TChannelMode] # value = @typing.overload - def __init__( - self, - topology: Topology, - cm_energy: typing.SupportsFloat, - leptonic: bool = False, - invariant_power: typing.SupportsFloat = 0.8, - t_channel_mode: PhaseSpaceMapping.TChannelMode = PhaseSpaceMapping.TChannelMode.TChannelMode.propagator, - cuts: madspace._madspace_py.Cuts | None = None, - permutations: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ] = [], - ) -> None: ... + def __init__(self, topology: Topology, cm_energy: typing.SupportsFloat, leptonic: bool = False, invariant_power: typing.SupportsFloat = 0.8, t_channel_mode: PhaseSpaceMapping.TChannelMode = PhaseSpaceMapping.TChannelMode.TChannelMode.propagator, cuts: madspace._madspace_py.Cuts | None = None, permutations: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] = []) -> None: + ... @typing.overload - def __init__( - self, - masses: collections.abc.Sequence[typing.SupportsFloat], - cm_energy: typing.SupportsFloat, - leptonic: bool = False, - invariant_power: typing.SupportsFloat = 0.8, - mode: PhaseSpaceMapping.TChannelMode = PhaseSpaceMapping.TChannelMode.TChannelMode.rambo, - cuts: madspace._madspace_py.Cuts | None = None, - ) -> None: ... - def channel_count(self) -> int: ... - def particle_count(self) -> int: ... - def random_dim(self) -> int: ... - + def __init__(self, masses: collections.abc.Sequence[typing.SupportsFloat], cm_energy: typing.SupportsFloat, leptonic: bool = False, invariant_power: typing.SupportsFloat = 0.8, mode: PhaseSpaceMapping.TChannelMode = PhaseSpaceMapping.TChannelMode.TChannelMode.rambo, cuts: madspace._madspace_py.Cuts | None = None) -> None: + ... + def channel_count(self) -> int: + ... + def particle_count(self) -> int: + ... + def random_dim(self) -> int: + ... class PrettyBox: - def __init__( - self, - title: str, - rows: typing.SupportsInt, - columns: collections.abc.Sequence[typing.SupportsInt], - offset: typing.SupportsInt = 0, - box_width: typing.SupportsInt = 91, - ) -> None: ... - def print_first(self) -> None: ... - def print_update(self) -> None: ... - def set_cell( - self, row: typing.SupportsInt, column: typing.SupportsInt, value: str - ) -> None: ... - def set_column( - self, column: typing.SupportsInt, values: collections.abc.Sequence[str] - ) -> None: ... - def set_row( - self, row: typing.SupportsInt, values: collections.abc.Sequence[str] - ) -> None: ... - @property - def line_count(self) -> int: ... - + def __init__(self, title: str, rows: typing.SupportsInt, columns: collections.abc.Sequence[typing.SupportsInt], offset: typing.SupportsInt = 0, box_width: typing.SupportsInt = 91) -> None: + ... + def print_first(self) -> None: + ... + def print_update(self) -> None: + ... + def set_cell(self, row: typing.SupportsInt, column: typing.SupportsInt, value: str) -> None: + ... + def set_column(self, column: typing.SupportsInt, values: collections.abc.Sequence[str]) -> None: + ... + def set_row(self, row: typing.SupportsInt, values: collections.abc.Sequence[str]) -> None: + ... + @property + def line_count(self) -> int: + ... class Propagator: - def __init__( - self, - mass: typing.SupportsFloat = 0.0, - width: typing.SupportsFloat = 0.0, - integration_order: typing.SupportsInt = 0, - e_min: typing.SupportsFloat = 0.0, - e_max: typing.SupportsFloat = 0.0, - pdg_id: typing.SupportsInt = 0, - ) -> None: ... + def __init__(self, mass: typing.SupportsFloat = 0.0, width: typing.SupportsFloat = 0.0, integration_order: typing.SupportsInt = 0, e_min: typing.SupportsFloat = 0.0, e_max: typing.SupportsFloat = 0.0, pdg_id: typing.SupportsInt = 0) -> None: + ... @property - def e_max(self) -> float: ... + def e_max(self) -> float: + ... @property - def e_min(self) -> float: ... + def e_min(self) -> float: + ... @property - def integration_order(self) -> int: ... + def integration_order(self) -> int: + ... @property - def mass(self) -> float: ... + def mass(self) -> float: + ... @property - def pdg_id(self) -> int: ... + def pdg_id(self) -> int: + ... @property - def width(self) -> float: ... - + def width(self) -> float: + ... class PropagatorChannelWeights(FunctionGenerator): - def __init__( - self, - topologies: collections.abc.Sequence[Topology], - permutations: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - channel_indices: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ], - ) -> None: ... - + def __init__(self, topologies: collections.abc.Sequence[Topology], permutations: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]], channel_indices: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]) -> None: + ... class RunningCoupling(FunctionGenerator): - def __init__(self, grid: AlphaSGrid, prefix: str = "") -> None: ... - + def __init__(self, grid: AlphaSGrid, prefix: str = '') -> None: + ... class SubchannelWeights(FunctionGenerator): - def __init__( - self, - topologies: collections.abc.Sequence[collections.abc.Sequence[Topology]], - permutations: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - channel_indices: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ], - ) -> None: ... - def channel_count(self) -> int: ... - + def __init__(self, topologies: collections.abc.Sequence[collections.abc.Sequence[Topology]], permutations: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]], channel_indices: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]) -> None: + ... + def channel_count(self) -> int: + ... class SubprocArgs: - def __init__( - self, - process_id: typing.SupportsInt = 0, - topologies: collections.abc.Sequence[Topology] = [], - permutations: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ] = [], - diagram_indices: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsInt] - ] = [], - diagram_color_indices: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ] = [], - color_flows: collections.abc.Sequence[ - collections.abc.Sequence[ - collections.abc.Sequence[tuple[typing.SupportsInt, typing.SupportsInt]] - ] - ] = [], - pdg_color_types: collections.abc.Mapping[ - typing.SupportsInt, typing.SupportsInt - ] = {}, - helicities: collections.abc.Sequence[ - collections.abc.Sequence[typing.SupportsFloat] - ] = [], - pdg_ids: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ] = [], - matrix_flavor_indices: collections.abc.Sequence[typing.SupportsInt] = [], - ) -> None: ... - @property - def color_flows(self) -> list[list[list[tuple[int, int]]]]: ... + def __init__(self, process_id: typing.SupportsInt = 0, topologies: collections.abc.Sequence[Topology] = [], permutations: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]] = [], diagram_indices: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] = [], diagram_color_indices: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]] = [], color_flows: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[tuple[typing.SupportsInt, typing.SupportsInt]]]] = [], pdg_color_types: collections.abc.Mapping[typing.SupportsInt, typing.SupportsInt] = {}, helicities: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsFloat]] = [], pdg_ids: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]] = [], matrix_flavor_indices: collections.abc.Sequence[typing.SupportsInt] = []) -> None: + ... + @property + def color_flows(self) -> list[list[list[tuple[int, int]]]]: + ... @color_flows.setter - def color_flows( - self, - arg0: collections.abc.Sequence[ - collections.abc.Sequence[ - collections.abc.Sequence[tuple[typing.SupportsInt, typing.SupportsInt]] - ] - ], - ) -> None: ... - @property - def diagram_color_indices(self) -> list[list[list[int]]]: ... + def color_flows(self, arg0: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[tuple[typing.SupportsInt, typing.SupportsInt]]]]) -> None: + ... + @property + def diagram_color_indices(self) -> list[list[list[int]]]: + ... @diagram_color_indices.setter - def diagram_color_indices( - self, - arg0: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - ) -> None: ... - @property - def diagram_indices(self) -> list[list[int]]: ... + def diagram_color_indices(self, arg0: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]]) -> None: + ... + @property + def diagram_indices(self) -> list[list[int]]: + ... @diagram_indices.setter - def diagram_indices( - self, - arg0: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]], - ) -> None: ... + def diagram_indices(self, arg0: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]) -> None: + ... @property - def helicities(self) -> list[list[float]]: ... + def helicities(self) -> list[list[float]]: + ... @helicities.setter - def helicities( - self, - arg0: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsFloat]], - ) -> None: ... + def helicities(self, arg0: collections.abc.Sequence[collections.abc.Sequence[typing.SupportsFloat]]) -> None: + ... @property - def matrix_flavor_indices(self) -> list[int]: ... + def matrix_flavor_indices(self) -> list[int]: + ... @matrix_flavor_indices.setter - def matrix_flavor_indices( - self, arg0: collections.abc.Sequence[typing.SupportsInt] - ) -> None: ... + def matrix_flavor_indices(self, arg0: collections.abc.Sequence[typing.SupportsInt]) -> None: + ... @property - def pdg_color_types(self) -> dict[int, int]: ... + def pdg_color_types(self) -> dict[int, int]: + ... @pdg_color_types.setter - def pdg_color_types( - self, arg0: collections.abc.Mapping[typing.SupportsInt, typing.SupportsInt] - ) -> None: ... + def pdg_color_types(self, arg0: collections.abc.Mapping[typing.SupportsInt, typing.SupportsInt]) -> None: + ... @property - def pdg_ids(self) -> list[list[list[int]]]: ... + def pdg_ids(self) -> list[list[list[int]]]: + ... @pdg_ids.setter - def pdg_ids( - self, - arg0: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - ) -> None: ... - @property - def permutations(self) -> list[list[list[int]]]: ... + def pdg_ids(self, arg0: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]]) -> None: + ... + @property + def permutations(self) -> list[list[list[int]]]: + ... @permutations.setter - def permutations( - self, - arg0: collections.abc.Sequence[ - collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]] - ], - ) -> None: ... - @property - def process_id(self) -> int: ... + def permutations(self, arg0: collections.abc.Sequence[collections.abc.Sequence[collections.abc.Sequence[typing.SupportsInt]]]) -> None: + ... + @property + def process_id(self) -> int: + ... @process_id.setter - def process_id(self, arg0: typing.SupportsInt) -> None: ... + def process_id(self, arg0: typing.SupportsInt) -> None: + ... @property - def topologies(self) -> list[Topology]: ... + def topologies(self) -> list[Topology]: + ... @topologies.setter - def topologies(self, arg0: collections.abc.Sequence[Topology]) -> None: ... - + def topologies(self, arg0: collections.abc.Sequence[Topology]) -> None: + ... class TPropagatorMapping(Mapping): - def __init__( - self, - integration_order: collections.abc.Sequence[typing.SupportsInt], - invariant_power: typing.SupportsFloat = 0.0, - ) -> None: ... - def random_dim(self) -> int: ... - + def __init__(self, integration_order: collections.abc.Sequence[typing.SupportsInt], invariant_power: typing.SupportsFloat = 0.0) -> None: + ... + def random_dim(self) -> int: + ... class Tensor: @staticmethod - def numpy(tensor): ... + def numpy(tensor): + ... @staticmethod - def torch(tensor): ... - def __dlpack__( - self, - stream: typing.SupportsInt | None = None, - max_version: tuple[typing.SupportsInt, typing.SupportsInt] | None = None, - dl_device: typing.SupportsInt | None = None, - copy: bool | None = None, - ) -> typing.Any: ... - def __dlpack_device__(self) -> tuple[int, int]: ... - + def torch(tensor): + ... + def __dlpack__(self, stream: typing.SupportsInt | None = None, max_version: tuple[typing.SupportsInt, typing.SupportsInt] | None = None, dl_device: typing.SupportsInt | None = None, copy: bool | None = None) -> typing.Any: + ... + def __dlpack_device__(self) -> tuple[int, int]: + ... class ThreeBodyDecay(Mapping): - def __init__(self, com: bool) -> None: ... - + def __init__(self, com: bool) -> None: + ... class Topology: @staticmethod - def topologies(diagram: Diagram) -> list[Topology]: ... - def __init__(self, diagram: Diagram) -> None: ... - def propagator_momentum_terms( - self, arg0: bool - ) -> list[tuple[list[int], float, float]]: ... + def topologies(diagram: Diagram) -> list[Topology]: + ... + def __init__(self, diagram: Diagram) -> None: + ... + def propagator_momentum_terms(self, arg0: bool) -> list[tuple[list[int], float, float]]: + ... @property - def decay_integration_order(self) -> list[int]: ... + def decay_integration_order(self) -> list[int]: + ... @property - def decays(self) -> list[Decay]: ... + def decays(self) -> list[Decay]: + ... @property - def incoming_masses(self) -> list[float]: ... + def incoming_masses(self) -> list[float]: + ... @property - def outgoing_indices(self) -> list[int]: ... + def outgoing_indices(self) -> list[int]: + ... @property - def outgoing_masses(self) -> list[float]: ... + def outgoing_masses(self) -> list[float]: + ... @property - def t_integration_order(self) -> list[int]: ... + def t_integration_order(self) -> list[int]: + ... @property - def t_propagator_count(self) -> int: ... + def t_propagator_count(self) -> int: + ... @property - def t_propagator_masses(self) -> list[float]: ... + def t_propagator_masses(self) -> list[float]: + ... @property - def t_propagator_widths(self) -> list[float]: ... - + def t_propagator_widths(self) -> list[float]: + ... class TwoBodyDecay(Mapping): - def __init__(self, com: bool) -> None: ... - + def __init__(self, com: bool) -> None: + ... class TwoToThreeParticleScattering(Mapping): - def __init__( - self, - t_invariant_power: typing.SupportsFloat = 0.0, - t_mass: typing.SupportsFloat = 0.0, - t_width: typing.SupportsFloat = 0.0, - s_invariant_power: typing.SupportsFloat = 0.0, - s_mass: typing.SupportsFloat = 0.0, - s_width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, t_invariant_power: typing.SupportsFloat = 0.0, t_mass: typing.SupportsFloat = 0.0, t_width: typing.SupportsFloat = 0.0, s_invariant_power: typing.SupportsFloat = 0.0, s_mass: typing.SupportsFloat = 0.0, s_width: typing.SupportsFloat = 0.0) -> None: + ... class TwoToTwoParticleScattering(Mapping): - def __init__( - self, - com: bool, - invariant_power: typing.SupportsFloat = 0.0, - mass: typing.SupportsFloat = 0.0, - width: typing.SupportsFloat = 0.0, - ) -> None: ... - + def __init__(self, com: bool, invariant_power: typing.SupportsFloat = 0.0, mass: typing.SupportsFloat = 0.0, width: typing.SupportsFloat = 0.0) -> None: + ... class Type: @typing.overload - def __init__( - self, - dtype: DataType, - batch_size: BatchSize, - shape: collections.abc.Sequence[typing.SupportsInt], - ) -> None: ... + def __init__(self, dtype: DataType, batch_size: BatchSize, shape: collections.abc.Sequence[typing.SupportsInt]) -> None: + ... @typing.overload - def __init__( - self, batch_size_list: collections.abc.Sequence[BatchSize] - ) -> None: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... + def __init__(self, batch_size_list: collections.abc.Sequence[BatchSize]) -> None: + ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... @property - def batch_size(self) -> BatchSize: ... + def batch_size(self) -> BatchSize: + ... @property - def dtype(self) -> DataType: ... + def dtype(self) -> DataType: + ... @property - def shape(self) -> list[int]: ... - + def shape(self) -> list[int]: + ... class Unweighter(FunctionGenerator): - def __init__(self, types: NamedTypes) -> None: ... - + def __init__(self, types: NamedTypes) -> None: + ... class Value: @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, value: typing.SupportsFloat) -> None: ... - def __repr__(self) -> str: ... - def __str__(self) -> str: ... + def __init__(self, value: typing.SupportsFloat) -> None: + ... + def __repr__(self) -> str: + ... + def __str__(self) -> str: + ... @property - def literal_value( - self, - ) -> int | float | tuple[list[int], list[int] | list[float]] | None: ... + def literal_value(self) -> int | float | tuple[list[int], list[int] | list[float]] | None: + ... @property - def local_index(self) -> int: ... + def local_index(self) -> int: + ... @property - def type(self) -> Type: ... - + def type(self) -> Type: + ... class VegasGridOptimizer: - def __init__( - self, - contexts: collections.abc.Sequence[Context], - grid_name: str, - damping: typing.SupportsFloat, - ) -> None: ... - def add_data(self, values: typing.Any, counts: typing.Any) -> None: ... - def optimize(self) -> None: ... - + def __init__(self, contexts: collections.abc.Sequence[Context], grid_name: str, damping: typing.SupportsFloat) -> None: + ... + def add_data(self, values: typing.Any, counts: typing.Any) -> None: + ... + def optimize(self) -> None: + ... class VegasHistogram(FunctionGenerator): - def __init__( - self, dimension: typing.SupportsInt, bin_count: typing.SupportsInt - ) -> None: ... - + def __init__(self, dimension: typing.SupportsInt, bin_count: typing.SupportsInt) -> None: + ... class VegasMapping(Mapping): - def __init__( - self, - dimension: typing.SupportsInt, - bin_count: typing.SupportsInt, - prefix: str = "", - ) -> None: ... - def grid_name(self) -> str: ... - def initialize_globals(self, context: Context) -> None: ... - + def __init__(self, dimension: typing.SupportsInt, bin_count: typing.SupportsInt, prefix: str = '') -> None: + ... + def grid_name(self) -> str: + ... + def initialize_globals(self, context: Context) -> None: + ... class Verbosity: """ Members: - + silent - + log - + pretty """ - - __members__: typing.ClassVar[ - dict[str, Verbosity] - ] # value = {'silent': , 'log': , 'pretty': } + __members__: typing.ClassVar[dict[str, Verbosity]] # value = {'silent': , 'log': , 'pretty': } log: typing.ClassVar[Verbosity] # value = pretty: typing.ClassVar[Verbosity] # value = silent: typing.ClassVar[Verbosity] # value = - def __eq__(self, other: typing.Any) -> bool: ... - def __getstate__(self) -> int: ... - def __hash__(self) -> int: ... - def __index__(self) -> int: ... + def __eq__(self, other: typing.Any) -> bool: + ... + def __getstate__(self) -> int: + ... + def __hash__(self) -> int: + ... + def __index__(self) -> int: + ... @typing.overload - def __init__(self, value: typing.SupportsInt) -> None: ... + def __init__(self, value: typing.SupportsInt) -> None: + ... @typing.overload - def __init__(self, name: str) -> None: ... - def __int__(self) -> int: ... - def __ne__(self, other: typing.Any) -> bool: ... - def __repr__(self) -> str: ... - def __setstate__(self, state: typing.SupportsInt) -> None: ... - def __str__(self) -> str: ... - @property - def name(self) -> str: ... - @property - def value(self) -> int: ... - -def batch_float_array(count: typing.SupportsInt) -> Type: ... -def batch_four_vec_array(count: typing.SupportsInt) -> Type: ... -def cpu_device() -> Device: ... -def cuda_device(index: typing.SupportsInt = 0) -> Device: ... -def default_context() -> Context: ... -def default_cuda_context(index: typing.SupportsInt = 0) -> Context: ... -def default_hip_context(index: typing.SupportsInt = 0) -> Context: ... -def format_progress( - progress: typing.SupportsFloat, width: typing.SupportsInt -) -> str: ... -def format_si_prefix(value: typing.SupportsFloat) -> str: ... -def format_with_error( - value: typing.SupportsFloat, error: typing.SupportsFloat -) -> str: ... -def hip_device(index: typing.SupportsInt = 0) -> Device: ... -def initialize_vegas_grid(context: Context, grid_name: str) -> None: ... -def multichannel_batch_size(count: typing.SupportsInt) -> Type: ... -def set_lib_path(lib_path: str) -> None: ... -def set_simd_vector_size(vector_size: typing.SupportsInt) -> None: ... - + def __init__(self, name: str) -> None: + ... + def __int__(self) -> int: + ... + def __ne__(self, other: typing.Any) -> bool: + ... + def __repr__(self) -> str: + ... + def __setstate__(self, state: typing.SupportsInt) -> None: + ... + def __str__(self) -> str: + ... + @property + def name(self) -> str: + ... + @property + def value(self) -> int: + ... +def batch_float_array(count: typing.SupportsInt) -> Type: + ... +def batch_four_vec_array(count: typing.SupportsInt) -> Type: + ... +def cpu_device() -> Device: + ... +def cuda_device(index: typing.SupportsInt = 0) -> Device: + ... +def default_context() -> Context: + ... +def default_cuda_context(index: typing.SupportsInt = 0) -> Context: + ... +def default_hip_context(index: typing.SupportsInt = 0) -> Context: + ... +def format_progress(progress: typing.SupportsFloat, width: typing.SupportsInt) -> str: + ... +def format_si_prefix(value: typing.SupportsFloat) -> str: + ... +def format_with_error(value: typing.SupportsFloat, error: typing.SupportsFloat) -> str: + ... +def hip_device(index: typing.SupportsInt = 0) -> Device: + ... +def initialize_vegas_grid(context: Context, grid_name: str) -> None: + ... +def multichannel_batch_size(count: typing.SupportsInt) -> Type: + ... +def set_lib_path(lib_path: str) -> None: + ... +def set_simd_vector_size(vector_size: typing.SupportsInt) -> None: + ... batch_float: Type # value = float[batch_size] batch_four_vec: Type # value = float[batch_size, 4] batch_int: Type # value = int[batch_size] diff --git a/madspace/pyproject.toml b/madspace/pyproject.toml index e271cdd3f..4047a259a 100644 --- a/madspace/pyproject.toml +++ b/madspace/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "scikit_build_core.build" [project] name = "madspace" -version = "0.2.0" +version = "0.1.3" description="Fast and flexible phase-space mappings" readme = "README.md" authors = [ @@ -19,6 +19,14 @@ classifiers = [ ] dependencies = [] +[project.optional-dependencies] +test = [ + "pytest", + "numpy", + "scipy", + "torch", +] + [tool.scikit-build] wheel.expand-macos-universal-tags = true minimum-version = "build-system.requires" diff --git a/madspace/src/phasespace/color_ordered_mapping.cpp b/madspace/src/phasespace/color_ordered_mapping.cpp index 58d2c0579..c405230f4 100644 --- a/madspace/src/phasespace/color_ordered_mapping.cpp +++ b/madspace/src/phasespace/color_ordered_mapping.cpp @@ -9,48 +9,48 @@ using namespace madspace; namespace { -// Cyclically rotate the colour order so that particle 0 is first, then +// Cyclically rotate the color order so that particle 0 is first, then // split the rest into two sets: particles strictly between 0 and 1 in the // cyclic order go into set1, particles after 1 (wrapping around) go into // set2. std::pair, std::vector> -split_sets_from_colour_order(const std::vector& colour_order) { - std::size_t n = colour_order.size(); +split_sets_from_color_order(const std::vector& color_order) { + std::size_t n = color_order.size(); if (n < 4) { throw std::invalid_argument( "ColorOrderedMapping requires at least 4 particles (2 beams + 2 outgoing)" ); } - auto it = std::find(colour_order.begin(), colour_order.end(), 0u); - if (it == colour_order.end()) { - throw std::invalid_argument("colour_order must contain particle 0"); + auto it = std::find(color_order.begin(), color_order.end(), 0u); + if (it == color_order.end()) { + throw std::invalid_argument("color_order must contain particle 0"); } - std::size_t i0 = std::distance(colour_order.begin(), it); + std::size_t i0 = std::distance(color_order.begin(), it); std::vector rotated; rotated.reserve(n); for (std::size_t k = 0; k < n; ++k) { - rotated.push_back(colour_order[(i0 + k) % n]); + rotated.push_back(color_order[(i0 + k) % n]); } auto it1 = std::find(rotated.begin(), rotated.end(), 1u); if (it1 == rotated.end()) { - throw std::invalid_argument("colour_order must contain particle 1"); + throw std::invalid_argument("color_order must contain particle 1"); } 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]; - if (p <= 1) throw std::invalid_argument("invalid colour_order"); + 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]; - if (p <= 1) throw std::invalid_argument("invalid colour_order"); + if (p <= 1) throw std::invalid_argument("invalid color_order"); set2.push_back(p - 2); } if (set1.empty() || set2.empty()) { throw std::invalid_argument( "ColorOrderedMapping requires both sets to be non-empty " - "(particles 0 and 1 must not be adjacent in the colour order)" + "(particles 0 and 1 must not be adjacent in the color order)" ); } return {set1, set2}; @@ -69,14 +69,14 @@ std::size_t n_block_randoms_for_set_size(std::size_t k) { } // namespace ColorOrderedMapping::ColorOrderedMapping( - const std::vector& colour_order, + const std::vector& color_order, double t_invariant_power, double s_invariant_power ) : Mapping( "ColorOrderedMapping", [&] { - auto [s1, s2] = split_sets_from_colour_order(colour_order); + auto [s1, s2] = split_sets_from_color_order(color_order); bool use_double_t = (s1.size() == 1) != (s2.size() == 1); // DoubleT branch: 0 set-masses + 3 central randoms. // Standard branch: 1 set-mass per multi-particle side + 2 central randoms. @@ -99,7 +99,7 @@ ColorOrderedMapping::ColorOrderedMapping( return input_types; }(), [&] { - std::size_t n_out = colour_order.size() - 2; + std::size_t n_out = color_order.size() - 2; NamedVector output_types; for (std::size_t i = 0; i < n_out + 2; ++i) { output_types.push_back(std::format("momentum{}", i), batch_four_vec); @@ -107,7 +107,7 @@ ColorOrderedMapping::ColorOrderedMapping( return output_types; }(), [&] { - std::size_t n_out = colour_order.size() - 2; + std::size_t n_out = color_order.size() - 2; NamedVector cond_types{{"com_energy", batch_float}}; for (std::size_t i = 0; i < n_out; ++i) { cond_types.push_back(std::format("mass{}", i), batch_float); @@ -115,12 +115,12 @@ ColorOrderedMapping::ColorOrderedMapping( return cond_types; }() ), - _n_out(colour_order.size() - 2), + _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.) { - auto [s1, s2] = split_sets_from_colour_order(colour_order); + auto [s1, s2] = split_sets_from_color_order(color_order); _set1 = s1; _set2 = s2; _use_double_t = (s1.size() == 1) != (s2.size() == 1); diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index 49323941c..c9cf69259 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -491,7 +491,7 @@ PYBIND11_MODULE(_madspace_py, m) { py::classh(m, "ColorOrderedMapping") .def( py::init, double, double>(), - py::arg("colour_order"), + py::arg("color_order"), py::arg("t_invariant_power") = 0., py::arg("s_invariant_power") = 0. ) diff --git a/madspace/tests/test_color_ordered_mapping.py b/madspace/tests/test_color_ordered_mapping.py index 4b86af1eb..953e7f585 100644 --- a/madspace/tests/test_color_ordered_mapping.py +++ b/madspace/tests/test_color_ordered_mapping.py @@ -38,7 +38,7 @@ def _mass_sets_for(n_out): return sets -# Each entry: (colour_order, int_order_for_old_class, topology_label). +# 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, 2, 3, 4, 1, 5], [0, 1, 2], "n=6: set1={2,3,4}, set2={5}"), @@ -50,13 +50,13 @@ def _mass_sets_for(n_out): def _expand(topologies): - """Flat list of (colour_order, int_order, masses, label) pairs.""" + """Flat list of (color_order, int_order, masses, label) pairs.""" out = [] - for colour_order, int_order, topo_label in topologies: - n_out = len(colour_order) - 2 + 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((colour_order, int_order, masses, label)) + out.append((color_order, int_order, masses, label)) return out @@ -83,9 +83,9 @@ def _filter_physical(det, p_ext): # Tests # ---------------------------- -@pytest.mark.parametrize("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_momentum_conservation(rng, colour_order, _int_order, masses, _label): - mapping = ms.ColorOrderedMapping(colour_order) +@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] @@ -99,9 +99,9 @@ def test_momentum_conservation(rng, colour_order, _int_order, masses, _label): assert p_in == approx(p_out_sum, abs=1e-3, rel=1e-6) -@pytest.mark.parametrize("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_on_shell_masses(rng, colour_order, _int_order, masses, _label): - mapping = ms.ColorOrderedMapping(colour_order) +@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] @@ -119,8 +119,8 @@ def test_on_shell_masses(rng, colour_order, _int_order, masses, _label): f"particle {i} mass off" -@pytest.mark.parametrize("colour_order,_int_order,masses,_label", CASES, ids=CASE_IDS) -def test_inverse(rng, colour_order, _int_order, masses, _label): +@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 @@ -130,7 +130,7 @@ def test_inverse(rng, colour_order, _int_order, masses, _label): weight in any MC integral) and check the rest. """ n_events = N - mapping = ms.ColorOrderedMapping(colour_order) + 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] @@ -170,8 +170,8 @@ def test_inverse(rng, colour_order, _int_order, masses, _label): ) -@pytest.mark.parametrize("colour_order,int_order,masses,_label", CASES, ids=CASE_IDS) -def test_phase_space_volume_matches_old(rng, colour_order, int_order, masses, _label): +@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). @@ -180,7 +180,7 @@ def test_phase_space_volume_matches_old(rng, colour_order, int_order, masses, _l test_inverse is a precision check; this one is the physics check. """ n_events = N - new = ms.ColorOrderedMapping(colour_order) + new = ms.ColorOrderedMapping(color_order) old = ms.TPropagatorMapping(int_order) r_new = rng.random((n_events, new.random_dim()))