madspace 0.3.1__cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

madspace/include/madspace/phasespace/mlp.h

@@ -0,0 +1,37 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/runtime/context.h"
+
+namespace madspace {
+
+class MLP : public FunctionGenerator {
+public:
+    enum Activation { relu, leaky_relu, elu, gelu, sigmoid, softplus, linear };
+    MLP(std::size_t input_dim,
+        std::size_t output_dim,
+        std::size_t hidden_dim = 32,
+        std::size_t layers = 3,
+        Activation activation = leaky_relu,
+        const std::string& prefix = "");
+
+    std::size_t input_dim() const { return _input_dim; }
+    std::size_t output_dim() const { return _output_dim; }
+    void initialize_globals(ContextPtr context) const;
+    std::string last_layer_bias_name() const {
+        return prefixed_name(_prefix, std::format("layer{}.bias", _layers));
+    }
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::size_t _input_dim;
+    std::size_t _output_dim;
+    std::size_t _hidden_dim;
+    std::size_t _layers;
+    Activation _activation;
+    std::string _prefix;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/multichannel.h

@@ -0,0 +1,49 @@
+#pragma once
+
+#include <format>
+#include <vector>
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/phasespace.h"
+
+namespace madspace {
+
+class MultiChannelMapping : public Mapping {
+public:
+    MultiChannelMapping(const std::vector<std::shared_ptr<Mapping>>& mappings);
+
+private:
+    Result build_impl(
+        FunctionBuilder& fb,
+        const ValueVec& inputs,
+        const ValueVec& conditions,
+        bool inverse
+    ) const;
+    Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override {
+        return build_impl(fb, inputs, conditions, false);
+    }
+    Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override {
+        return build_impl(fb, inputs, conditions, true);
+    }
+
+    std::vector<std::shared_ptr<Mapping>> _mappings;
+};
+
+class MultiChannelFunction : public FunctionGenerator {
+public:
+    MultiChannelFunction(
+        const std::vector<std::shared_ptr<FunctionGenerator>>& functions
+    );
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::vector<std::shared_ptr<FunctionGenerator>> _functions;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/observable.h

@@ -0,0 +1,85 @@
+#pragma once
+
+#include "madspace/madcode.h"
+#include "madspace/phasespace/base.h"
+#include "madspace/util.h"
+
+#include <vector>
+
+namespace madspace {
+
+class Observable : public FunctionGenerator {
+public:
+    static const std::vector<int> jet_pids;
+    static const std::vector<int> bottom_pids;
+    static const std::vector<int> lepton_pids;
+    static const std::vector<int> missing_pids;
+    static const std::vector<int> photon_pids;
+
+    enum ObservableOption {
+        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
+    };
+
+    Observable(
+        const std::vector<int>& pids,
+        ObservableOption observable,
+        const nested_vector2<int>& select_pids,
+        bool sum_momenta = false,
+        bool sum_observable = false,
+        const std::optional<ObservableOption>& order_observable = std::nullopt,
+        const std::vector<int>& order_indices = {},
+        bool ignore_incoming = true,
+        const std::string& name = ""
+    );
+    ObservableOption observable() const { return _observable; }
+    std::vector<std::size_t> simple_observable_indices() const {
+        if (_sum_momenta || _sum_observable || _indices.size() != 1) {
+            return {};
+        } else {
+            return {_indices.at(0).begin(), _indices.at(0).end()};
+        }
+    }
+    std::string name() const { return _name; }
+    bool not_found() const;
+
+private:
+    Observable(
+        std::tuple<nested_vector2<me_int_t>, nested_vector2<me_int_t>, Type>
+            indices_and_type,
+        const std::vector<int>& pids,
+        ObservableOption observable,
+        bool sum_momenta,
+        bool sum_observable,
+        const std::optional<ObservableOption>& order_observable,
+        bool ignore_incoming,
+        const std::string& name
+    );
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    ObservableOption _observable;
+    nested_vector2<me_int_t> _indices;
+    std::optional<ObservableOption> _order_observable;
+    nested_vector2<me_int_t> _order_indices;
+    bool _sum_momenta;
+    bool _sum_observable;
+    std::string _name;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/pdf.h

@@ -0,0 +1,78 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/runtime/context.h"
+
+namespace madspace {
+
+struct PdfGrid {
+    std::vector<double> x;
+    std::vector<double> logx;
+    std::vector<double> q;
+    std::vector<double> logq2;
+    std::vector<int> pids;
+    std::vector<std::vector<double>> values;
+    std::vector<std::size_t> region_sizes;
+
+    PdfGrid(const std::string& file);
+    std::size_t grid_point_count() const;
+    std::size_t q_count() const;
+    void initialize_coefficients(Tensor tensor) const;
+    void initialize_logx(Tensor tensor) const;
+    void initialize_logq2(Tensor tensor) const;
+    std::vector<std::size_t> coefficients_shape(bool batch_dim = false) const;
+    std::vector<std::size_t> logx_shape(bool batch_dim = false) const;
+    std::vector<std::size_t> logq2_shape(bool batch_dim = false) const;
+    void initialize_globals(ContextPtr context, const std::string& prefix = "") const;
+};
+
+class PartonDensity : public FunctionGenerator {
+public:
+    PartonDensity(
+        const PdfGrid& grid,
+        const std::vector<int>& pids,
+        bool dynamic_pid = false,
+        const std::string& prefix = ""
+    );
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::vector<me_int_t> _pid_indices;
+    bool _dynamic_pid;
+    std::string _prefix;
+    std::vector<std::size_t> _logx_shape;
+    std::vector<std::size_t> _logq2_shape;
+    std::vector<std::size_t> _coeffs_shape;
+};
+
+struct AlphaSGrid {
+    std::vector<double> q;
+    std::vector<double> logq2;
+    std::vector<double> values;
+    std::vector<std::size_t> region_sizes;
+
+    AlphaSGrid(const std::string& file);
+    std::size_t q_count() const;
+    void initialize_coefficients(Tensor tensor) const;
+    void initialize_logq2(Tensor tensor) const;
+    std::vector<std::size_t> coefficients_shape(bool batch_dim = false) const;
+    std::vector<std::size_t> logq2_shape(bool batch_dim = false) const;
+    void initialize_globals(ContextPtr context, const std::string& prefix = "") const;
+};
+
+class RunningCoupling : public FunctionGenerator {
+public:
+    RunningCoupling(const AlphaSGrid& grid, const std::string& prefix = "");
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::string _prefix;
+    std::vector<std::size_t> _logq2_shape;
+    std::vector<std::size_t> _coeffs_shape;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/phasespace.h

@@ -0,0 +1,67 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/chili.h"
+#include "madspace/phasespace/cuts.h"
+#include "madspace/phasespace/invariants.h"
+#include "madspace/phasespace/luminosity.h"
+#include "madspace/phasespace/rambo.h"
+#include "madspace/phasespace/t_propagator_mapping.h"
+#include "madspace/phasespace/three_particle.h"
+#include "madspace/phasespace/topology.h"
+
+namespace madspace {
+
+class PhaseSpaceMapping : public Mapping {
+public:
+    enum TChannelMode { propagator, rambo, chili };
+
+    PhaseSpaceMapping(
+        const Topology& topology,
+        double cm_energy,
+        bool leptonic = false,
+        double invariant_power = 0.8,
+        TChannelMode t_channel_mode = propagator,
+        const std::optional<Cuts>& cuts = std::nullopt,
+        const std::vector<std::vector<std::size_t>>& permutations = {}
+    );
+
+    PhaseSpaceMapping(
+        const std::vector<double>& external_masses,
+        double cm_energy,
+        bool leptonic = false,
+        double invariant_power = 0.8,
+        TChannelMode mode = rambo,
+        const std::optional<Cuts>& cuts = std::nullopt
+    );
+
+    std::size_t random_dim() const {
+        return 3 * _topology.outgoing_masses().size() - (_leptonic ? 4 : 2);
+    }
+    std::size_t particle_count() const {
+        return _topology.outgoing_masses().size() + 2;
+    }
+    std::size_t channel_count() const { return _permutations.size(); }
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+
+    Topology _topology;
+    Cuts _cuts;
+    double _pi_factors;
+    double _sqrt_s_lab;
+    bool _leptonic;
+    bool _map_luminosity;
+    std::vector<Invariant> _s_invariants;
+    std::variant<TPropagatorMapping, FastRamboMapping, ChiliMapping, std::monostate>
+        _t_mapping;
+    std::vector<std::variant<TwoBodyDecay, ThreeBodyDecay, FastRamboMapping>> _s_decays;
+    nested_vector2<me_int_t> _permutations;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/rambo.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+
+namespace madspace {
+
+class FastRamboMapping : public Mapping {
+public:
+    FastRamboMapping(std::size_t n_particles, bool massless, bool com = true);
+
+    std::size_t random_dim() const { return 3 * _n_particles - 4; }
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+
+    std::size_t _n_particles;
+    bool _massless;
+    double _com;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/scale.h

@@ -0,0 +1,52 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+
+namespace madspace {
+
+class EnergyScale : public FunctionGenerator {
+public:
+    enum DynamicalScaleType {
+        transverse_energy,
+        transverse_mass,
+        half_transverse_mass,
+        partonic_energy
+    };
+
+    EnergyScale(std::size_t particle_count) :
+        EnergyScale(particle_count, half_transverse_mass, false, false, 0., 0., 0.) {}
+    EnergyScale(std::size_t particle_count, DynamicalScaleType type) :
+        EnergyScale(particle_count, type, false, false, 0., 0., 0.) {}
+    EnergyScale(std::size_t particle_count, double fixed_scale) :
+        EnergyScale(
+            particle_count,
+            half_transverse_mass,
+            true,
+            true,
+            fixed_scale,
+            fixed_scale,
+            fixed_scale
+        ) {}
+    EnergyScale(
+        std::size_t particle_count,
+        DynamicalScaleType dynamical_scale_type,
+        bool ren_scale_fixed,
+        bool fact_scale_fixed,
+        double ren_scale,
+        double fact_scale1,
+        double fact_scale2
+    );
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    DynamicalScaleType _dynamical_scale_type;
+    bool _ren_scale_fixed;
+    bool _fact_scale_fixed;
+    double _ren_scale;
+    double _fact_scale1;
+    double _fact_scale2;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/t_propagator_mapping.h

@@ -0,0 +1,34 @@
+#pragma once
+
+#include <vector>
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/invariants.h"
+#include "madspace/phasespace/topology.h"
+#include "madspace/phasespace/two_particle.h"
+
+namespace madspace {
+
+class TPropagatorMapping : public Mapping {
+public:
+    TPropagatorMapping(
+        const std::vector<std::size_t>& integration_order, double invariant_power = 0.8
+    );
+    std::size_t random_dim() const { return 3 * _integration_order.size() - 1; }
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+
+    std::vector<std::size_t> _integration_order;
+    std::vector<bool> _sample_sides;
+    Invariant _uniform_invariant;
+    TwoToTwoParticleScattering _com_scattering;
+    TwoToTwoParticleScattering _lab_scattering;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/three_particle.h

@@ -0,0 +1,68 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/invariants.h"
+
+namespace madspace {
+
+class ThreeBodyDecay : public Mapping {
+public:
+    ThreeBodyDecay(bool com) :
+        Mapping(
+            "ThreeBodyDecay",
+            [&] {
+                TypeVec input_types(9, batch_float);
+                if (!com) {
+                    input_types.push_back(batch_four_vec);
+                }
+                return input_types;
+            }(),
+            {batch_four_vec, batch_four_vec, batch_four_vec},
+            {}
+        ),
+        _com(com) {}
+    std::size_t random_dim() const { return 5; }
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+
+    bool _com;
+};
+
+class TwoToThreeParticleScattering : public Mapping {
+public:
+    TwoToThreeParticleScattering(
+        double t_invariant_power = 0,
+        double t_mass = 0,
+        double t_width = 0,
+        double s_invariant_power = 0,
+        double s_mass = 0,
+        double s_width = 0
+    ) :
+        Mapping(
+            "TwoToThreeParticleScattering",
+            {batch_float, batch_float, batch_float, batch_float, batch_float},
+            {batch_four_vec, batch_four_vec},
+            {batch_four_vec, batch_four_vec, batch_four_vec}
+        ),
+        _t_invariant(t_invariant_power, t_mass, t_width),
+        _s_invariant(s_invariant_power, s_mass, s_width) {}
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+
+    Invariant _t_invariant;
+    Invariant _s_invariant;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/topology.h

@@ -0,0 +1,116 @@
+#pragma once
+
+#include <array>
+#include <ostream>
+#include <string>
+#include <vector>
+
+namespace madspace {
+
+struct Propagator {
+    double mass;
+    double width;
+    int integration_order;
+    double e_min;
+    double e_max;
+    int pdg_id;
+};
+
+class Diagram {
+public:
+    enum LineType { incoming, outgoing, propagator };
+    class LineRef {
+    public:
+        LineRef(LineType type, std::size_t index) : _type(type), _index(index) {}
+        LineRef(std::string str);
+        LineType type() const { return _type; }
+        std::size_t index() const { return _index; }
+
+    private:
+        LineType _type;
+        std::size_t _index;
+    };
+    using Vertex = std::vector<LineRef>;
+
+    Diagram(
+        const std::vector<double>& incoming_masses,
+        const std::vector<double>& outgoing_masses,
+        const std::vector<Propagator>& propagators,
+        const std::vector<Vertex>& vertices
+    );
+
+    const std::vector<double>& incoming_masses() const { return _incoming_masses; }
+    const std::vector<double>& outgoing_masses() const { return _outgoing_masses; }
+    const std::vector<Propagator>& propagators() const { return _propagators; }
+    const std::vector<Vertex>& vertices() const { return _vertices; }
+    const std::array<int, 2>& incoming_vertices() const { return _incoming_vertices; };
+    const std::vector<int>& outgoing_vertices() const { return _outgoing_vertices; };
+    const std::vector<std::vector<std::size_t>>& propagator_vertices() const {
+        return _propagator_vertices;
+    }
+
+private:
+    std::vector<double> _incoming_masses;
+    std::vector<double> _outgoing_masses;
+    std::vector<Propagator> _propagators;
+    std::vector<Vertex> _vertices;
+    std::array<int, 2> _incoming_vertices;
+    std::vector<int> _outgoing_vertices;
+    std::vector<std::vector<std::size_t>> _propagator_vertices;
+};
+
+std::ostream& operator<<(std::ostream& out, const Diagram::LineRef& value);
+
+class Topology {
+public:
+    struct Decay {
+        std::size_t index;
+        std::size_t parent_index;
+        std::vector<std::size_t> child_indices;
+        double mass;
+        double width;
+        double e_min;
+        double e_max;
+        int pdg_id;
+        bool on_shell;
+    };
+
+    static std::vector<Topology> topologies(const Diagram& diagram);
+    Topology(const Diagram& diagram);
+
+    std::size_t t_propagator_count() const { return _t_integration_order.size(); }
+    const std::vector<std::size_t>& t_integration_order() const {
+        return _t_integration_order;
+    }
+    const std::vector<double>& t_propagator_masses() const {
+        return _t_propagator_masses;
+    }
+    const std::vector<double>& t_propagator_widths() const {
+        return _t_propagator_widths;
+    }
+    const std::vector<Decay>& decays() const { return _decays; }
+    const std::vector<std::size_t>& decay_integration_order() const {
+        return _decay_integration_order;
+    }
+    const std::vector<std::size_t>& outgoing_indices() const {
+        return _outgoing_indices;
+    }
+    const std::vector<double>& incoming_masses() const { return _incoming_masses; }
+    const std::vector<double>& outgoing_masses() const { return _outgoing_masses; }
+    std::vector<std::tuple<std::vector<int>, double, double>>
+    propagator_momentum_terms(bool only_decays = false) const;
+
+private:
+    Topology() = default;
+
+    std::vector<std::size_t> _t_integration_order;
+    std::vector<double> _t_propagator_masses;
+    std::vector<double> _t_propagator_widths;
+    std::vector<Decay> _decays;
+    std::vector<std::size_t> _decay_integration_order;
+    std::vector<std::size_t> _outgoing_indices;
+    std::vector<double> _incoming_masses;
+    std::vector<double> _outgoing_masses;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/two_particle.h

@@ -0,0 +1,63 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/invariants.h"
+
+namespace madspace {
+
+class TwoBodyDecay : public Mapping {
+public:
+    TwoBodyDecay(bool com) :
+        Mapping(
+            "TwoBodyDecay",
+            [&] {
+                TypeVec input_types(5, batch_float);
+                if (!com) {
+                    input_types.push_back(batch_four_vec);
+                }
+                return input_types;
+            }(),
+            {batch_four_vec, batch_four_vec},
+            {}
+        ),
+        _com(com) {}
+    std::size_t random_dim() const { return 2; }
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+
+    bool _com;
+};
+
+class TwoToTwoParticleScattering : public Mapping {
+public:
+    TwoToTwoParticleScattering(
+        bool com, double invariant_power = 0, double mass = 0, double width = 0
+    ) :
+        Mapping(
+            "TwoToTwoParticleScattering",
+            {batch_float, batch_float, batch_float, batch_float},
+            {batch_four_vec, batch_four_vec},
+            {batch_four_vec, batch_four_vec}
+        ),
+        _com(com),
+        _invariant(invariant_power, mass, width) {}
+
+private:
+    Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+    Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const override;
+
+    bool _com;
+    Invariant _invariant;
+};
+
+} // namespace madspace