madspace 0.3.1__cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

madspace/include/madspace/phasespace/channel_weights.h

@@ -0,0 +1,51 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/topology.h"
+#include "madspace/util.h"
+
+namespace madspace {
+
+class PropagatorChannelWeights : public FunctionGenerator {
+public:
+    PropagatorChannelWeights(
+        const std::vector<Topology>& topologies,
+        const nested_vector3<std::size_t>& permutations,
+        const nested_vector2<std::size_t>& channel_indices
+    );
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    nested_vector2<double> _momentum_factors;
+    nested_vector2<me_int_t> _invariant_indices;
+    nested_vector2<double> _masses;
+    nested_vector2<double> _widths;
+};
+
+class SubchannelWeights : public FunctionGenerator {
+public:
+    SubchannelWeights(
+        const nested_vector2<Topology>& topologies,
+        const nested_vector3<std::size_t>& permutations,
+        const nested_vector2<std::size_t>& channel_indices
+    );
+
+    std::size_t channel_count() const { return _channel_indices.size(); }
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    nested_vector2<double> _momentum_factors;
+    nested_vector2<double> _masses;
+    nested_vector2<double> _widths;
+    nested_vector2<me_int_t> _invariant_indices;
+    nested_vector2<me_int_t> _on_shell;
+    std::vector<me_int_t> _group_sizes;
+    std::vector<me_int_t> _channel_indices;
+    std::vector<me_int_t> _subchannel_indices;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/chili.h

@@ -0,0 +1,32 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+
+#include <vector>
+
+namespace madspace {
+
+class ChiliMapping : public Mapping {
+public:
+    ChiliMapping(
+        std::size_t n_particles,
+        const std::vector<double>& y_max,
+        const std::vector<double>& pt_min
+    );
+
+    std::size_t random_dim() const { return 3 * _n_particles - 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;
+
+    std::size_t _n_particles;
+    std::vector<double> _y_max;
+    std::vector<double> _pt_min;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/cross_section.h

@@ -0,0 +1,47 @@
+#pragma once
+
+#include "madspace/phasespace/matrix_element.h"
+#include "madspace/phasespace/pdf.h"
+#include "madspace/phasespace/scale.h"
+
+namespace madspace {
+
+class DifferentialCrossSection : public FunctionGenerator {
+public:
+    DifferentialCrossSection(
+        const MatrixElement& matrix_element,
+        double cm_energy,
+        const RunningCoupling& running_coupling,
+        const EnergyScale& energy_scale,
+        const nested_vector2<me_int_t>& pid_options = {},
+        bool has_pdf1 = false,
+        bool has_pdf2 = false,
+        const std::optional<PdfGrid>& pdf_grid1 = std::nullopt,
+        const std::optional<PdfGrid>& pdf_grid2 = std::nullopt,
+        bool has_mirror = false,
+        bool input_momentum_fraction = true
+    );
+
+    const nested_vector2<me_int_t>& pid_options() const { return _pid_options; }
+    bool has_mirror() const { return _has_mirror; }
+    bool has_pdf(std::size_t pdf_index) const {
+        return _pdfs.at(pdf_index) != std::nullopt;
+    }
+    const MatrixElement& matrix_element() const { return _matrix_element; }
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    nested_vector2<me_int_t> _pid_options;
+    MatrixElement _matrix_element;
+    std::array<std::optional<PartonDensity>, 2> _pdfs;
+    std::array<std::vector<me_int_t>, 2> _pdf_indices;
+    RunningCoupling _running_coupling;
+    double _e_cm;
+    EnergyScale _energy_scale;
+    bool _has_mirror;
+    bool _input_momentum_fraction;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/cuts.h

@@ -0,0 +1,34 @@
+#pragma once
+
+#include "madspace/madcode.h"
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/observable.h"
+
+#include <vector>
+
+namespace madspace {
+
+class Cuts : public FunctionGenerator {
+public:
+    enum CutMode { any, all };
+    struct CutItem {
+        Observable observable;
+        double min = -std::numeric_limits<double>::infinity();
+        double max = std::numeric_limits<double>::infinity();
+        CutMode mode = CutMode::all;
+    };
+
+    Cuts(const std::vector<CutItem>& cut_data);
+    Cuts(std::size_t particle_count);
+    double sqrt_s_min() const;
+    std::vector<double> eta_max() const;
+    std::vector<double> pt_min() const;
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::vector<CutItem> _cut_data;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/discrete_flow.h

@@ -0,0 +1,44 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/mlp.h"
+
+namespace madspace {
+
+class DiscreteFlow : public Mapping {
+public:
+    DiscreteFlow(
+        const std::vector<std::size_t>& option_counts,
+        const std::string& prefix = "",
+        const std::vector<std::size_t>& dims_with_prior = {},
+        std::size_t condition_dim = 0,
+        std::size_t subnet_hidden_dim = 32,
+        std::size_t subnet_layers = 3,
+        MLP::Activation subnet_activation = MLP::leaky_relu
+    );
+    const std::vector<std::size_t>& option_counts() const { return _option_counts; }
+    std::size_t condition_dim() const { return _condition_dim; }
+    void initialize_globals(ContextPtr context) const;
+
+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;
+    Result build_transform(
+        FunctionBuilder& fb,
+        const ValueVec& inputs,
+        const ValueVec& conditions,
+        bool inverse
+    ) const;
+
+    std::vector<std::size_t> _option_counts;
+    std::size_t _condition_dim;
+    std::optional<std::string> _first_prob_name;
+    std::vector<MLP> _subnets;
+    std::vector<bool> _dim_has_prior;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/discrete_sampler.h

@@ -0,0 +1,53 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/runtime/context.h"
+
+namespace madspace {
+
+class DiscreteHistogram : public FunctionGenerator {
+public:
+    DiscreteHistogram(const std::vector<std::size_t>& option_counts);
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::vector<std::size_t> _option_counts;
+};
+
+class DiscreteSampler : public Mapping {
+public:
+    DiscreteSampler(
+        const std::vector<std::size_t>& option_counts,
+        const std::string& prefix = "",
+        const std::vector<std::size_t>& dims_with_prior = {}
+    );
+    const std::vector<std::size_t>& option_counts() const { return _option_counts; }
+    const std::vector<std::string>& prob_names() const { return _prob_names; }
+    void initialize_globals(ContextPtr context) const;
+
+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;
+    Result build_transform(
+        FunctionBuilder& fb,
+        const ValueVec& inputs,
+        const ValueVec& conditions,
+        bool inverse
+    ) const;
+
+    std::vector<std::size_t> _option_counts;
+    std::vector<bool> _dim_has_prior;
+    std::vector<std::string> _prob_names;
+};
+
+void initialize_uniform_probs(
+    ContextPtr context, const std::string& name, std::size_t option_count
+);
+
+} // namespace madspace

madspace/include/madspace/phasespace/flow.h

@@ -0,0 +1,53 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/mlp.h"
+
+namespace madspace {
+
+class Flow : public Mapping {
+public:
+    Flow(
+        std::size_t input_dim,
+        std::size_t condition_dim = 0,
+        const std::string& prefix = "",
+        std::size_t bin_count = 10,
+        std::size_t subnet_hidden_dim = 32,
+        std::size_t subnet_layers = 3,
+        MLP::Activation subnet_activation = MLP::leaky_relu,
+        bool invert_spline = true
+    );
+    std::size_t input_dim() const { return _input_dim; }
+    std::size_t condition_dim() const { return _condition_dim; }
+    void initialize_globals(ContextPtr context) const;
+    void initialize_from_vegas(ContextPtr context, const std::string& grid_name) const;
+
+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;
+    Result build_transform(
+        FunctionBuilder& fb,
+        const ValueVec& inputs,
+        const ValueVec& conditions,
+        bool inverse
+    ) const;
+
+    struct CouplingBlock {
+        MLP subnet1;
+        MLP subnet2;
+        std::vector<me_int_t> indices1;
+        std::vector<me_int_t> indices2;
+    };
+
+    std::vector<CouplingBlock> _coupling_blocks;
+    std::size_t _input_dim;
+    std::size_t _condition_dim;
+    std::size_t _bin_count;
+    bool _invert_spline;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/histograms.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/observable.h"
+
+namespace madspace {
+
+class ObservableHistograms : public FunctionGenerator {
+public:
+    struct HistItem {
+        Observable observable;
+        double min;
+        double max;
+        std::size_t bin_count;
+    };
+    ObservableHistograms(const std::vector<HistItem>& observables);
+    const std::vector<HistItem>& observables() const { return _observables; }
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::vector<HistItem> _observables;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/integrand.h

@@ -0,0 +1,204 @@
+#pragma once
+
+#include "madspace/phasespace/channel_weight_network.h"
+#include "madspace/phasespace/channel_weights.h"
+#include "madspace/phasespace/cross_section.h"
+#include "madspace/phasespace/discrete_flow.h"
+#include "madspace/phasespace/discrete_sampler.h"
+#include "madspace/phasespace/flow.h"
+#include "madspace/phasespace/matrix_element.h"
+#include "madspace/phasespace/pdf.h"
+#include "madspace/phasespace/phasespace.h"
+#include "madspace/phasespace/vegas.h"
+#include "madspace/util.h"
+
+namespace madspace {
+
+class Unweighter : public FunctionGenerator {
+public:
+    Unweighter(const TypeVec& types);
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+};
+
+class Integrand : public FunctionGenerator {
+public:
+    using AdaptiveMapping = std::variant<std::monostate, VegasMapping, Flow>;
+    using AdaptiveDiscrete =
+        std::variant<std::monostate, DiscreteSampler, DiscreteFlow>;
+    inline static const int sample = 1;
+    inline static const int unweight = 2;
+    inline static const int return_momenta = 4;
+    inline static const int return_x1_x2 = 8;
+    inline static const int return_indices = 16;
+    inline static const int return_random = 32;
+    inline static const int return_latent = 64;
+    inline static const int return_channel = 128;
+    inline static const int return_chan_weights = 256;
+    inline static const int return_cwnet_input = 512;
+    inline static const int return_discrete = 1024;
+    inline static const int return_discrete_latent = 2048;
+
+    inline static const std::vector<MatrixElement::MatrixElementInput>
+        matrix_element_inputs = {
+            MatrixElement::momenta_in,
+            MatrixElement::alpha_s_in,
+            MatrixElement::flavor_in,
+            MatrixElement::random_color_in,
+            MatrixElement::random_helicity_in,
+            MatrixElement::random_diagram_in,
+        };
+    inline static const std::vector<MatrixElement::MatrixElementOutput>
+        matrix_element_outputs = {
+            MatrixElement::matrix_element_out,
+            MatrixElement::diagram_amp2_out,
+            MatrixElement::color_index_out,
+            MatrixElement::helicity_index_out,
+            MatrixElement::diagram_index_out,
+        };
+
+    Integrand(
+        const PhaseSpaceMapping& mapping,
+        const DifferentialCrossSection& diff_xs,
+        const AdaptiveMapping& adaptive_map = std::monostate{},
+        const AdaptiveDiscrete& discrete_before = std::monostate{},
+        const AdaptiveDiscrete& discrete_after = std::monostate{},
+        const std::optional<PdfGrid>& pdf_grid = std::nullopt,
+        const std::optional<EnergyScale>& energy_scale = std::nullopt,
+        const std::optional<PropagatorChannelWeights>& prop_chan_weights = std::nullopt,
+        const std::optional<SubchannelWeights>& subchan_weights = std::nullopt,
+        const std::optional<ChannelWeightNetwork>& chan_weight_net = std::nullopt,
+        const std::vector<me_int_t>& chan_weight_remap = {},
+        std::size_t remapped_chan_count = 0,
+        int flags = 0,
+        const std::vector<std::size_t>& channel_indices = {},
+        const std::vector<std::size_t>& active_flavors = {}
+    );
+    std::size_t particle_count() const { return _mapping.particle_count(); }
+    int flags() const { return _flags; }
+    std::optional<std::string> vegas_grid_name() const {
+        if (auto vegas = std::get_if<VegasMapping>(&_adaptive_map)) {
+            return vegas->grid_name();
+        } else {
+            return std::nullopt;
+        }
+    }
+    std::size_t vegas_dimension() const {
+        if (auto vegas = std::get_if<VegasMapping>(&_adaptive_map)) {
+            return vegas->dimension();
+        } else {
+            return 0;
+        }
+    }
+    std::size_t vegas_bin_count() const {
+        if (auto vegas = std::get_if<VegasMapping>(&_adaptive_map)) {
+            return vegas->bin_count();
+        } else {
+            return 0;
+        }
+    }
+    const PhaseSpaceMapping& mapping() const { return _mapping; }
+    const DifferentialCrossSection& diff_xs() const { return _diff_xs; }
+    const AdaptiveMapping& adaptive_map() const { return _adaptive_map; }
+    const AdaptiveDiscrete& discrete_before() const { return _discrete_before; }
+    const AdaptiveDiscrete& discrete_after() const { return _discrete_after; }
+    const std::optional<EnergyScale>& energy_scale() const { return _energy_scale; }
+    const std::optional<PropagatorChannelWeights>& prop_chan_weights() const {
+        return _prop_chan_weights;
+    }
+    const std::optional<ChannelWeightNetwork>& chan_weight_net() const {
+        return _chan_weight_net;
+    }
+    const std::size_t random_dim() const { return _random_dim; }
+    std::tuple<std::vector<std::size_t>, std::vector<bool>> latent_dims() const;
+
+private:
+    struct ChannelArgs {
+        Value r, batch_size;
+        bool has_permutations, has_multi_flavor, has_mirror, has_pdf_prior;
+        Value max_weight;
+    };
+    struct ChannelResult {
+        std::array<Value, 21> values;
+
+        Value& r() { return values[0]; }
+        Value& latent() { return values[1]; }
+        Value& momenta() { return values[2]; }
+        Value& momenta_mirror() { return values[3]; }
+        Value& momenta_acc() { return values[4]; }
+        Value& x(std::size_t pdf_index) { return values[5 + pdf_index]; }
+        Value& x_acc(std::size_t pdf_index) { return values[7 + pdf_index]; }
+        Value& pdf_prior() { return values[9]; }
+        Value& chan_index() { return values[10]; }
+        Value& chan_index_in_group() { return values[11]; }
+        Value& flavor_id() { return values[12]; }
+        Value& mirror_id() { return values[13]; }
+        Value& indices_acc() { return values[14]; }
+        Value& weight_before_cuts() { return values[15]; }
+        Value& weight_after_cuts() { return values[16]; }
+        Value& adaptive_prob() { return values[17]; }
+        Value& pdf_cache(std::size_t pdf_index) { return values[18 + pdf_index]; }
+        Value& scale_cache() { return values[20]; }
+    };
+
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+    ChannelResult
+    build_channel_part(FunctionBuilder& fb, const ChannelArgs& args) const;
+    ValueVec build_common_part(
+        FunctionBuilder& fb, const ChannelArgs& args, ChannelResult& result
+    ) const;
+
+    PhaseSpaceMapping _mapping;
+    DifferentialCrossSection _diff_xs;
+    AdaptiveMapping _adaptive_map;
+    AdaptiveDiscrete _discrete_before;
+    AdaptiveDiscrete _discrete_after;
+    std::array<std::optional<PartonDensity>, 2> _pdfs;
+    std::array<std::vector<me_int_t>, 2> _pdf_indices;
+    std::optional<EnergyScale> _energy_scale;
+    std::optional<PropagatorChannelWeights> _prop_chan_weights;
+    std::optional<SubchannelWeights> _subchan_weights;
+    std::optional<ChannelWeightNetwork> _chan_weight_net;
+    std::vector<me_int_t> _chan_weight_remap;
+    me_int_t _remapped_chan_count;
+    int _flags;
+    std::vector<me_int_t> _channel_indices;
+    me_int_t _random_dim;
+    std::size_t _latent_dim;
+    std::vector<double> _active_flavors;
+
+    friend class IntegrandProbability;
+    friend class MultiChannelIntegrand;
+};
+
+class MultiChannelIntegrand : public FunctionGenerator {
+public:
+    MultiChannelIntegrand(const std::vector<std::shared_ptr<Integrand>>& integrands);
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::vector<std::shared_ptr<Integrand>> _integrands;
+};
+
+class IntegrandProbability : public FunctionGenerator {
+public:
+    IntegrandProbability(const Integrand& integrand);
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    Integrand::AdaptiveMapping _adaptive_map;
+    Integrand::AdaptiveDiscrete _discrete_before;
+    Integrand::AdaptiveDiscrete _discrete_after;
+    std::size_t _permutation_count;
+    std::size_t _flavor_count;
+    bool _has_pdf_prior;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/invariants.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+
+namespace madspace {
+
+class Invariant : public Mapping {
+public:
+    Invariant(double power = 0, double mass = 0, double width = 0) :
+        Mapping("Invariant", {batch_float}, {batch_float}, {batch_float, batch_float}),
+        _power(power),
+        _mass(mass),
+        _width(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;
+
+    double _power, _mass, _width;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/luminosity.h

@@ -0,0 +1,41 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/invariants.h"
+
+namespace madspace {
+
+class Luminosity : public Mapping {
+public:
+    Luminosity(
+        double s_lab,
+        double s_hat_min,
+        double s_hat_max = 0,
+        double invariant_power = 1,
+        double mass = 0,
+        double width = 0
+    ) :
+        Mapping(
+            "Luminosity",
+            {batch_float, batch_float},
+            {batch_float, batch_float, batch_float},
+            {}
+        ),
+        _s_lab(s_lab),
+        _s_hat_min(s_hat_min),
+        _s_hat_max(s_hat_max == 0 ? s_lab : s_hat_max),
+        _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;
+
+    double _s_lab, _s_hat_min, _s_hat_max;
+    Invariant _invariant;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/matrix_element.h

@@ -0,0 +1,70 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/runtime/context.h"
+
+namespace madspace {
+
+class MatrixElement : public FunctionGenerator {
+public:
+    enum MatrixElementInput {
+        momenta_in,
+        alpha_s_in,
+        flavor_in,
+        random_color_in,
+        random_helicity_in,
+        random_diagram_in,
+        helicity_in,
+        diagram_in
+    };
+
+    enum MatrixElementOutput {
+        matrix_element_out,
+        diagram_amp2_out,
+        color_index_out,
+        helicity_index_out,
+        diagram_index_out
+    };
+
+    MatrixElement(
+        std::size_t matrix_element_index,
+        std::size_t particle_count,
+        const std::vector<MatrixElementInput>& inputs = {momenta_in},
+        const std::vector<MatrixElementOutput>& outputs = {matrix_element_out},
+        std::size_t diagram_count = 1,
+        bool sample_random_inputs = false
+    );
+    MatrixElement(
+        const MatrixElementApi& matrix_element_api,
+        const std::vector<MatrixElementInput>& inputs = {momenta_in},
+        const std::vector<MatrixElementOutput>& outputs = {matrix_element_out},
+        bool sample_random_inputs = false
+    ) :
+        MatrixElement(
+            matrix_element_api.index(),
+            matrix_element_api.particle_count(),
+            inputs,
+            outputs,
+            matrix_element_api.diagram_count(),
+            sample_random_inputs
+        ) {};
+    std::size_t matrix_element_index() const { return _matrix_element_index; }
+    std::size_t diagram_count() const { return _diagram_count; }
+    std::size_t particle_count() const { return _particle_count; }
+    const std::vector<MatrixElementInput>& inputs() const { return _inputs; }
+    const std::vector<MatrixElementOutput>& outputs() const { return _outputs; }
+    std::vector<MatrixElementInput> external_inputs() const;
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::size_t _matrix_element_index;
+    std::size_t _particle_count;
+    std::size_t _diagram_count;
+    std::vector<MatrixElementInput> _inputs;
+    std::vector<MatrixElementOutput> _outputs;
+    bool _sample_random_inputs;
+};
+
+} // namespace madspace