diff --git a/madspace/include/madspace/phasespace/color_ordered_mapping.hpp b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp index 86edf359b5..a59e4d11d3 100644 --- a/madspace/include/madspace/phasespace/color_ordered_mapping.hpp +++ b/madspace/include/madspace/phasespace/color_ordered_mapping.hpp @@ -1,5 +1,7 @@ #pragma once +#include +#include #include #include "madspace/phasespace/base.hpp" @@ -53,14 +55,34 @@ namespace madspace { class ColorOrderedMapping : public Mapping { public: /** - * @param color_order 0-based permutation of `{0, ..., n-1}` - * (`n = n_out + 2`) giving the chain order; - * particles 0 and 1 are the incoming beams. See - * @ref Topology. - * @param t_invariant_power Exponent of the `1/|t|^p` sampling of every - * momentum transfer; see @ref Invariant. - * @param s_invariant_power Exponent of the `1/s^p` sampling of every - * time-like invariant; see @ref Invariant. + * Settings of the chain. Besides these, the chain follows AmpliCol's + * gen23 [2]: a side of exactly two particles (with the other side not + * empty) is one 2->3 step with the opposite beam as the previous emission, + * the composite and rest masses are sampled as `1/sqrt(s)` (flat in the + * mass), and the s23 random numbers of the 2->3 steps are remapped with the + * arcsine map of @ref TwoToThreeParticleScattering. + */ + struct Options { + /// 0-based permutation of `{0, ..., n-1}` (`n = n_out + 2`) giving the + /// chain order; particles 0 and 1 are the incoming beams. See + /// @ref Topology. + std::vector color_order; + /// Exponent of the `1/|x|^p` sampling of the momentum transfers and + /// the s23 invariants of the chain; see @ref Invariant. Unset: 1 with + /// cuts (AmpliCol) and 0.8 without, where these ranges reach zero and + /// `1/|x|` would pile the points up at the regulator of + /// @ref Invariant. + std::optional chain_invariant_power = std::nullopt; + /// In every 2->3 step, keep the azimuthal solution that passes the pt + /// cuts, and pick with the discrete choice only when both pass, as + /// AmpliCol does. The weight assumes uniformly sampled discrete + /// choices: turn it off when they are sampled with a learned + /// distribution. + bool cut_aware_branch = true; + }; + + /** + * @param options Chain order and settings; see @ref Options. * @param pt_min Per-outgoing-particle minimum transverse * momentum; empty disables the cut. See @ref Cuts. * @param m_inv_min Symmetric matrix of minimum pair invariant @@ -68,10 +90,6 @@ class ColorOrderedMapping : public Mapping { * @param dr_min Symmetric matrix of minimum pair `delta_r` * separations; empty disables the cut. See * @ref Cuts. - * @param arcsine_s23 Remap the random number of every 2->3 - * invariant with the arcsine map in the azimuth, - * removing the `1/|sin(phi)|` edge peak of the - * weight; see @ref TwoToThreeParticleScattering. * @param y_max Per-outgoing-particle bound on the absolute lab * rapidity, negative (or empty) for none. Used * where a block scatters the two beams in the @@ -87,13 +105,10 @@ class ColorOrderedMapping : public Mapping { * conditions. */ ColorOrderedMapping( - const std::vector& color_order, - double t_invariant_power = 0.8, - double s_invariant_power = 0.8, + const Options& options, const std::vector& pt_min = {}, const std::vector>& m_inv_min = {}, const std::vector>& dr_min = {}, - bool arcsine_s23 = true, const std::vector& y_max = {} ); @@ -125,13 +140,10 @@ class ColorOrderedMapping : public Mapping { // The public constructor with the rapidity bounds worked out once. ColorOrderedMapping( const RapidityBounds& bounds, - const std::vector& color_order, - double t_invariant_power, - double s_invariant_power, + const Options& options, const std::vector& pt_min, const std::vector>& m_inv_min, - const std::vector>& dr_min, - bool arcsine_s23 + const std::vector>& dr_min ); Result build_forward_impl( @@ -151,6 +163,22 @@ class ColorOrderedMapping : public Mapping { // term, which is applied separately) for a subset of outgoing particles. // Returns 0 when cuts are disabled or the subset has fewer than 2 members. double cut_floor(const std::vector& subset) const; + // Floor on the squared mass of a set of outgoing particles (masses and + // cuts), and floor on its transverse mass: the sum of those of the + // particles with their pt cuts, at least the square root of the former. + std::array set_floors( + FunctionBuilder& fb, const std::vector& set, const ValueVec& m_out + ) const; + // Range of the squared mass of the rest system of a walk over the set s + // after its peel j, from a system of mass prev_mass. + std::array rest_mass_range( + FunctionBuilder& fb, + const std::vector& s, + std::size_t j, + Value prev_mass, + Value e_cm, + const ValueVec& m_out + ) const; // 0-indexed outgoing-particle indices (values in {0,...,n_out-1}). // _set1 contains the outgoing particles attached to beam 0's side, @@ -179,12 +207,14 @@ class ColorOrderedMapping : public Mapping { bool _has_cut; bool _rapidity_window; - Invariant _uniform_invariant; + Invariant _mass_invariant; TwoToTwoParticleScattering _com_scattering; TwoToTwoParticleScattering _lab_scattering; // first peel of a single chain, whose incoming momenta are the beams TwoToTwoParticleScattering _chain_scattering; TwoToThreeParticleScattering _two_to_three; + // the 2->3 step of a side of two particles, with an incoming recoil + TwoToThreeParticleScattering _two_to_three_beam; DoubleT _double_t; }; diff --git a/madspace/include/madspace/phasespace/phasespace.hpp b/madspace/include/madspace/phasespace/phasespace.hpp index 7c779ae46f..b1b5693fef 100644 --- a/madspace/include/madspace/phasespace/phasespace.hpp +++ b/madspace/include/madspace/phasespace/phasespace.hpp @@ -80,14 +80,21 @@ class PhaseSpaceMapping : public Mapping { * @param leptonic If true, skip the PDF convolution (no `x1`/`x2` * sampling). * @param invariant_power Exponent of the `1/s^p` sampling of every - * time-like invariant. See @ref Invariant. + * time-like invariant. See @ref Invariant. The + * chain of @ref TChannelMode::color_ordered takes + * its powers from @p color_ordered. * @param t_channel_mode The @ref TChannelMode strategy. * @param cuts Fiducial cuts to fold into the sampling. See * @ref Cuts. * @param permutations One outgoing-particle permutation per channel. * More than one enables the `permutation_index` * condition. - * @param color_order Color ordering for @ref TChannelMode::color_ordered. + * @param color_ordered Chain order and settings of + * @ref TChannelMode::color_ordered, required there + * and not allowed for the other modes. An unset + * `chain_invariant_power` is 1 with cuts and + * @p invariant_power without. See + * @ref ColorOrderedMapping::Options. * @param beam_rapidity Rapidity of the beams' centre-of-mass frame in the * lab frame the momenta are returned in: * 0.5 ln(E1 / E2) for beam energies E1 (along +z) @@ -100,6 +107,14 @@ class PhaseSpaceMapping : public Mapping { * beams' frame (leg 1 from beam 2, leg 2 from beam 1) * before the lab boost and the cuts. With identical * beams that is a rotation of the lab event too. + * @param s_hat_power Exponent of the `1/s_hat^p` sampling of the + * partonic invariant s_hat in the PDF convolution + * (see @ref Luminosity), in every mode. Unset, it is + * 1.2 for @ref TChannelMode::color_ordered when + * masses or cuts bound s_hat from below (as + * AmpliCol), and @p invariant_power otherwise. A + * resonant s_hat is sampled with its Breit-Wigner + * either way. */ PhaseSpaceMapping( const Topology& topology, @@ -109,9 +124,10 @@ class PhaseSpaceMapping : public Mapping { TChannelMode t_channel_mode = propagator, const std::optional& cuts = std::nullopt, const std::vector>& permutations = {}, - const std::optional>& color_order = std::nullopt, + const std::optional& color_ordered = std::nullopt, double beam_rapidity = 0., - bool mirror_beams = false + bool mirror_beams = false, + std::optional s_hat_power = std::nullopt ); /** @@ -121,11 +137,18 @@ class PhaseSpaceMapping : public Mapping { * @param cm_energy Total collision energy. * @param leptonic If true, skip the PDF convolution. * @param invariant_power Exponent of the `1/s^p` invariant sampling. See - * @ref Invariant. + * @ref Invariant. The chain of + * @ref TChannelMode::color_ordered takes its + * powers from @p color_ordered. * @param mode The @ref TChannelMode strategy (default * @ref TChannelMode::rambo). * @param cuts Fiducial cuts. See @ref Cuts. - * @param color_order Color ordering for @ref TChannelMode::color_ordered. + * @param color_ordered Chain order and settings of + * @ref TChannelMode::color_ordered, required there + * and not allowed for the other modes. An unset + * `chain_invariant_power` is 1 with cuts and + * @p invariant_power without. See + * @ref ColorOrderedMapping::Options. * @param beam_rapidity Rapidity of the beams' centre-of-mass frame in the * lab frame the momenta are returned in: * 0.5 ln(E1 / E2) for beam energies E1 (along +z) @@ -138,6 +161,14 @@ class PhaseSpaceMapping : public Mapping { * beams' frame (leg 1 from beam 2, leg 2 from beam 1) * before the lab boost and the cuts. With identical * beams that is a rotation of the lab event too. + * @param s_hat_power Exponent of the `1/s_hat^p` sampling of the + * partonic invariant s_hat in the PDF convolution + * (see @ref Luminosity), in every mode. Unset, it is + * 1.2 for @ref TChannelMode::color_ordered when + * masses or cuts bound s_hat from below (as + * AmpliCol), and @p invariant_power otherwise. A + * resonant s_hat is sampled with its Breit-Wigner + * either way. */ PhaseSpaceMapping( const std::vector& external_masses, @@ -146,9 +177,10 @@ class PhaseSpaceMapping : public Mapping { double invariant_power = 0.8, TChannelMode mode = rambo, const std::optional& cuts = std::nullopt, - const std::optional>& color_order = std::nullopt, + const std::optional& color_ordered = std::nullopt, double beam_rapidity = 0., - bool mirror_beams = false + bool mirror_beams = false, + std::optional s_hat_power = std::nullopt ); /** diff --git a/madspace/include/madspace/phasespace/three_particle.hpp b/madspace/include/madspace/phasespace/three_particle.hpp index 86953a254f..2ba277ccd3 100644 --- a/madspace/include/madspace/phasespace/three_particle.hpp +++ b/madspace/include/madspace/phasespace/three_particle.hpp @@ -174,6 +174,23 @@ class TwoToThreeParticleScattering : public Mapping { * holds it more precisely than pa + pb - p3 (a * soft system next to the beams) keeps that * precision. + * @param incoming_recoil If true, `momentum3` is an incoming momentum + * k (the other beam), and the second invariant is + * the momentum transfer @f$(p_2 - k)^2 \le 0@f$ + * instead of @f$(p_2 + p_3)^2@f$; its absolute + * value is sampled. With @p has_cut it is bounded + * by @f$|(p_2 - k)^2| \ge \texttt{etmin\_2}^2@f$, + * the `drcut` and `s23_min_cut` conditions are + * not used. Requires @p p12_condition. AmpliCol + * builds a side of two particles this way, with + * the opposite beam as the previous emission. + * @param cut_aware_branch With @p has_cut, form both azimuthal solutions + * and keep the one passing the transverse-energy + * floors `etmin_1`, `etmin_2` (AmpliCol's + * gen23_one_step): `discrete_choice` only picks + * when both pass. The weight assumes + * `discrete_choice` is sampled uniformly; with a + * learned distribution of it, leave this off. */ TwoToThreeParticleScattering( double t_invariant_power = 0, @@ -184,7 +201,9 @@ class TwoToThreeParticleScattering : public Mapping { double s_width = 0, bool has_cut = false, bool arcsine_s23 = true, - bool p12_condition = false + bool p12_condition = false, + bool incoming_recoil = false, + bool cut_aware_branch = false ); /// Number of discrete inputs (1): which of the two two-body solutions. @@ -204,12 +223,23 @@ class TwoToThreeParticleScattering : public Mapping { std::array split_conditions(FunctionBuilder& fb, const NamedVector& conditions) const; + // Sampled and kinematic range of the second invariant: s23 itself, or + // |(p2 - k)^2| for an incoming recoil. Arguments: the kinematic range of + // s23 and the conditions. + std::array sampling_ranges( + FunctionBuilder& fb, + Value s23_min, + Value s23_max, + const NamedVector& conditions + ) const; Invariant _t_invariant; Invariant _s_invariant; double _s_power, _s_mass; bool _has_cut; bool _arcsine_s23; bool _p12_condition; + bool _incoming_recoil; + bool _cut_aware_branch; }; } // namespace madspace diff --git a/madspace/instruction_set.yaml b/madspace/instruction_set.yaml index 30125d9b21..29155d2e8e 100644 --- a/madspace/instruction_set.yaml +++ b/madspace/instruction_set.yaml @@ -1212,6 +1212,53 @@ two_to_three_particle_scattering_inverse: desc: "Jacobian of the mapping" desc: "Inverse of `two_to_three_particle_scattering`. Building block of @ref TwoToThreeParticleScattering." +two_to_three_branch_choice: + inputs: + - name: phi_choice_in + type: [int] + desc: "sampled branch selector, used when both solutions pass" + - name: pa + type: [float, 4] + desc: "first incoming momentum" + - name: p12 + type: [float, 4] + desc: "momentum of the system of the first and second outgoing particles" + - name: p3 + type: [float, 4] + desc: "momentum of the already-generated third particle" + - name: u + type: [float] + desc: "position of s23 in its kinematic range" + - name: u_c + type: [float] + desc: "1 - u" + - name: t1_abs + type: [float] + desc: "absolute value of the first momentum transfer" + - name: m1 + type: [float] + desc: "mass of the first outgoing particle" + - name: m2 + type: [float] + desc: "mass of the second outgoing particle" + - name: etmin_1 + type: [float] + desc: "transverse-energy floor of the first outgoing momentum" + - name: etmin_2 + type: [float] + desc: "transverse-energy floor of the second outgoing momentum" + outputs: + - name: phi_choice + type: [int] + desc: "branch selector to use: a passing solution" + - name: factor + type: [float] + desc: "weight factor for a uniformly sampled selector: 1 if both solutions pass, 1/2 if one does, 0 if none" + - name: inv_factor + type: [float] + desc: "inverse of the factor, 0 if no solution passes" + desc: "Choice of the azimuthal solution of the 2 to 3 step that passes the transverse-energy cuts, as in AmpliCol. Building block of @ref TwoToThreeParticleScattering." + double_t_scattering: inputs: - name: r_phi diff --git a/madspace/src/compgraphs/function.cpp b/madspace/src/compgraphs/function.cpp index cc6ce199d6..a407ac827f 100644 --- a/madspace/src/compgraphs/function.cpp +++ b/madspace/src/compgraphs/function.cpp @@ -348,11 +348,16 @@ FunctionBuilder::instruction(InstructionPtr instruction, const ValueVec& args) { auto& const_arg = args.at(1 - var_index); switch (opcode) { case opcodes::add: - case opcodes::sub: if (std::get(const_arg.literal_value) == 0.) { return {var_arg}; } break; + case opcodes::sub: + // x - 0 = x, but 0 - x is not x + if (var_index == 0 && std::get(const_arg.literal_value) == 0.) { + return {var_arg}; + } + break; case opcodes::mul: if (std::get(const_arg.literal_value) == 1.) { return {var_arg}; diff --git a/madspace/src/kernels/threeparticle.hpp b/madspace/src/kernels/threeparticle.hpp index 72a3e202d1..2233d0d8eb 100644 --- a/madspace/src/kernels/threeparticle.hpp +++ b/madspace/src/kernels/threeparticle.hpp @@ -667,12 +667,7 @@ KERNELSPEC void kernel_s23_position( // edges, which is where the arcsine map puts its points. template KERNELSPEC Triplet, FVal, FVal> s23_power_offsets( - FVal x, - FVal x_c, - FVal s_min, - FVal s_max, - FVal power, - FVal mass + FVal x, FVal x_c, FVal s_min, FVal s_max, FVal power, FVal mass ) { // returns s23 - s_min, s_max - s23 and d s23 / d x auto m2 = mass * mass - 1e-2; @@ -785,8 +780,7 @@ KERNELSPEC void kernel_s23_arcsine_sample_inverse( auto big_r = expm1(p_safe * log_ratio); auto big_r_safe = where(fabs(big_r) > 0., big_r, FVal(1.)); auto pow_x = expm1(p_safe * log1p(d_lo / qa)) / big_r_safe; - auto pow_x_c = - -expm1(p_safe * log1p(-d_hi / qb)) * (1. + big_r) / big_r_safe; + auto pow_x_c = -expm1(p_safe * log1p(-d_hi / qb)) * (1. + big_r) / big_r_safe; auto is_flat = pw == 0.; auto is_log = pw == 1.; auto x = min(max(where(is_flat, flat_x, where(is_log, log_x, pow_x)), 0.), 1.); @@ -809,9 +803,8 @@ KERNELSPEC void kernel_s23_arcsine_sample_inverse( auto v_c = max(cos_b * cos_b + x_c * den_safe, 0.); auto theta = atan2(sqrt(v), sqrt(v_c)); auto r_arc = (theta - theta_a) / dtheta_safe; - auto dr_dx = where( - v * v_c > 0., den_safe / (dtheta_safe * 2. * sqrt(v * v_c)), FVal(0.) - ); + auto dr_dx = + where(v * v_c > 0., den_safe / (dtheta_safe * 2. * sqrt(v * v_c)), FVal(0.)); r = where(ok, min(max(r_arc, 0.), 1.), x); dr_dx = where(ok, dr_dx, FVal(1.)); det = where(offsets.third > 0., dr_dx / offsets.third, FVal(0.)); @@ -850,15 +843,14 @@ KERNELSPEC void kernel_two_to_three_particle_scattering( // cos(phi) = 2 u - 1, u = (s23 - s23_min) / (s23_max - s23_min), // -G4 = lambda(s12, ma^2, t2) / 16 * (s23 - s23_min) * (s23_max - s23). // Both are evaluated through u and 1 - u, which come in as inputs (see - // kernel_s23_position and kernel_s23_arcsine_sample). Forming G4 (or the equivalent cos(phi) - // expression) from the invariants instead cancels terms of order s^4 down + // kernel_s23_position and kernel_s23_arcsine_sample). Forming G4 (or the equivalent + // cos(phi) expression) from the invariants instead cancels terms of order s^4 down // to a result that can be many orders of magnitude smaller whenever the // s23 range is narrow (a soft particle), and returned Jacobians of up to // 1e20 times the typical weight. The range itself comes from momenta for // the same reason, see s23_min_max. - auto s23_range = s23_min_max( - load_mom(pa), load_mom(p3), p_12, s12, t1_abs, m1, m2 - ); + auto s23_range = + s23_min_max(load_mom(pa), load_mom(p3), p_12, s12, t1_abs, m1, m2); auto s23_width = s23_range.second - s23_range.first; FVal u_val(u), u_c_val(u_c); // cos(phi) = u - (1 - u), |sin(phi)| = 2 sqrt(u (1 - u)), both at full @@ -876,9 +868,8 @@ KERNELSPEC void kernel_two_to_three_particle_scattering( // The edges of the range (sin(phi) = 0) are a set of measure zero. auto det_2to3 = where(sqrt_neg_gram4_x8 > 0., 1. / sqrt_neg_gram4_x8, FVal(0.)); - auto scatter_out = p1com_from_tabs_cs( - pa_com, s12, cos_phi, sin_phi, t1_abs, m1, m2, ma_2, t2 - ); + auto scatter_out = + p1com_from_tabs_cs(pa_com, s12, cos_phi, sin_phi, t1_abs, m1, m2, ma_2, t2); auto p1_com = scatter_out.first; auto p3_p12 = boost(load_mom(p3), p_12, -1.); auto p1_rot = rotate_two_ref(p1_com, pa_com, p3_p12); @@ -966,5 +957,122 @@ KERNELSPEC void kernel_two_to_three_particle_scattering_inverse( det = 2 * det_2to3; } +// Whether the solution of the 2->3 step with the given sign of sin(phi) passes +// the transverse-energy floors, E^2 - p_z^2 >= etmin^2, of p1 and p2. The +// momenta are formed as in kernel_two_to_three_particle_scattering. +template +KERNELSPEC BVal two_to_three_branch_passes( + FourMom pa_com, + FourMom p3_p12, + FourMom p_12, + FVal s12, + FVal cos_phi, + FVal sin_phi, + FVal t1_abs, + FVal m1, + FVal m2, + FVal ma_2, + FVal t2, + FVal etmin_1, + FVal etmin_2 +) { + auto p1_com = + p1com_from_tabs_cs(pa_com, s12, cos_phi, sin_phi, t1_abs, m1, m2, ma_2, t2) + .first; + auto p1 = boost(rotate_two_ref(p1_com, pa_com, p3_p12), p_12, 1.); + FourMom p2; + for (int i = 0; i < 4; ++i) { + p2[i] = p_12[i] - p1[i]; + } + auto mt2_1 = (p1[0] - p1[3]) * (p1[0] + p1[3]); + auto mt2_2 = (p2[0] - p2[3]) * (p2[0] + p2[3]); + return ((etmin_1 <= 0.) | (mt2_1 >= etmin_1 * etmin_1)) & + ((etmin_2 <= 0.) | (mt2_2 >= etmin_2 * etmin_2)); +} + +// Choice between the two azimuthal solutions of the 2->3 step as in AmpliCol's +// gen23_one_step: both are formed, and only those passing the +// transverse-energy floors of p1 and p2 are kept. If both pass, index_in picks +// one (factor 1); if one passes it is taken whatever index_in is, and the two +// values of index_in lead to the same point (factor 1/2); if none passes the +// weight is zero. The factor is that of a uniformly sampled index_in. +template +KERNELSPEC void kernel_two_to_three_branch_choice( + IIn index_in, + FIn pa, + FIn p12, + FIn p3, + FIn u, + FIn u_c, + FIn t1_abs, + FIn m1, + FIn m2, + FIn etmin_1, + FIn etmin_2, + IOut index_out, + FOut factor, + FOut inv_factor +) { + auto p_12 = load_mom(p12); + auto p_a = load_mom(pa); + FourMom p_c; + for (int i = 0; i < 4; ++i) { + p_c[i] = p_12[i] - p_a[i]; + } + auto pa_com = boost(p_a, p_12, -1.); + auto p3_p12 = boost(load_mom(p3), p_12, -1.); + auto ma_2 = lsquare(p_a); + auto s12 = lsquare(p_12); + auto t2 = lsquare(p_c); + FVal u_val(u), u_c_val(u_c), e1(etmin_1), e2(etmin_2); + auto u_sum = u_val + u_c_val; + auto u_sum_safe = where(u_sum > 0., u_sum, FVal(1.)); + auto cos_phi = (u_val - u_c_val) / u_sum_safe; + auto sin_phi_abs = 2. * sqrt(max(u_val * u_c_val, 0.)) / u_sum_safe; + // index 0: sin(phi) >= 0, index 1: sin(phi) < 0 (as the scattering kernel) + auto pass_0 = two_to_three_branch_passes( + pa_com, + p3_p12, + p_12, + s12, + cos_phi, + sin_phi_abs, + t1_abs, + m1, + m2, + ma_2, + t2, + e1, + e2 + ); + auto pass_1 = two_to_three_branch_passes( + pa_com, + p3_p12, + p_12, + s12, + cos_phi, + -sin_phi_abs, + t1_abs, + m1, + m2, + ma_2, + t2, + e1, + e2 + ); + IVal index(index_in); + index_out = where( + pass_0 & pass_1, + index, + where(pass_0, IVal(0), where(pass_1, IVal(1), index)) + ); + factor = where( + pass_0 & pass_1, FVal(1.), where(pass_0 | pass_1, FVal(0.5), FVal(0.)) + ); + inv_factor = where( + pass_0 & pass_1, FVal(1.), where(pass_0 | pass_1, FVal(2.), FVal(0.)) + ); +} + } // namespace kernels } // namespace madspace diff --git a/madspace/src/phasespace/color_ordered_mapping.cpp b/madspace/src/phasespace/color_ordered_mapping.cpp index cd7402b66b..2022a253ae 100644 --- a/madspace/src/phasespace/color_ordered_mapping.cpp +++ b/madspace/src/phasespace/color_ordered_mapping.cpp @@ -1,7 +1,9 @@ #include "madspace/phasespace/color_ordered_mapping.hpp" #include +#include #include +#include #include #include "madspace/util.hpp" @@ -76,11 +78,79 @@ std::size_t n_block_randoms_for_set_size(std::size_t k) { return 2 + 2 * (k - 2); } -std::size_t n_discrete_for_set_size(std::size_t k) { - // One discrete two-solution choice per 2->3 peel (all peels after the first). +// A side of exactly two particles is one 2->3 step off the opposite beam when +// the other side is not empty (AmpliCol), else a 2->2 peel. +bool two_particle_step(std::size_t k, std::size_t k_other) { + return k == 2 && k_other > 0; +} + +std::size_t n_discrete_for_set_size(std::size_t k, std::size_t k_other) { + // One discrete two-solution choice per 2->3 step: every peel after the first, + // and the single step of a side of two particles. + if (two_particle_step(k, k_other)) { + return 1; + } return (k >= 2) ? (k - 2) : 0; } +// s in [s_min, s_max], sampled with the given Invariant, and its weight. The +// cut-aware bounds of a mass can leave an empty range, which gets zero weight +// (and a small width, so that the sampling stays finite). +std::array sample_mass( + FunctionBuilder& fb, const Invariant& invariant, Value r, Value s_min, Value s_max +) { + auto ok = fb.cut_one(fb.sub(s_max, s_min), Value(0.), Value(1e300)); + auto width = fb.max(fb.mul(s_min, Value(1e-10)), Value(1e-300)); + auto res = + invariant.build_forward(fb, {r}, {s_min, fb.max(s_max, fb.add(s_min, width))}); + return {res["invariant"], fb.mul(res["det"], ok)}; +} + +std::array recover_mass( + FunctionBuilder& fb, const Invariant& invariant, Value s, Value s_min, Value s_max +) { + auto ok = fb.cut_one(fb.sub(s_max, s_min), Value(0.), Value(1e300)); + auto width = fb.max(fb.mul(s_min, Value(1e-10)), Value(1e-300)); + auto res = + invariant.build_inverse(fb, {s}, {s_min, fb.max(s_max, fb.add(s_min, width))}); + return {res["random"], fb.mul(res["det"], ok)}; +} + +// Range of the squared mass s_a of a system a that makes up the full final +// state together with a system b, from the energies in the partonic frame, +// E_a = (s_hat + s_a - s_b) / (2 sqrt(s_hat)) >= e_min_a and the same for b +// (AmpliCol's gent_one_step). e_min is a floor on the transverse mass: the sum +// of those of the particles, and the square root of the mass floor. With s_b +// still to be sampled (m_b empty) its floor and e_min_b bound it. +std::array full_system_mass_range( + FunctionBuilder& fb, + Value e_cm, + Value floor_a, + Value e_min_a, + Value floor_b, + Value e_min_b, + std::optional m_b +) { + auto lower = fb.add( + m_b ? fb.square(*m_b) : floor_b, + fb.mul(e_cm, fb.sub(fb.add(e_min_a, e_min_a), e_cm)) + ); + auto s_min = fb.max(floor_a, lower); + Value s_max; + if (m_b) { + // s_b + s_hat - 2 sqrt(s_hat) e_b, as (sqrt(s_hat) - e_b)^2 + s_b - e_b^2: + // without a cut e_b = m_b, and the range of a heavy b keeps its + // relative precision + auto e_b = fb.max(*m_b, e_min_b); + s_max = fb.add( + fb.square(fb.sub(e_cm, e_b)), fb.mul(fb.sub(*m_b, e_b), fb.add(*m_b, e_b)) + ); + } else { + s_max = fb.square(fb.sub(e_cm, e_min_b)); + } + return {s_min, s_max}; +} + double mat_at(const std::vector>& m, std::size_t i, std::size_t j) { if (i < m.size() && j < m.at(i).size()) { return m.at(i).at(j); @@ -130,6 +200,20 @@ static bool has_any_cut( return false; } +// Power of the momentum transfers and s23 invariants of the chain: as given, +// else 1 with cuts (AmpliCol) and 0.8 without, where their ranges reach zero +// and 1/|x| would pile the points up at the regulator of Invariant. +static double chain_power( + const ColorOrderedMapping::Options& options, + const std::vector& pt_min, + const std::vector>& m_inv_min, + const std::vector>& dr_min +) { + return options.chain_invariant_power.value_or( + has_any_cut(pt_min, m_inv_min, dr_min) ? 1. : 0.8 + ); +} + double ColorOrderedMapping::pt2(std::size_t i) const { double p = (i < _pt_min.size()) ? _pt_min.at(i) : 0.0; return p * p; @@ -159,6 +243,53 @@ double ColorOrderedMapping::cut_floor(const std::vector& subset) co return cut * scaling; } +std::array ColorOrderedMapping::set_floors( + FunctionBuilder& fb, const std::vector& set, const ValueVec& m_out +) const { + Value mass_sum = m_out.at(set.at(0)); + for (std::size_t i = 1; i < set.size(); ++i) { + mass_sum = fb.add(mass_sum, m_out.at(set.at(i))); + } + auto floor = fb.square(mass_sum); + double fl = cut_floor(set); + if (fl > 0.0) { + floor = fb.max(floor, Value(fl)); + } + if (!_has_cut) { + return {floor, fb.sqrt(floor)}; + } + auto e_min = fb.sqrt(fb.add(fb.square(m_out.at(set.at(0))), Value(pt2(set.at(0))))); + for (std::size_t i = 1; i < set.size(); ++i) { + e_min = fb.add( + e_min, + fb.sqrt(fb.add(fb.square(m_out.at(set.at(i))), Value(pt2(set.at(i))))) + ); + } + return {floor, fb.max(e_min, fb.sqrt(floor))}; +} + +std::array ColorOrderedMapping::rest_mass_range( + FunctionBuilder& fb, + const std::vector& s, + std::size_t j, + Value prev_mass, + Value e_cm, + const ValueVec& m_out +) const { + // rest system {s[j+1], ..., s[k-1]} after peeling s[j] off a system of + // mass prev_mass + std::vector rest(s.begin() + j + 1, s.end()); + auto [floor, e_min] = set_floors(fb, rest, m_out); + if (_use_single_chain && j == 0) { + // the rest and the first particle make up the full final state + auto [floor_p, e_min_p] = set_floors(fb, {s.at(0)}, m_out); + return full_system_mass_range( + fb, e_cm, floor, e_min, floor_p, e_min_p, m_out.at(s.at(0)) + ); + } + return {floor, fb.square(fb.sub(prev_mass, m_out.at(s.at(j))))}; +} + // A set of particles is bounded by the largest of their bounds when each has // one: the rapidity of a sum of momenta is a weighted mean of theirs. ColorOrderedMapping::RapidityBounds ColorOrderedMapping::rapidity_bounds( @@ -200,40 +331,27 @@ ColorOrderedMapping::RapidityBounds ColorOrderedMapping::rapidity_bounds( } ColorOrderedMapping::ColorOrderedMapping( - const std::vector& color_order, - double t_invariant_power, - double s_invariant_power, + const Options& options, const std::vector& pt_min, const std::vector>& m_inv_min, const std::vector>& dr_min, - bool arcsine_s23, const std::vector& y_max ) : ColorOrderedMapping( - rapidity_bounds(color_order, y_max), - color_order, - t_invariant_power, - s_invariant_power, - pt_min, - m_inv_min, - dr_min, - arcsine_s23 + rapidity_bounds(options.color_order, y_max), options, pt_min, m_inv_min, dr_min ) {} ColorOrderedMapping::ColorOrderedMapping( const RapidityBounds& bounds, - const std::vector& color_order, - double t_invariant_power, - double s_invariant_power, + const Options& options, const std::vector& pt_min, const std::vector>& m_inv_min, - const std::vector>& dr_min, - bool arcsine_s23 + const std::vector>& dr_min ) : Mapping( "ColorOrderedMapping", [&] { - auto [s1, s2] = split_sets_from_color_order(color_order); + auto [s1, s2] = split_sets_from_color_order(options.color_order); bool use_single_chain = s1.empty() || s2.empty(); bool use_double_t = !use_single_chain && ((s1.size() == 1) != (s2.size() == 1)); @@ -251,8 +369,8 @@ ColorOrderedMapping::ColorOrderedMapping( n_block_randoms_for_set_size(s2.size()); std::size_t total = n_set_masses + n_intermediate_masses + n_central + n_walk; - std::size_t n_discrete = - n_discrete_for_set_size(s1.size()) + n_discrete_for_set_size(s2.size()); + std::size_t n_discrete = n_discrete_for_set_size(s1.size(), s2.size()) + + n_discrete_for_set_size(s2.size(), s1.size()); NamedVector input_types; for (std::size_t i = 0; i < total; ++i) { input_types.push_back(std::format("random{}", i), batch_float); @@ -265,7 +383,7 @@ ColorOrderedMapping::ColorOrderedMapping( return input_types; }(), [&] { - std::size_t n_out = color_order.size() - 2; + std::size_t n_out = options.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); @@ -273,7 +391,7 @@ ColorOrderedMapping::ColorOrderedMapping( return output_types; }(), [&] { - std::size_t n_out = color_order.size() - 2; + std::size_t n_out = options.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); @@ -285,15 +403,16 @@ ColorOrderedMapping::ColorOrderedMapping( return cond_types; }() ), - _n_out(color_order.size() - 2), + _n_out(options.color_order.size() - 2), _pt_min(pt_min), _m_inv_min(m_inv_min), _dr_min(dr_min), _has_cut(has_any_cut(pt_min, m_inv_min, dr_min)), _rapidity_window(bounds.active()), + _mass_invariant(0.5), _com_scattering( true, - t_invariant_power, + chain_power(options, pt_min, m_inv_min, dr_min), 0., 0., has_any_cut(pt_min, m_inv_min, dr_min), @@ -302,11 +421,15 @@ ColorOrderedMapping::ColorOrderedMapping( 1. ), _lab_scattering( - false, t_invariant_power, 0., 0., has_any_cut(pt_min, m_inv_min, dr_min) + false, + chain_power(options, pt_min, m_inv_min, dr_min), + 0., + 0., + has_any_cut(pt_min, m_inv_min, dr_min) ), _chain_scattering( false, - t_invariant_power, + chain_power(options, pt_min, m_inv_min, dr_min), 0., 0., has_any_cut(pt_min, m_inv_min, dr_min), @@ -315,21 +438,36 @@ ColorOrderedMapping::ColorOrderedMapping( bounds.chain_beam_sign ), _two_to_three( - t_invariant_power, + chain_power(options, pt_min, m_inv_min, dr_min), 0., 0., - s_invariant_power, + chain_power(options, pt_min, m_inv_min, dr_min), 0., 0., has_any_cut(pt_min, m_inv_min, dr_min), - arcsine_s23, - true + true, + true, + false, + options.cut_aware_branch + ), + _two_to_three_beam( + chain_power(options, pt_min, m_inv_min, dr_min), + 0., + 0., + chain_power(options, pt_min, m_inv_min, dr_min), + 0., + 0., + has_any_cut(pt_min, m_inv_min, dr_min), + true, + true, + true, + options.cut_aware_branch ), _double_t( - t_invariant_power, + chain_power(options, pt_min, m_inv_min, dr_min), 0., 0., - t_invariant_power, + chain_power(options, pt_min, m_inv_min, dr_min), 0., 0., has_any_cut(pt_min, m_inv_min, dr_min), @@ -337,7 +475,7 @@ ColorOrderedMapping::ColorOrderedMapping( bounds.double_t ? bounds.central2 : -1., 1. ) { - auto [s1, s2] = split_sets_from_color_order(color_order); + auto [s1, s2] = split_sets_from_color_order(options.color_order); _set1 = s1; _set2 = s2; _use_single_chain = s1.empty() || s2.empty(); @@ -351,8 +489,8 @@ ColorOrderedMapping::ColorOrderedMapping( std::size_t n_walk = n_block_randoms_for_set_size(s1.size()) + n_block_randoms_for_set_size(s2.size()); _random_dim = n_set_masses + n_intermediate_masses + n_central + n_walk; - _discrete_dim = - n_discrete_for_set_size(s1.size()) + n_discrete_for_set_size(s2.size()); + _discrete_dim = n_discrete_for_set_size(s1.size(), s2.size()) + + n_discrete_for_set_size(s2.size(), s1.size()); } Mapping::Result ColorOrderedMapping::build_forward_impl( @@ -470,35 +608,31 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( m_set1 = fb.sqrt(m2); } } else { - Value mass_sum_set1 = masses_of(_set1); - Value mass_sum_set2 = masses_of(_set2); + // Composite masses: set1 first, then set2 in what is left. Both make up + // the full final state, so the energies of the two sides bound them + // (AmpliCol's shatminmax and gent_one_step), with the mass and + // transverse-mass floors of the cuts. + auto [floor1, e_min1] = set_floors(fb, _set1, m_out); + auto [floor2, e_min2] = set_floors(fb, _set2, m_out); if (_set1.size() >= 2) { - auto s_min = fb.square(mass_sum_set1); - // Cut floor: composite invariant mass must clear the cut-implied - // minimum (gen23 invm_min) in addition to the kinematic minimum. - double floor1 = cut_floor(_set1); - if (floor1 > 0.0) { - s_min = fb.max(s_min, Value(floor1)); - } - auto s_max = fb.square(fb.sub(e_cm, mass_sum_set2)); - auto res = - _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); - m_set1 = fb.sqrt(res["invariant"]); - dets.push_back(res["det"]); + auto [s_min, s_max] = full_system_mass_range( + fb, e_cm, floor1, e_min1, floor2, e_min2, std::nullopt + ); + auto [s, det] = + sample_mass(fb, _mass_invariant, next_random(), s_min, s_max); + m_set1 = fb.sqrt(s); + dets.push_back(det); } else { m_set1 = m_out.at(_set1.at(0)); } if (_set2.size() >= 2) { - auto s_min = fb.square(mass_sum_set2); - double floor2 = cut_floor(_set2); - if (floor2 > 0.0) { - s_min = fb.max(s_min, Value(floor2)); - } - auto s_max = fb.square(fb.sub(e_cm, m_set1)); - auto res = - _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); - m_set2 = fb.sqrt(res["invariant"]); - dets.push_back(res["det"]); + auto [s_min, s_max] = full_system_mass_range( + fb, e_cm, floor2, e_min2, floor1, e_min1, m_set1 + ); + auto [s, det] = + sample_mass(fb, _mass_invariant, next_random(), s_min, s_max); + m_set2 = fb.sqrt(s); + dets.push_back(det); } else { m_set2 = m_out.at(_set2.at(0)); } @@ -555,23 +689,12 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( } Value prev_mass = m_set; for (std::size_t j = 0; j < k - 2; ++j) { - Value m_min = m_out.at(s.at(j + 1)); - for (std::size_t i = j + 2; i < k; ++i) { - m_min = fb.add(m_min, m_out.at(s.at(i))); - } - auto s_min = fb.square(m_min); - // Cut floor for the rest system {s[j+1], ..., s[k-1]}. - std::vector rest_sub(s.begin() + j + 1, s.end()); - double fl = cut_floor(rest_sub); - if (fl > 0.0) { - s_min = fb.max(s_min, Value(fl)); - } - auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s.at(j)))); - auto r = - _uniform_invariant.build_forward(fb, {next_random()}, {s_min, s_max}); - Value m_rest = fb.sqrt(r["invariant"]); + auto [s_min, s_max] = rest_mass_range(fb, s, j, prev_mass, e_cm, m_out); + auto [s_rest, det] = + sample_mass(fb, _mass_invariant, next_random(), s_min, s_max); + Value m_rest = fb.sqrt(s_rest); res.push_back(m_rest); - dets.push_back(r["det"]); + dets.push_back(det); prev_mass = m_rest; } return res; @@ -592,8 +715,10 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( auto walk = [&](const std::vector& s, + std::size_t k_other, Value P_set, Value R_b, + Value beam_own, Value beam_other, const ValueVec& rest_masses) { std::size_t k = s.size(); @@ -601,10 +726,36 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( p_out.at(s.at(0)) = P_set; return; } - Value R_a = fb.sub(P_set, R_b); auto etmin_particle = [&](std::size_t idx) { return fb.sqrt(fb.add(fb.square(m_out.at(idx)), Value(pt2(idx)))); }; + if (two_particle_step(k, k_other)) { + // AmpliCol's gen23_one_step(set(i), inext, i, 3 - i): p1 = s[0] + // with t1 = (beam_own - p1)^2, p2 = s[1] with the momentum + // transfer (p2 - beam_other)^2, the two invariants next to the + // beams in the color order. + ValueVec cond{beam_own, P_set, beam_other}; + if (_has_cut) { + cond.push_back(etmin_particle(s.at(0))); + cond.push_back(etmin_particle(s.at(1))); + cond.push_back(Value(0.)); + cond.push_back(Value(0.)); + } + auto ks = _two_to_three_beam.build_forward( + fb, + {next_discrete(), + next_random(), + next_random(), + m_out.at(s.at(0)), + m_out.at(s.at(1))}, + cond + ); + p_out.at(s.at(0)) = ks.at(0); + p_out.at(s.at(1)) = ks.at(1); + dets.push_back(ks["det"]); + return; + } + Value R_a = fb.sub(P_set, R_b); ValueVec etmin_suffix; if (_has_cut) { etmin_suffix.resize(k); @@ -677,10 +828,10 @@ Mapping::Result ColorOrderedMapping::build_forward_impl( }; if (!_set1.empty()) { - walk(_set1, P_set1, R_b_for_set1, pb, rest_masses_set1); + walk(_set1, _set2.size(), P_set1, R_b_for_set1, pa, pb, rest_masses_set1); } if (!_set2.empty()) { - walk(_set2, P_set2, R_b_for_set2, pa, rest_masses_set2); + walk(_set2, _set1.size(), P_set2, R_b_for_set2, pb, pa, rest_masses_set2); } // Assemble outputs: momentum0, momentum1 = beams; momentum_{2+i} = outgoing i. @@ -795,34 +946,28 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( m_set1 = fb.sqrt(invariants.at(idx_m2_set1)); } } else { - Value mass_sum_set1 = masses_of(_set1); - Value mass_sum_set2 = masses_of(_set2); + auto [floor1, e_min1] = set_floors(fb, _set1, m_out); + auto [floor2, e_min2] = set_floors(fb, _set2, m_out); if (_set1.size() >= 2) { - auto s_min = fb.square(mass_sum_set1); - double floor1 = cut_floor(_set1); - if (floor1 > 0.0) { - s_min = fb.max(s_min, Value(floor1)); - } - auto s_max = fb.square(fb.sub(e_cm, mass_sum_set2)); + auto [s_min, s_max] = full_system_mass_range( + fb, e_cm, floor1, e_min1, floor2, e_min2, std::nullopt + ); 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"]); + auto [r, det] = recover_mass(fb, _mass_invariant, m2_set1, s_min, s_max); + random_out.push_back(r); + dets.push_back(det); m_set1 = fb.sqrt(m2_set1); } else { m_set1 = m_out.at(_set1.at(0)); } if (_set2.size() >= 2) { - auto s_min = fb.square(mass_sum_set2); - double floor2 = cut_floor(_set2); - if (floor2 > 0.0) { - s_min = fb.max(s_min, Value(floor2)); - } - auto s_max = fb.square(fb.sub(e_cm, m_set1)); + auto [s_min, s_max] = full_system_mass_range( + fb, e_cm, floor2, e_min2, floor1, e_min1, 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"]); + auto [r, det] = recover_mass(fb, _mass_invariant, m2_set2, s_min, s_max); + random_out.push_back(r); + dets.push_back(det); m_set2 = fb.sqrt(m2_set2); } else { m_set2 = m_out.at(_set2.at(0)); @@ -904,22 +1049,11 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( // the (sqrt of the) actual recovered m2 at each step. Value prev_mass = m_set; for (std::size_t j = 0; j < k - 2; ++j) { - Value m_min = m_out.at(s.at(j + 1)); - for (std::size_t i = j + 2; i < k; ++i) { - m_min = fb.add(m_min, m_out.at(s.at(i))); - } - auto s_min = fb.square(m_min); - // Identical cut floor to the forward pass. - std::vector rest_sub(s.begin() + j + 1, s.end()); - double fl = cut_floor(rest_sub); - if (fl > 0.0) { - s_min = fb.max(s_min, Value(fl)); - } - auto s_max = fb.square(fb.sub(prev_mass, m_out.at(s.at(j)))); + auto [s_min, s_max] = rest_mass_range(fb, s, j, prev_mass, e_cm, m_out); 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"]); + auto [r, det] = recover_mass(fb, _mass_invariant, m2, s_min, s_max); + random_out.push_back(r); + dets.push_back(det); prev_mass = fb.sqrt(m2); } }; @@ -937,19 +1071,39 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( // Phase 3 inverse: peel-off walks (mirror the forward restructure: first peel // projects the own-side beam R_a, 2->3 peels project beam_other = 3-i side). auto walk_inverse = - [&]( - const std::vector& s, Value P_set, Value R_b, Value beam_other - ) { + [&](const std::vector& s, + std::size_t k_other, + Value P_set, + Value R_b, + Value beam_own, + Value beam_other) { std::size_t k = s.size(); if (k == 1) { return; } - Value R_a = fb.sub(P_set, R_b); - // ETmin suffix sums for the recoil chain (see forward walk); only built - // when cuts are active. auto etmin_particle = [&](std::size_t idx) { return fb.sqrt(fb.add(fb.square(m_out.at(idx)), Value(pt2(idx)))); }; + if (two_particle_step(k, k_other)) { + ValueVec cond{beam_own, P_set, beam_other}; + if (_has_cut) { + cond.push_back(etmin_particle(s.at(0))); + cond.push_back(etmin_particle(s.at(1))); + cond.push_back(Value(0.)); + cond.push_back(Value(0.)); + } + auto rs = _two_to_three_beam.build_inverse( + fb, {p_outgoing(s.at(0)), p_outgoing(s.at(1))}, cond + ); + discrete_out.push_back(rs.at(0)); + random_out.push_back(rs.at(1)); + random_out.push_back(rs.at(2)); + dets.push_back(rs["det"]); + return; + } + Value R_a = fb.sub(P_set, R_b); + // ETmin suffix sums for the recoil chain (see forward walk); only built + // when cuts are active. ValueVec etmin_suffix; if (_has_cut) { etmin_suffix.resize(k); @@ -1011,10 +1165,10 @@ Mapping::Result ColorOrderedMapping::build_inverse_impl( }; if (!_set1.empty()) { - walk_inverse(_set1, P_set1, R_b_for_set1, pb); + walk_inverse(_set1, _set2.size(), P_set1, R_b_for_set1, pa, pb); } if (!_set2.empty()) { - walk_inverse(_set2, P_set2, R_b_for_set2, pa); + walk_inverse(_set2, _set1.size(), P_set2, R_b_for_set2, pb, pa); } // input_types is [random0..random_{N-1}, discrete0..discrete_{M-1}], so the diff --git a/madspace/src/phasespace/phasespace.cpp b/madspace/src/phasespace/phasespace.cpp index 11fff4c9d3..0a7e43a66f 100644 --- a/madspace/src/phasespace/phasespace.cpp +++ b/madspace/src/phasespace/phasespace.cpp @@ -109,36 +109,16 @@ nested_vector2 invert_permutations(nested_vector2 perms_in) } // namespace namespace { -// Chain (color) order for the t-channel ColorOrderedMapping: the externally -// supplied order if given, else the default single chain [0, 2, ..., n+1, 1]. -std::vector ps_chain_order( - const Topology& topology, const std::optional>& color_order -) { - if (color_order) { - return *color_order; - } - std::size_t n_t_out = topology.decays().at(0).child_indices.size(); - std::vector chain; - chain.reserve(n_t_out + 2); - chain.push_back(0); - for (std::size_t i = 0; i < n_t_out; ++i) { - chain.push_back(i + 2); - } - chain.push_back(1); - return chain; -} - // Number of discrete two-solution choices for the t-channel (opt-in r_disc): -// non-zero only for color_ordered with at least one 2->3 peel. +// non-zero only for color_ordered with at least one 2->3 step. std::size_t ps_discrete_dim( const Topology& topology, PhaseSpaceMapping::TChannelMode mode, - const std::optional>& color_order + const std::optional& color_ordered ) { - if (mode == PhaseSpaceMapping::color_ordered && + if (mode == PhaseSpaceMapping::color_ordered && color_ordered && topology.t_propagator_count() >= 1) { - ColorOrderedMapping co(ps_chain_order(topology, color_order)); - return co.discrete_dim(); + return ColorOrderedMapping(*color_ordered).discrete_dim(); } return 0; } @@ -257,9 +237,10 @@ PhaseSpaceMapping::PhaseSpaceMapping( TChannelMode t_channel_mode, const std::optional& cuts, const std::vector>& permutations, - const std::optional>& color_order, + const std::optional& color_ordered, double beam_rapidity, - bool mirror_beams + bool mirror_beams, + std::optional s_hat_power ) : Mapping( "PhaseSpaceMapping", @@ -272,7 +253,7 @@ PhaseSpaceMapping::PhaseSpaceMapping( }; // Opt-in discrete channel: only declared when the t-channel strategy // actually has discrete two-solution choices (color_ordered). - std::size_t nd = ps_discrete_dim(topology, t_channel_mode, color_order); + std::size_t nd = ps_discrete_dim(topology, t_channel_mode, color_ordered); if (nd > 0) { in.push_back( "discrete", @@ -319,6 +300,17 @@ PhaseSpaceMapping::PhaseSpaceMapping( ), _two_to_one(topology.outgoing_masses().size() == 1), _t_mapping(std::monostate{}) { + // the chain order and settings belong to the color-ordered mode only + if ((t_channel_mode == PhaseSpaceMapping::color_ordered) != + color_ordered.has_value()) { + throw std::invalid_argument( + t_channel_mode == PhaseSpaceMapping::color_ordered + ? "PhaseSpaceMapping: the color_ordered mode needs color_ordered " + "options (the chain order)" + : "PhaseSpaceMapping: color_ordered options are only used by the " + "color_ordered mode" + ); + } if ((_beam_rapidity != 0. || _mirror_beams) && _topology.is_decay()) { throw std::invalid_argument("a decay has no beams to give a rapidity"); } @@ -681,7 +673,18 @@ PhaseSpaceMapping::PhaseSpaceMapping( if (!is_com_decay || _map_luminosity) { double mass = decay.width == 0. ? 0. : decay.mass; double width = decay.width; - info.invariant = Invariant(invariant_power, mass, width, decay.flat_window); + // The root is s_hat, which the PDF convolution samples. A + // color-ordered chain samples it as AmpliCol does, with + // 1/s_hat^1.2, when masses or cuts bound it from below; without a + // bound that density piles the points up at the regulator of + // Invariant (and AmpliCol falls back to a flat one). + double s_hat_default = + t_channel_mode == PhaseSpaceMapping::color_ordered && info.m_min > 0. + ? 1.2 + : invariant_power; + double power = + is_com_decay ? s_hat_power.value_or(s_hat_default) : invariant_power; + info.invariant = Invariant(power, mass, width, decay.flat_window); } } for (std::size_t index : _topology.decay_integration_order()) { @@ -710,14 +713,6 @@ PhaseSpaceMapping::PhaseSpaceMapping( _t_mapping = ChiliMapping(_topology.t_propagator_count() + 1, eta_max, pt_min); } else if (t_channel_mode == PhaseSpaceMapping::color_ordered) { - // color_order is optional in general but REQUIRED here: the chain is - // built in the externally supplied color order so the t-channel - // topology matches the known color structure of the process. - if (!color_order) { - throw std::invalid_argument( - "PhaseSpaceMapping: color_ordered mode requires a color_order" - ); - } // Reorder the per-pair cut matrices (indexed by topology outgoing // position) into the child order in which masses/pt are handed to the // chain, mirroring the pt_min reordering above. Composite (non-leaf) @@ -762,16 +757,24 @@ PhaseSpaceMapping::PhaseSpaceMapping( } // The blocks of the chain that scatter the two beams (central 2->2, // double-t, first peel of a single chain) take the rapidity bounds. - _t_mapping = ColorOrderedMapping( - ps_chain_order(topology, color_order), - invariant_power, - invariant_power, - pt_min, - m_inv_co, - dr_co, - true, - y_max - ); + // An unset chain power is 1 with cuts, as in AmpliCol, and + // invariant_power without: then the ranges of the momentum transfers + // and s23 invariants reach zero, where 1/|x| piles the points up at + // the regulator of Invariant (AmpliCol falls back to a flat density). + auto options = *color_ordered; + auto positive = [](double x) { return x > 0.; }; + bool chain_cuts = std::any_of(pt_min.begin(), pt_min.end(), positive); + for (std::size_t a = 0; a < nc; ++a) { + chain_cuts |= + std::any_of( + m_inv_co.at(a).begin(), m_inv_co.at(a).end(), positive + ) || + std::any_of(dr_co.at(a).begin(), dr_co.at(a).end(), positive); + } + if (!options.chain_invariant_power) { + options.chain_invariant_power = chain_cuts ? 1. : invariant_power; + } + _t_mapping = ColorOrderedMapping(options, pt_min, m_inv_co, dr_co, y_max); } else if (t_channel_mode == PhaseSpaceMapping::propagator || topology.t_propagator_count() < 2) { // The first scattering of the chain is between the two beams in @@ -787,7 +790,7 @@ PhaseSpaceMapping::PhaseSpaceMapping( } // Random-number budget: identical computation to the declared input shape. - _n_discrete = ps_discrete_dim(_topology, t_channel_mode, color_order); + _n_discrete = ps_discrete_dim(_topology, t_channel_mode, color_ordered); for (auto& perm : permutations) { _permutations.emplace_back(perm.begin(), perm.end()); @@ -801,9 +804,10 @@ PhaseSpaceMapping::PhaseSpaceMapping( double invariant_power, TChannelMode mode, const std::optional& cuts, - const std::optional>& color_order, + const std::optional& color_ordered, double beam_rapidity, - bool mirror_beams + bool mirror_beams, + std::optional s_hat_power ) : PhaseSpaceMapping( Topology([&] { @@ -854,9 +858,10 @@ PhaseSpaceMapping::PhaseSpaceMapping( mode, cuts, {}, - color_order, + color_ordered, beam_rapidity, - mirror_beams + mirror_beams, + s_hat_power ) {} Mapping::Result PhaseSpaceMapping::build_forward_impl( diff --git a/madspace/src/phasespace/three_particle.cpp b/madspace/src/phasespace/three_particle.cpp index 965b2088db..98b0f4ac9c 100644 --- a/madspace/src/phasespace/three_particle.cpp +++ b/madspace/src/phasespace/three_particle.cpp @@ -1,6 +1,7 @@ #include "madspace/phasespace/three_particle.hpp" #include +#include using namespace madspace; @@ -100,7 +101,9 @@ TwoToThreeParticleScattering::TwoToThreeParticleScattering( double s_width, bool has_cut, bool arcsine_s23, - bool p12_condition + bool p12_condition, + bool incoming_recoil, + bool cut_aware_branch ) : Mapping( "TwoToThreeParticleScattering", @@ -131,7 +134,9 @@ TwoToThreeParticleScattering::TwoToThreeParticleScattering( _s_mass(s_mass), _has_cut(has_cut), _arcsine_s23(arcsine_s23), - _p12_condition(p12_condition) { + _p12_condition(p12_condition), + _incoming_recoil(incoming_recoil), + _cut_aware_branch(cut_aware_branch && has_cut) { // The block peels a particle off a t-channel chain, where no s-channel // resonance is expected; the arcsine sampling covers power laws only. if (arcsine_s23 && s_width != 0.) { @@ -140,8 +145,19 @@ TwoToThreeParticleScattering::TwoToThreeParticleScattering( "Breit-Wigner s23 sampling (s_width != 0)" ); } + if (incoming_recoil && !p12_condition) { + throw std::invalid_argument( + "TwoToThreeParticleScattering: incoming_recoil requires p12_condition" + ); + } } +namespace { + +Value negated(FunctionBuilder& fb, Value x) { return fb.sub(fb.sub(x, x), x); } + +} // namespace + std::array TwoToThreeParticleScattering::split_conditions( FunctionBuilder& fb, const NamedVector& conditions ) const { @@ -158,6 +174,30 @@ std::array TwoToThreeParticleScattering::split_conditions( return {p_a, fb.sub(fb.add(p_a, p_b), p_3), fb.sub(p_b, p_3)}; } +std::array TwoToThreeParticleScattering::sampling_ranges( + FunctionBuilder& fb, + Value s23_min, + Value s23_max, + const NamedVector& conditions +) const { + // returns the sampled and the kinematic range of the sampled variable + if (!_incoming_recoil) { + return {s23_min, s23_max, s23_min, s23_max}; + } + // v = -(p2 - k)^2 >= 0 runs over [-s23_max, -s23_min] + auto v_phys_min = negated(fb, s23_max), v_phys_max = negated(fb, s23_min); + auto v_min = v_phys_min, v_max = v_phys_max; + if (_has_cut) { + // |(p2 - k)^2| >= m_T(p2)^2 for a massless k along the beam axis + v_min = fb.max(v_min, fb.square(conditions.at(4))); + // an empty range keeps a small width; the weight is zero there, as + // the point then lies outside the kinematic range + auto width = fb.max(fb.mul(v_min, Value(1e-10)), Value(1e-300)); + v_max = fb.max(v_max, fb.add(v_min, width)); + } + return {v_min, v_max, v_phys_min, v_phys_max}; +} + Mapping::Result TwoToThreeParticleScattering::build_forward_impl( FunctionBuilder& fb, const NamedVector& inputs, @@ -165,67 +205,81 @@ Mapping::Result TwoToThreeParticleScattering::build_forward_impl( ) const { auto index_choice = inputs.at(0), r_s23 = inputs.at(1), r_t1 = inputs.at(2), m1 = inputs.at(3), m2 = inputs.at(4); - auto p_3 = conditions.at(2); + auto p_3 = _incoming_recoil ? negated(fb, conditions.at(2)) : conditions.at(2); auto [p_a, p_12, p_c] = split_conditions(fb, conditions); auto [t1_min, t1_max] = _has_cut ? fb.t_inv_min_max_cut(p_a, p_c, m1, m2, conditions.at(3), conditions.at(4)) : fb.t_inv_min_max(p_a, p_c, m1, m2); auto t_inv_result = _t_invariant.build_forward(fb, {r_t1}, {t1_min, t1_max}); - auto [s23_min, s23_max] = _has_cut - ? fb.s23_min_max_cut( - p_a, - p_12, - p_3, - t_inv_result["invariant"], - m1, - m2, - conditions.at(3), - conditions.at(4), - conditions.at(5), - conditions.at(6) - ) - : fb.s23_min_max(p_a, p_12, p_3, t_inv_result["invariant"], m1, m2); + auto t1 = t_inv_result["invariant"]; // The scattering kernel takes the position of s23 in its kinematic range, // u and 1 - u. Without cuts the sampled and the kinematic range coincide. - auto [s23_phys_min, s23_phys_max] = _has_cut - ? fb.s23_min_max(p_a, p_12, p_3, t_inv_result["invariant"], m1, m2) - : std::array{s23_min, s23_max}; + std::array ranges; + if (_incoming_recoil) { + auto [lo, hi] = fb.s23_min_max(p_a, p_12, p_3, t1, m1, m2); + ranges = sampling_ranges(fb, lo, hi, conditions); + } else if (_has_cut) { + auto [s_min, s_max] = fb.s23_min_max_cut( + p_a, + p_12, + p_3, + t1, + m1, + m2, + conditions.at(3), + conditions.at(4), + conditions.at(5), + conditions.at(6) + ); + auto [lo, hi] = fb.s23_min_max(p_a, p_12, p_3, t1, m1, m2); + ranges = {s_min, s_max, lo, hi}; + } else { + auto [lo, hi] = fb.s23_min_max(p_a, p_12, p_3, t1, m1, m2); + ranges = {lo, hi, lo, hi}; + } + auto [v_min, v_max, v_phys_min, v_phys_max] = ranges; Value u, u_c, det_s23; if (_arcsine_s23) { // arcsine map and importance sampling in one step, carrying u and // 1 - u at full relative precision up to the edges auto [u_out, u_c_out, det_out] = fb.s23_arcsine_sample( - r_s23, - s23_min, - s23_max, - s23_phys_min, - s23_phys_max, - Value(_s_power), - Value(_s_mass) + r_s23, v_min, v_max, v_phys_min, v_phys_max, Value(_s_power), Value(_s_mass) ); u = u_out; u_c = u_c_out; det_s23 = det_out; } else { - auto s23_inv_result = - _s_invariant.build_forward(fb, {r_s23}, {s23_min, s23_max}); + auto s23_inv_result = _s_invariant.build_forward(fb, {r_s23}, {v_min, v_max}); det_s23 = s23_inv_result["det"]; auto [u_out, u_c_out] = - fb.s23_position(s23_inv_result["invariant"], s23_phys_min, s23_phys_max); + fb.s23_position(s23_inv_result["invariant"], v_phys_min, v_phys_max); u = u_out; u_c = u_c_out; } + if (_incoming_recoil) { + // v = -s23 runs the other way through the range of s23 + std::swap(u, u_c); + } auto det_inv = fb.mul(t_inv_result["det"], det_s23); + if (_cut_aware_branch) { + auto [index_pass, factor, inv_factor] = fb.two_to_three_branch_choice( + index_choice, + p_a, + p_12, + p_3, + u, + u_c, + t1, + m1, + m2, + conditions.at(3), + conditions.at(4) + ); + index_choice = index_pass; + det_inv = fb.mul(det_inv, factor); + } auto [p1, p2, det_scatter] = fb.two_to_three_particle_scattering( - index_choice, - p_a, - p_12, - p_3, - u, - u_c, - t_inv_result["invariant"], - m1, - m2 + index_choice, p_a, p_12, p_3, u, u_c, t1, m1, m2 ); return {{{"momentum1", p1}, {"momentum2", p2}}, fb.mul(det_inv, det_scatter)}; } @@ -236,7 +290,7 @@ Mapping::Result TwoToThreeParticleScattering::build_inverse_impl( const NamedVector& conditions ) const { auto p1 = inputs.at(0), p2 = inputs.at(1); - auto p_3 = conditions.at(2); + auto p_3 = _incoming_recoil ? negated(fb, conditions.at(2)) : conditions.at(2); auto [p_a, p_12, p_c] = split_conditions(fb, conditions); auto [t1_abs, t1_min, t1_max] = _has_cut ? fb.t_inv_value_and_min_max_cut( @@ -244,50 +298,82 @@ Mapping::Result TwoToThreeParticleScattering::build_inverse_impl( ) : fb.t_inv_value_and_min_max(p_a, p_c, p1, p2); auto t_inv_result = _t_invariant.build_inverse(fb, {t1_abs}, {t1_min, t1_max}); - auto [s23, s23_min, s23_max] = _has_cut - ? fb.s23_value_and_min_max_cut( - p_a, - p_3, - t1_abs, - p1, - p2, - conditions.at(3), - conditions.at(4), - conditions.at(5), - conditions.at(6) - ) - : fb.s23_value_and_min_max(p_a, p_3, t1_abs, p1, p2); + Value v; + std::array ranges; + if (_incoming_recoil) { + auto [s23, lo, hi] = fb.s23_value_and_min_max(p_a, p_3, t1_abs, p1, p2); + v = negated(fb, s23); + ranges = sampling_ranges(fb, lo, hi, conditions); + } else if (_has_cut) { + auto [s23, s_min, s_max] = fb.s23_value_and_min_max_cut( + p_a, + p_3, + t1_abs, + p1, + p2, + conditions.at(3), + conditions.at(4), + conditions.at(5), + conditions.at(6) + ); + v = s23; + ranges = {s_min, s_max, s_min, s_max}; + if (_arcsine_s23) { + auto [s23_phys, lo, hi] = + fb.s23_value_and_min_max(p_a, p_3, t1_abs, p1, p2); + ranges = {s_min, s_max, lo, hi}; + } + } else { + auto [s23, lo, hi] = fb.s23_value_and_min_max(p_a, p_3, t1_abs, p1, p2); + v = s23; + ranges = {lo, hi, lo, hi}; + } + auto [v_min, v_max, v_phys_min, v_phys_max] = ranges; auto [m1, m2, index_choice, u, u_c, det_scatter] = fb.two_to_three_particle_scattering_inverse(p1, p2, p_3, p_a, p_12, t1_abs); - Value s23_phys_min = s23_min, s23_phys_max = s23_max; - if (_has_cut && _arcsine_s23) { - auto [s23_phys, phys_min, phys_max] = - fb.s23_value_and_min_max(p_a, p_3, t1_abs, p1, p2); - s23_phys_min = phys_min; - s23_phys_max = phys_max; + Value det_branch = Value(1.); + if (_cut_aware_branch) { + auto [index_pass, factor, inv_factor] = fb.two_to_three_branch_choice( + index_choice, + p_a, + p_12, + p_3, + u, + u_c, + t1_abs, + m1, + m2, + conditions.at(3), + conditions.at(4) + ); + det_branch = inv_factor; } + // position of the sampled variable in its kinematic range + Value u_v = _incoming_recoil ? u_c : u, u_c_v = _incoming_recoil ? u : u_c; Value r_s23, det_s23; if (_arcsine_s23) { // u and 1 - u read off the azimuth, at full relative precision auto [r, det_r] = fb.s23_arcsine_sample_inverse( - u, - u_c, - s23_min, - s23_max, - s23_phys_min, - s23_phys_max, + u_v, + u_c_v, + v_min, + v_max, + v_phys_min, + v_phys_max, Value(_s_power), Value(_s_mass) ); r_s23 = r; det_s23 = det_r; } else { - auto s23_inv_result = - _s_invariant.build_inverse(fb, {s23}, {s23_min, s23_max}); + auto s23_inv_result = _s_invariant.build_inverse(fb, {v}, {v_min, v_max}); r_s23 = s23_inv_result["random"]; det_s23 = s23_inv_result["det"]; } auto det_inv = fb.mul(t_inv_result["det"], det_s23); + if (_cut_aware_branch) { + det_inv = fb.mul(det_inv, det_branch); + } return { {{"discrete_choice", index_choice}, {"random_s23", r_s23}, diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp index dae9db04f3..ceb6955a8f 100644 --- a/madspace/src/python/madspace.cpp +++ b/madspace/src/python/madspace.cpp @@ -723,7 +723,18 @@ PYBIND11_MODULE(_madspace_py, m) { m, "TwoToThreeParticleScattering", pydoc::doc("TwoToThreeParticleScattering") ) .def( - py::init(), + py::init< + double, + double, + double, + double, + double, + double, + bool, + bool, + bool, + bool, + bool>(), py::arg("t_invariant_power") = 0., py::arg("t_mass") = 0., py::arg("t_width") = 0., @@ -732,6 +743,9 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("s_width") = 0., py::arg("has_cut") = false, py::arg("arcsine_s23") = true, + py::arg("p12_condition") = false, + py::arg("incoming_recoil") = false, + py::arg("cut_aware_branch") = false, pydoc::doc("TwoToThreeParticleScattering::TwoToThreeParticleScattering") ) .def( @@ -788,26 +802,51 @@ PYBIND11_MODULE(_madspace_py, m) { pydoc::doc("TPropagatorMapping::random_dim") ); - py::classh( + auto color_ordered_mapping = py::classh( m, "ColorOrderedMapping", pydoc::doc("ColorOrderedMapping") + ); + py::class_( + color_ordered_mapping, "Options", pydoc::doc("ColorOrderedMapping::Options") ) + .def( + py::init([](const std::vector& color_order, + std::optional chain_invariant_power, + bool cut_aware_branch) { + return ColorOrderedMapping::Options{ + color_order, chain_invariant_power, cut_aware_branch + }; + }), + py::arg("color_order"), + py::arg("chain_invariant_power") = std::nullopt, + py::arg("cut_aware_branch") = true + ) + .def_readwrite( + "color_order", + &ColorOrderedMapping::Options::color_order, + pydoc::doc("ColorOrderedMapping::Options::color_order") + ) + .def_readwrite( + "chain_invariant_power", + &ColorOrderedMapping::Options::chain_invariant_power, + pydoc::doc("ColorOrderedMapping::Options::chain_invariant_power") + ) + .def_readwrite( + "cut_aware_branch", + &ColorOrderedMapping::Options::cut_aware_branch, + pydoc::doc("ColorOrderedMapping::Options::cut_aware_branch") + ); + color_ordered_mapping .def( py::init< - std::vector, - double, - double, + const ColorOrderedMapping::Options&, std::vector, std::vector>, std::vector>, - bool, std::vector>(), - py::arg("color_order"), - py::arg("t_invariant_power") = 0.8, - py::arg("s_invariant_power") = 0.8, + py::arg("options"), py::arg("pt_min") = std::vector{}, py::arg("m_inv_min") = std::vector>{}, py::arg("dr_min") = std::vector>{}, - py::arg("arcsine_s23") = true, py::arg("y_max") = std::vector{}, pydoc::doc("ColorOrderedMapping::ColorOrderedMapping") ) @@ -1264,9 +1303,10 @@ PYBIND11_MODULE(_madspace_py, m) { PhaseSpaceMapping::TChannelMode, const std::optional&, const nested_vector2&, - const std::optional>&, + const std::optional&, double, - bool>(), + bool, + std::optional>(), py::arg("topology"), py::arg("cm_energy"), py::arg("leptonic") = false, @@ -1274,9 +1314,10 @@ PYBIND11_MODULE(_madspace_py, m) { py::arg("t_channel_mode") = PhaseSpaceMapping::propagator, py::arg("cuts") = std::nullopt, py::arg("permutations") = nested_vector2{}, - py::arg("color_order") = std::nullopt, + py::arg("color_ordered") = std::nullopt, py::arg("beam_rapidity") = 0., py::arg("mirror_beams") = false, + py::arg("s_hat_power") = std::nullopt, pydoc::doc("PhaseSpaceMapping::PhaseSpaceMapping") ) .def( @@ -1287,18 +1328,20 @@ PYBIND11_MODULE(_madspace_py, m) { double, PhaseSpaceMapping::TChannelMode, std::optional, - const std::optional>&, + const std::optional&, double, - bool>(), + bool, + std::optional>(), py::arg("external_masses"), py::arg("cm_energy"), py::arg("leptonic") = false, py::arg("invariant_power") = 0.8, py::arg("mode") = PhaseSpaceMapping::rambo, py::arg("cuts") = std::nullopt, - py::arg("color_order") = std::nullopt, + py::arg("color_ordered") = std::nullopt, py::arg("beam_rapidity") = 0., py::arg("mirror_beams") = false, + py::arg("s_hat_power") = std::nullopt, pydoc::doc("PhaseSpaceMapping::PhaseSpaceMapping#2") ) .def( diff --git a/madspace/tests/test_color_ordered_amplicol_steps.py b/madspace/tests/test_color_ordered_amplicol_steps.py new file mode 100644 index 0000000000..6dca3fbfaa --- /dev/null +++ b/madspace/tests/test_color_ordered_amplicol_steps.py @@ -0,0 +1,281 @@ +"""ColorOrderedMapping steps aligned with AmpliCol's gen23 (issue #292). + +- A side of exactly two particles, with the other side not empty, is one 2->3 + step with the opposite beam as the previous emission: both invariants next + to the beams in the color order are sampled, instead of one and an azimuth. + The block then samples a momentum transfer with an incoming recoil. +- Composite and rest masses are sampled with 1/sqrt(s) (flat in the mass), + within the bounds the energies of the two sides set. +- With cut_aware_branch the 2->3 steps keep the azimuthal solution that passes + the pt cuts, and pick at random only when both do. +- The fiducial volume agrees with RAMBO's, whatever the options. +- The function builder no longer folds 0 - x to x. +""" + +import math + +import pytest +import torch +from madspace_testing import default_context, generator, rand, randint + +import madspace as ms + +E_BEAM = 500.0 +PA = torch.tensor([E_BEAM, 0.0, 0.0, E_BEAM]) +PB = torch.tensor([E_BEAM, 0.0, 0.0, -E_BEAM]) +CUTS = dict( + pt_min=[20.0] * 4, + m_inv_min=[[0.0 if i == j else 10.0 for j in range(4)] for i in range(4)], + dr_min=[[0.0 if i == j else 0.4 for j in range(4)] for i in range(4)], +) + + +def lsquare(p): + return p[..., 0] ** 2 - torch.sum(p[..., 1:] ** 2, dim=-1) + + +def test_sub_of_zero_literal(): + """sub(0, x) was folded to x, which turned the bounds built from massless + literal masses upside down.""" + fb = ms.FunctionBuilder( + ms.NamedTypes([("x", ms.batch_float)]), + ms.NamedTypes([("minus_x", ms.batch_float), ("x_minus_0", ms.batch_float)]), + ) + x = fb.input(0) + fb.output(0, fb.sub(ms.Value(0.0), x)) + fb.output(1, fb.sub(x, ms.Value(0.0))) + runtime = ms.FunctionRuntime(fb.function(), default_context()) + values = torch.tensor([1.0, -2.0, 3.5]) + minus_x, x_minus_0 = [torch.from_dlpack(t) for t in runtime.call([values])] + torch.testing.assert_close(minus_x, -values) + torch.testing.assert_close(x_minus_0, values) + + +@pytest.mark.parametrize( + "order, n_discrete, n_random", + [ + ([0, 2, 3, 1, 4, 5], 2, 8), + ([0, 2, 3, 1, 4], 1, 5), + ([0, 2, 1, 3, 4], 1, 5), + ([0, 2, 3, 4, 1, 5], 1, 8), + ([0, 2, 3, 1], 0, 2), + ], + ids=["2+2", "2+1", "1+2", "3+1", "2+0"], +) +def test_discrete_dim(order, n_discrete, n_random): + """One discrete choice per side of two particles next to a non-empty side + (a 2->3 step instead of a 2->2 peel, with the same random numbers); a + single chain of two keeps its 2->2 step, as in AmpliCol.""" + mapping = ms.ColorOrderedMapping(ms.ColorOrderedMapping.Options(order)) + assert mapping.discrete_dim() == n_discrete + assert mapping.random_dim() == n_random + + +def beam_block(r_s23, r_t1, arcsine, has_cut=False, s_power=0.0): + """The 2->3 step of a side of two massless particles with the incoming + recoil pb, for a side momentum p12 recoiling against the other side.""" + mapping = ms.TwoToThreeParticleScattering( + 0.8, + 0.0, + 0.0, + s_power, + 0.0, + 0.0, + has_cut=has_cut, + arcsine_s23=arcsine, + p12_condition=True, + incoming_recoil=True, + ) + n = len(r_s23) + # p12 = pa + pb - k, k a massless momentum of the other side + k = torch.tensor([300.0, 300.0 * math.sin(0.9), 0.0, 300.0 * math.cos(0.9)]) + p12 = PA + PB - k + inputs = [ + torch.zeros(n, dtype=torch.int32), + r_s23, + r_t1, + torch.zeros(n), + torch.zeros(n), + ] + conditions = [PA.repeat(n, 1), p12.repeat(n, 1), PB.repeat(n, 1)] + if has_cut: + conditions += [torch.full((n,), 20.0), torch.full((n,), 20.0)] + conditions += [torch.zeros(n), torch.zeros(n)] + p1, p2, det = mapping.map_forward(inputs, conditions) + return mapping, inputs, conditions, p1, p2, det, p12 + + +def test_beam_step_kinematics(): + """p1 + p2 = p12 on shell, t1 = (pa - p1)^2 set by r_t1 alone, and the + sampled invariant (p2 - pb)^2 runs monotonically with r_s23.""" + r_s23 = torch.linspace(0.01, 0.99, 50) + r_t1 = torch.full((50,), 0.37) + *_, p1, p2, det, p12 = beam_block(r_s23, r_t1, arcsine=False) + assert torch.all(det > 0) + torch.testing.assert_close(p1 + p2, p12.repeat(50, 1), rtol=0, atol=1e-10) + torch.testing.assert_close(lsquare(p1), torch.zeros(50), rtol=0, atol=1e-8) + torch.testing.assert_close(lsquare(p2), torch.zeros(50), rtol=0, atol=1e-8) + t1 = lsquare(PA - p1) + torch.testing.assert_close(t1, t1[0].expand(50), rtol=1e-10, atol=0) + v = -lsquare(p2 - PB) + assert torch.all(v > 0) + assert torch.all(torch.diff(v) > 0) + # flat sampling of v: equally spaced in r + torch.testing.assert_close( + torch.diff(v), torch.diff(v)[0].expand(49), rtol=1e-8, atol=0 + ) + + +def test_beam_step_flat_weight_with_arcsine(): + """The measure is flat in the azimuth: with the arcsine map and flat + sampling the weight does not depend on r_s23, as for an outgoing recoil.""" + edge = 10.0 ** -torch.arange(1.0, 6.0) + r_s23 = torch.cat([edge, torch.tensor([0.5]), 1.0 - edge]) + r_t1 = torch.full((len(r_s23),), 0.37) + *_, det, _ = beam_block(r_s23, r_t1, arcsine=True) + assert torch.all(torch.isfinite(det)) and torch.all(det > 0) + torch.testing.assert_close( + det / det[len(edge)], torch.ones_like(det), rtol=1e-6, atol=0 + ) + + +@pytest.mark.parametrize("arcsine", [True, False], ids=["arcsine", "flat"]) +@pytest.mark.parametrize("has_cut", [False, True], ids=["no cut", "cut"]) +@pytest.mark.parametrize("s_power", [0.0, 0.8]) +def test_beam_step_round_trip(arcsine, has_cut, s_power): + g = generator(11) + r = rand((20_000, 2), generator=g) + mapping, inputs, conditions, p1, p2, det, _ = beam_block( + r[:, 0].contiguous(), r[:, 1].contiguous(), arcsine, has_cut, s_power + ) + ok = det > 0 + assert ok.float().mean() > 0.5 + *inv, det_inv = mapping.map_inverse([p1, p2], conditions) + # without the arcsine map the inverse goes through the invariant itself, + # which leaves ~1e-8 near the edges of its range + torch.testing.assert_close( + inv[1][ok], r[ok, 0], rtol=0, atol=1e-8 if arcsine else 1e-7 + ) + torch.testing.assert_close(inv[2][ok], r[ok, 1], rtol=0, atol=1e-8) + assert torch.max(torch.abs(det[ok] * det_inv[ok] - 1.0)) < 1e-7 + + +def passes(p, cuts): + """Outgoing momenta p (batch, n_out, 4) inside the cuts.""" + n_out = p.shape[1] + pt = torch.hypot(p[..., 1], p[..., 2]) + ok = torch.all(pt > cuts["pt_min"][0], dim=1) + eta = torch.asinh(p[..., 3] / torch.clamp(pt, min=1e-300)) + phi = torch.atan2(p[..., 2], p[..., 1]) + for i in range(n_out): + for j in range(i + 1, n_out): + q = p[:, i] + p[:, j] + dphi = phi[:, i] - phi[:, j] + dphi = torch.atan2(torch.sin(dphi), torch.cos(dphi)) + ok &= lsquare(q) > cuts["m_inv_min"][i][j] ** 2 + ok &= torch.hypot(eta[:, i] - eta[:, j], dphi) > cuts["dr_min"][i][j] + return ok + + +def sample(order, seed, n, cuts=None, **options): + """Weights of ColorOrderedMapping for the massless phase space at 1 TeV, + zero outside the cuts if given.""" + cuts = cuts or {} + cm_energy = 1000.0 + n_out = len(order) - 2 + mapping = ms.ColorOrderedMapping( + ms.ColorOrderedMapping.Options(order, **options), **cuts + ) + g = generator(seed) + r = rand((n, mapping.random_dim()), generator=g) + d = randint(0, 2, (n, mapping.discrete_dim()), generator=g, dtype=torch.int32) + inputs = [r[:, i].contiguous() for i in range(r.shape[1])] + inputs += [d[:, j].contiguous() for j in range(d.shape[1])] + conditions = [torch.full((n,), cm_energy)] + [torch.zeros(n)] * n_out + *momenta, det = mapping.map_forward(inputs, conditions) + p = torch.stack(momenta, dim=1)[:, 2:] + det = det * 2.0 ** mapping.discrete_dim() + ok = torch.isfinite(p).all(dim=2).all(dim=1) & torch.isfinite(det) + if cuts: + ok &= passes(p, cuts) + return torch.where(ok, det, 0.0) + + +@pytest.mark.parametrize( + "order", + [[0, 2, 3, 1, 4, 5], [0, 2, 3, 1, 4], [0, 2, 1, 3, 4]], + ids=["2+2", "2+1", "1+2"], +) +def test_massless_volume(order): + """The analytic massless volume, (pi/2)^(n-1) s^(n-2) / ((n-1)! (n-2)!).""" + n = len(order) - 2 + volume = (math.pi / 2) ** (n - 1) * 1000.0 ** (2 * n - 4) + volume /= math.factorial(n - 1) * math.factorial(n - 2) + w = torch.cat([sample(order, seed, 1_000_000) for seed in range(400, 404)]) + err = w.std(correction=0).item() / math.sqrt(len(w)) + assert abs(w.mean().item() - volume) < 4.0 * err + assert err / volume < 0.01 + + +def rambo_acceptance(n_out, n): + """Fraction of the massless phase space at 1 TeV inside the cuts, from + RAMBO (a constant weight).""" + mapping = ms.PhaseSpaceMapping( + [0.0] * (n_out + 2), 1000.0, leptonic=True, mode=ms.PhaseSpaceMapping.rambo + ) + g = generator(600) + p = mapping.map_forward([rand((n, mapping.random_dim()), generator=g)])[0] + frac = passes(p[:, 2:], CUTS).double().mean().item() + return frac, math.sqrt(frac * (1.0 - frac) / n) + + +@pytest.mark.parametrize( + "order, options", + [ + ([0, 2, 3, 1, 4, 5], {}), + ([0, 2, 3, 4, 5, 1], {}), + ([0, 2, 3, 4, 1, 5], {}), + ([0, 2, 3, 1, 4, 5], dict(cut_aware_branch=False)), + ([0, 2, 3, 4, 5, 1], dict(cut_aware_branch=False, chain_invariant_power=0.7)), + ], + ids=["2+2", "chain", "3+1", "2+2 no branch", "chain no branch, 0.7"], +) +def test_cut_volume(order, options): + """With cuts, the fiducial volume is the massless volume times the + fraction RAMBO finds inside the cuts, whatever the options.""" + n_out = len(order) - 2 + volume = (math.pi / 2) ** (n_out - 1) * 1000.0 ** (2 * n_out - 4) + volume /= math.factorial(n_out - 1) * math.factorial(n_out - 2) + frac, frac_err = rambo_acceptance(n_out, 2_000_000) + n = 1_000_000 + w = torch.cat([sample(order, s, n, CUTS, **options) for s in range(500, 504)]) + ratio = w.mean().item() / volume + err = math.hypot(w.std(correction=0).item() / math.sqrt(len(w)) / volume, frac_err) + assert abs(ratio - frac) < 4.0 * err + + +@pytest.mark.parametrize( + "order", [[0, 2, 3, 1, 4, 5], [0, 2, 3, 4, 5, 1]], ids=["2+2", "chain"] +) +def test_round_trip_with_cuts(order): + """Forward and inverse agree with the new steps and the branch choice; + where only one solution passes, the inverse returns the one taken.""" + n = 200_000 + mapping = ms.ColorOrderedMapping(ms.ColorOrderedMapping.Options(order), **CUTS) + g = generator(21) + r = rand((n, mapping.random_dim()), generator=g) + d = randint(0, 2, (n, mapping.discrete_dim()), generator=g, dtype=torch.int32) + inputs = [r[:, i].contiguous() for i in range(r.shape[1])] + inputs += [d[:, j].contiguous() for j in range(d.shape[1])] + conditions = [torch.full((n,), 1000.0)] + [torch.zeros(n)] * 4 + *momenta, det = mapping.map_forward(inputs, conditions) + out = mapping.map_inverse(momenta, conditions) + r_inv = torch.stack(out[: mapping.random_dim()], dim=1) + ok = det > 0 + rt = torch.abs(det * out[-1] - 1.0)[ok] + assert torch.mean((rt > 1e-2).double()) < 1e-4 + assert torch.quantile(rt, 0.999) < 1e-6 + assert ( + torch.mean((torch.max(torch.abs(r_inv - r), dim=1).values[ok] > 1e-6).double()) + < 1e-4 + ) diff --git a/madspace/tests/test_color_ordered_pair_cuts.py b/madspace/tests/test_color_ordered_pair_cuts.py index f2d58a3de5..9f304856ee 100644 --- a/madspace/tests/test_color_ordered_pair_cuts.py +++ b/madspace/tests/test_color_ordered_pair_cuts.py @@ -79,7 +79,7 @@ def psm_momenta(order, cut_items, r, d): leptonic=True, mode=PSM.color_ordered, cuts=ms.Cuts(cut_items), - color_order=order, + color_ordered=ms.ColorOrderedMapping.Options(order), ) assert mapping.random_dim() == r.shape[1] assert mapping.discrete_dim() == d.shape[1] @@ -88,7 +88,9 @@ def psm_momenta(order, cut_items, r, d): def co_momenta(order, pt_min, m_inv_min, dr_min, r, d): - mapping = ms.ColorOrderedMapping(order, 0.8, 0.8, pt_min, m_inv_min, dr_min) + mapping = ms.ColorOrderedMapping( + ms.ColorOrderedMapping.Options(order), pt_min, m_inv_min, dr_min + ) n = r.shape[0] inputs = [r[:, i].contiguous() for i in range(r.shape[1])] inputs += [d[:, j].contiguous() for j in range(d.shape[1])] @@ -137,7 +139,7 @@ def test_pair_cut_reaches_color_ordered_mapping( """PhaseSpaceMapping must hand ColorOrderedMapping the full pair-cut matrices: its momenta equal those of a ColorOrderedMapping built by hand with the cut, and differ from one built without it.""" - probe = ms.ColorOrderedMapping(order) + probe = ms.ColorOrderedMapping(ms.ColorOrderedMapping.Options(order)) r, d = random_inputs(probe.random_dim(), probe.discrete_dim(), 2000, SEED) p_psm = psm_momenta(order, cut_items, r, d) @@ -163,7 +165,7 @@ def sample(order, cuts, n, seed): leptonic=True, mode=PSM.color_ordered, cuts=cuts, - color_order=order, + color_ordered=ms.ColorOrderedMapping.Options(order), ) r, d = random_inputs(mapping.random_dim(), mapping.discrete_dim(), n, seed) inputs = [r, d] if mapping.discrete_dim() else [r] diff --git a/madspace/tests/test_color_ordered_precision.py b/madspace/tests/test_color_ordered_precision.py index 63e4b2dc5b..1e28475359 100644 --- a/madspace/tests/test_color_ordered_precision.py +++ b/madspace/tests/test_color_ordered_precision.py @@ -27,7 +27,9 @@ def run(color_order, cm_energy, seed, cuts=None): cuts = cuts or {} n_out = len(color_order) - 2 - mapping = ms.ColorOrderedMapping(color_order, **cuts) + mapping = ms.ColorOrderedMapping( + ms.ColorOrderedMapping.Options(color_order), **cuts + ) g = generator(seed) r = rand((N, mapping.random_dim()), generator=g) d = randint(0, 2, (N, mapping.discrete_dim()), generator=g, dtype=torch.int32) @@ -48,9 +50,12 @@ def test_massless_on_shell(color_order): *_, momenta, _, det = run(color_order, cm_energy, 5) p = torch.stack(momenta, dim=1)[:, 2:] p = p[det > 0] - m2 = p[..., 0] ** 2 - torch.sum(p[..., 1:] ** 2, dim=-1) - # sqrt(|m^2|) / E is ~1e-8 from rounding alone; it reached 4e-5 before - assert torch.max(torch.sqrt(torch.abs(m2)) / p[..., 0]) < 1e-6 + # |p^2| / E^2 from (E - |p|) (E + |p|): ~1e-16 from rounding, up to ~1e-12 + # for particles with E ~ 1e-8 sqrt(s), which the 1/sqrt(s) sampling of the + # rest masses reaches; it was up to ~2e-9 before + p_abs = torch.sqrt(torch.sum(p[..., 1:] ** 2, dim=-1)) + m2 = (p[..., 0] - p_abs) * (p[..., 0] + p_abs) + assert torch.max(torch.abs(m2) / p[..., 0] ** 2) < 1e-10 # momentum conservation to a few ulps of the total energy p_sum = torch.sum(p, dim=1) p_sum[:, 0] -= cm_energy @@ -87,7 +92,7 @@ def test_lab_frame_on_shell(masses): masses, cm_energy, mode=ms.PhaseSpaceMapping.color_ordered, - color_order=[0, 2, 3, 4, 5, 1], + color_ordered=ms.ColorOrderedMapping.Options([0, 2, 3, 4, 5, 1]), ) g = generator(3) inputs = [rand((N, mapping.random_dim()), generator=g)] diff --git a/madspace/tests/test_cut_rapidity_windows.py b/madspace/tests/test_cut_rapidity_windows.py index 46554d0f53..5c7c572a50 100644 --- a/madspace/tests/test_cut_rapidity_windows.py +++ b/madspace/tests/test_cut_rapidity_windows.py @@ -568,7 +568,7 @@ def mapping(c): CM_ENERGY, mode=ms.PhaseSpaceMapping.color_ordered, cuts=c, - color_order=order, + color_ordered=ms.ColorOrderedMapping.Options(order), ) n = 200_000 @@ -604,7 +604,7 @@ def test_color_ordered_blocks_respect_the_bounds(name): def outside(y_max): mapping = ms.ColorOrderedMapping( - order, 0.8, 0.8, [PT_JET] * 4, [], [], True, y_max + ms.ColorOrderedMapping.Options(order), [PT_JET] * 4, [], [], y_max ) r = rand((n, mapping.random_dim()), generator=g) d = randint(0, 2, (n, mapping.discrete_dim()), generator=g, dtype=torch.int32) diff --git a/madspace/tests/test_cuts.py b/madspace/tests/test_cuts.py index 2ea20ce91b..f9d278ccf6 100644 --- a/madspace/tests/test_cuts.py +++ b/madspace/tests/test_cuts.py @@ -125,7 +125,9 @@ def _cut_volume( CM_ENERGY, mode=mode, leptonic=True, - color_order=color_order, + color_ordered=( + None if color_order is None else ms.ColorOrderedMapping.Options(color_order) + ), cuts=_cuts(pids, pt_min) if piped else None, ) diff --git a/madspace/tests/test_doc_convention.py b/madspace/tests/test_doc_convention.py index 83dc4b904f..6b3b0b7dc1 100644 --- a/madspace/tests/test_doc_convention.py +++ b/madspace/tests/test_doc_convention.py @@ -51,12 +51,11 @@ "FastRamboMapping": dict(n_particles=4, massless=True, com=True), "TPropagatorMapping": dict(integration_order=[0], invariant_power=0.8, pt_min=[]), "ColorOrderedMapping": dict( - color_order=[0, 1, 2, 3], - t_invariant_power=0.8, - s_invariant_power=0.8, + options=ms.ColorOrderedMapping.Options([0, 1, 2, 3]), pt_min=[], m_inv_min=[], dr_min=[], + y_max=[], ), "MultiChannelMapping": dict( mappings=[ms.TwoBodyDecay(com=True), ms.TwoBodyDecay(com=True)] @@ -132,7 +131,7 @@ def test_phasespace_mapping_both_constructors(): invariant_power=0.8, mode=ms.PhaseSpaceMapping.TChannelMode.rambo, cuts=None, - color_order=None, + color_ordered=None, ) diff --git a/madspace/tests/test_s_hat_power.py b/madspace/tests/test_s_hat_power.py new file mode 100644 index 0000000000..0523ee2b93 --- /dev/null +++ b/madspace/tests/test_s_hat_power.py @@ -0,0 +1,148 @@ +"""The power of the s_hat sampling of PhaseSpaceMapping, separate from the +invariant_power of the propagators, and the options of a color-ordered +PhaseSpaceMapping. Unset, the s_hat power is invariant_power, except for a +color-ordered chain with a lower bound on s_hat, where it is AmpliCol's 1.2.""" + +import math + +import pytest +import torch +from madspace_testing import generator, rand, randint + +import madspace as ms + +MASSES = [0.0] * 6 +CM_ENERGY = 13000.0 +ORDER = [0, 2, 3, 1, 4, 5] +MODES = [ + (ms.PhaseSpaceMapping.rambo, None), + (ms.PhaseSpaceMapping.propagator, None), + (ms.PhaseSpaceMapping.color_ordered, ms.ColorOrderedMapping.Options(ORDER)), +] +IDS = ["rambo", "propagator", "color_ordered"] + + +def run(mode, color_ordered, seed, n, cuts=None, **kwargs): + mapping = ms.PhaseSpaceMapping( + MASSES, + CM_ENERGY, + mode=mode, + color_ordered=color_ordered, + invariant_power=0.8, + cuts=cuts, + **kwargs, + ) + g = generator(seed) + inputs = [rand((n, mapping.random_dim()), generator=g)] + if mapping.discrete_dim(): + inputs.append( + randint(0, 2, (n, mapping.discrete_dim()), generator=g, dtype=torch.int32) + ) + p, x1, x2, det = mapping.map_forward(inputs) + return p, x1 * x2, det * 2.0 ** mapping.discrete_dim() + + +@pytest.mark.parametrize("mode, color_ordered", MODES, ids=IDS) +def test_unset_is_invariant_power(mode, color_ordered): + """Without s_hat_power and without a lower bound on s_hat (massless, no + cuts), s_hat is sampled with invariant_power, in every mode.""" + p0, tau0, det0 = run(mode, color_ordered, 1, 1000) + p1, tau1, det1 = run(mode, color_ordered, 1, 1000, s_hat_power=0.8) + assert torch.equal(p0, p1) and torch.equal(tau0, tau1) and torch.equal(det0, det1) + + +@pytest.mark.parametrize("mode, color_ordered", MODES, ids=IDS) +def test_only_s_hat_changes(mode, color_ordered): + """A different s_hat power moves the points in s_hat, not the integral.""" + n = 400_000 + _, tau_a, det_a = run(mode, color_ordered, 2, n, s_hat_power=0.8) + _, tau_b, det_b = run(mode, color_ordered, 3, n, s_hat_power=1.2) + # more points at small s_hat + assert torch.median(tau_b) < 0.5 * torch.median(tau_a) + # the massless volume (~ s_hat^2) convolved with the luminosity, weighted + # by 1/s_hat^2: finite variance for both powers down to s_hat = 0 + w_a = det_a / (tau_a * CM_ENERGY**2) ** 2 + w_b = det_b / (tau_b * CM_ENERGY**2) ** 2 + err = math.hypot( + w_a.std(correction=0).item() / math.sqrt(n), + w_b.std(correction=0).item() / math.sqrt(n), + ) + assert abs(w_a.mean() - w_b.mean()).item() < 4.0 * err + + +@pytest.mark.parametrize( + "masses", + [[0.0, 0.0, 173.0, 173.0, 0.0, 0.0], [0.0] * 6], + ids=["top masses", "jet cuts"], +) +def test_unset_color_ordered_bounded(masses): + """With masses or cuts bounding s_hat from below, a color-ordered chain + samples s_hat as 1/s_hat^1.2, like AmpliCol, unless told otherwise.""" + cuts = None if masses[2] > 0.0 else jet_cuts() + options = ms.ColorOrderedMapping.Options(ORDER) + tau = {} + for power in (None, 1.2, 0.8): + kwargs = {} if power is None else dict(s_hat_power=power) + mapping = ms.PhaseSpaceMapping( + masses, + CM_ENERGY, + mode=ms.PhaseSpaceMapping.color_ordered, + color_ordered=options, + cuts=cuts, + **kwargs, + ) + g = generator(4) + r = rand((1000, mapping.random_dim()), generator=g) + d = randint( + 0, 2, (1000, mapping.discrete_dim()), generator=g, dtype=torch.int32 + ) + tau[power] = mapping.map_forward([r, d])[1] + assert torch.equal(tau[None], tau[1.2]) + assert not torch.equal(tau[None], tau[0.8]) + + +def jet_cuts(): + pids = [21] * 6 + obs = ms.Observable + return ms.Cuts( + [ + ms.CutItem(obs(pids, obs.obs_pt, [obs.jet_pids]), min=20.0), + ms.CutItem(obs(pids, obs.obs_delta_r, [obs.jet_pids]), min=0.4), + ] + ) + + +def test_color_ordered_options(): + """The chain power and the branch choice of a color-ordered + PhaseSpaceMapping can be set; with cuts the defaults are AmpliCol's (chain + power 1, cut-aware branch).""" + options = ms.ColorOrderedMapping.Options + mode = ms.PhaseSpaceMapping.color_ordered + + def momenta(color_ordered): + return run(mode, color_ordered, 5, 1000, cuts=jet_cuts())[0] + + default = momenta(options(ORDER)) + explicit = momenta(options(ORDER, chain_invariant_power=1.0, cut_aware_branch=True)) + assert torch.equal(default, explicit) + assert not torch.equal(default, momenta(options(ORDER, chain_invariant_power=0.7))) + # the two solutions of a side of two particles mirror each other and pass + # the cuts together, so the branch choice shows in a chain of four + chain = [0, 2, 3, 4, 5, 1] + assert not torch.equal( + momenta(options(chain)), momenta(options(chain, cut_aware_branch=False)) + ) + + +def test_color_ordered_options_match_mode(): + """The options are required by the color-ordered mode and refused by the + others.""" + with pytest.raises(ValueError): + ms.PhaseSpaceMapping(MASSES, CM_ENERGY, mode=ms.PhaseSpaceMapping.color_ordered) + with pytest.raises(ValueError): + ms.PhaseSpaceMapping( + MASSES, + CM_ENERGY, + mode=ms.PhaseSpaceMapping.rambo, + color_ordered=ms.ColorOrderedMapping.Options(ORDER), + ) diff --git a/madspace/tests/test_t_channel.py b/madspace/tests/test_t_channel.py index 6e6d18e32c..5ad050e1e9 100644 --- a/madspace/tests/test_t_channel.py +++ b/madspace/tests/test_t_channel.py @@ -19,7 +19,7 @@ def _color_order(n_out, mode): """Single-chain color order for color_ordered mode; None otherwise.""" if mode != ms.PhaseSpaceMapping.color_ordered: return None - return [0] + [i + 2 for i in range(n_out)] + [1] + return ms.ColorOrderedMapping.Options([0] + [i + 2 for i in range(n_out)] + [1]) def _fwd_inputs(mapping, rng, n): @@ -95,7 +95,7 @@ def test_t_channel_masses(masses, rng, mode): masses, CM_ENERGY, mode=mode, - color_order=_color_order(len(masses) - 2, mode), + color_ordered=_color_order(len(masses) - 2, mode), ) inputs, r, _ = _fwd_inputs(mapping, rng, BATCH_SIZE) p_ext, x1, x2, det = mapping.map_forward(inputs) @@ -128,7 +128,7 @@ def test_t_channel_incoming(masses, rng, mode): masses, CM_ENERGY, mode=mode, - color_order=_color_order(len(masses) - 2, mode), + color_ordered=_color_order(len(masses) - 2, mode), ) inputs, r, _ = _fwd_inputs(mapping, rng, BATCH_SIZE) p_ext, x1, x2, det = mapping.map_forward(inputs) @@ -163,7 +163,7 @@ def test_t_channel_momentum_conservation(masses, rng, mode): masses, CM_ENERGY, mode=mode, - color_order=_color_order(len(masses) - 2, mode), + color_ordered=_color_order(len(masses) - 2, mode), ) inputs, r, _ = _fwd_inputs(mapping, rng, BATCH_SIZE) p_ext, x1, x2, det = mapping.map_forward(inputs) @@ -185,7 +185,7 @@ def test_t_channel_inverse(masses, rng, mode): CM_ENERGY, mode=mode, invariant_power=0.3, - color_order=_color_order(len(masses) - 2, mode), + color_ordered=_color_order(len(masses) - 2, mode), ) inputs, r, disc = _fwd_inputs(mapping, rng, BATCH_SIZE) p_ext, x1, x2, det = mapping.map_forward(inputs) @@ -228,7 +228,7 @@ def test_t_channel_phase_space_volume(particle_count, energy, rng, mode): energy, mode=mode, leptonic=False, - color_order=co, + color_ordered=co, ) else: mapping = ms.PhaseSpaceMapping( @@ -236,7 +236,7 @@ def test_t_channel_phase_space_volume(particle_count, energy, rng, mode): energy, mode=mode, leptonic=True, - color_order=co, + color_ordered=co, ) sample_count = 100000 inputs, r, _ = _fwd_inputs(mapping, rng, sample_count) @@ -290,7 +290,7 @@ def test_color_ordered_color_orders(rng, color_order, out_masses, _label): CM_ENERGY, mode=ms.PhaseSpaceMapping.color_ordered, invariant_power=0.3, - color_order=color_order, + color_ordered=ms.ColorOrderedMapping.Options(color_order), ) # random_dim is color-order invariant; discrete_dim tracks the peel count. assert mapping.random_dim() == 3 * len(out_masses) - 2 diff --git a/madspace/tests/test_two_to_three_arcsine.py b/madspace/tests/test_two_to_three_arcsine.py index 0ee6b8f287..f40a651338 100644 --- a/madspace/tests/test_two_to_three_arcsine.py +++ b/madspace/tests/test_two_to_three_arcsine.py @@ -10,7 +10,7 @@ - with flat s23 sampling the block weight becomes exactly independent of r_s23, - without the map the edge peak is still there, - forward and inverse agree, with and without cuts, -- ColorOrderedMapping integrates to the same volume either way, with cuts. +- ColorOrderedMapping, which always uses the map, integrates to the volume. """ import math @@ -128,20 +128,15 @@ def test_round_trip_at_the_edges(s_power): assert torch.max(torch.abs(det * det_inv - 1.0)) < 1e-9 -CUTS = dict( - pt_min=[20.0] * 4, - m_inv_min=[[0.0 if i == j else 10.0 for j in range(4)] for i in range(4)], - dr_min=[[0.0 if i == j else 0.4 for j in range(4)] for i in range(4)], -) - - -def sample(color_order, arcsine, seed, n, cuts=None): +def sample(color_order, seed, n, cuts=None): """Weights of ColorOrderedMapping for the massless n-body phase space at 1 TeV, zero outside the cuts if given.""" cuts = cuts or {} cm_energy = 1000.0 n_out = len(color_order) - 2 - mapping = ms.ColorOrderedMapping(color_order, arcsine_s23=arcsine, **cuts) + mapping = ms.ColorOrderedMapping( + ms.ColorOrderedMapping.Options(color_order), **cuts + ) g = generator(seed) r = rand((n, mapping.random_dim()), generator=g) d = randint(0, 2, (n, mapping.discrete_dim()), generator=g, dtype=torch.int32) @@ -179,21 +174,6 @@ def test_color_ordered_massless_volume(color_order): """The analytic massless 4-body volume, (pi/2)^3 s^2 / 12, with the map.""" volume = (math.pi / 2) ** 3 * 1000.0**4 / 12 means = [ - sample(color_order, True, seed, 1_000_000).mean().item() - for seed in range(200, 204) + sample(color_order, seed, 1_000_000).mean().item() for seed in range(200, 204) ] assert sum(means) / len(means) / volume == pytest.approx(1.0, abs=0.02) - - -def test_color_ordered_cut_volume_unchanged(): - """With cuts the sampled s23 range is a part of the kinematic one and the - map only covers that part. The fiducial volume must not change.""" - order = [0, 2, 3, 4, 5, 1] - n = 1_000_000 - on = torch.cat([sample(order, True, s, n, CUTS) for s in range(300, 304)]) - off = torch.cat([sample(order, False, s, n, CUTS) for s in range(310, 314)]) - err = math.hypot( - on.std(correction=0).item() / math.sqrt(len(on)), - off.std(correction=0).item() / math.sqrt(len(off)), - ) - assert abs(on.mean() - off.mean()).item() < 4.0 * err diff --git a/madspace/tests/test_two_to_three_soft_limit.py b/madspace/tests/test_two_to_three_soft_limit.py index 282e3b65e0..b7fbd68ee2 100644 --- a/madspace/tests/test_two_to_three_soft_limit.py +++ b/madspace/tests/test_two_to_three_soft_limit.py @@ -107,7 +107,7 @@ def test_color_ordered_massless_volume(color_order): * cm_energy ** (2 * (n_out - 2)) / (math.factorial(n_out - 1) * math.factorial(n_out - 2)) ) - mapping = ms.ColorOrderedMapping(color_order) + mapping = ms.ColorOrderedMapping(ms.ColorOrderedMapping.Options(color_order)) n = 2_000_000 means = [] for seed in range(100, 108):