madspace 0.3.1__cp311-cp311-macosx_14_0_arm64.whl

madspace/include/madspace/madcode/instruction.h

@@ -0,0 +1,208 @@
+#pragma once
+
+#include <initializer_list>
+#include <map>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+#include "type.h"
+
+namespace madspace {
+
+namespace opcodes {
+enum Opcode {
+#include "opcode_mixin.h"
+};
+} // namespace opcodes
+
+class Instruction {
+public:
+    Instruction(const std::string& name, int opcode, bool differentiable) :
+        _name(name), _opcode(opcode), _differentiable(differentiable) {}
+    virtual ~Instruction() = default;
+    virtual TypeVec signature(const ValueVec& args) const = 0;
+    const std::string& name() const { return _name; }
+    int opcode() const { return _opcode; }
+    bool differentiable() const { return _differentiable; }
+
+protected:
+    void check_arg_count(const ValueVec& args, std::size_t count) const;
+    me_int_t int_literal_arg(
+        const ValueVec& args, std::size_t index, bool check_non_negative = true
+    ) const;
+
+private:
+    std::string _name;
+    int _opcode;
+    bool _differentiable;
+};
+
+class ShapeExpr {
+public:
+    ShapeExpr(const char* expr);
+    bool check_and_update(std::map<char, int>& variables, int value) const;
+    std::optional<int> evaluate(const std::map<char, int>& variables) const;
+    char first_var_name() const { return std::get<0>(terms.at(0)); }
+
+private:
+    std::vector<std::tuple<char, int>> terms;
+};
+
+class SimpleInstruction : public Instruction {
+public:
+    using DynShape = std::vector<std::variant<int, ShapeExpr, std::monostate>>;
+    using SigType = std::tuple<DataType, bool, DynShape, bool>;
+
+    SimpleInstruction(
+        std::string name,
+        int opcode,
+        bool differentiable,
+        std::initializer_list<SigType> _inputs,
+        std::initializer_list<SigType> _outputs
+    ) :
+        Instruction(name, opcode, differentiable), inputs(_inputs), outputs(_outputs) {}
+
+    TypeVec signature(const ValueVec& args) const override;
+
+private:
+    const std::vector<SigType> inputs;
+    const std::vector<SigType> outputs;
+};
+
+class StackInstruction : public Instruction {
+public:
+    StackInstruction(int opcode, bool differentiable) :
+        Instruction("stack", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class UnstackInstruction : public Instruction {
+public:
+    UnstackInstruction(int opcode, bool differentiable) :
+        Instruction("unstack", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class UnstackSizesInstruction : public Instruction {
+public:
+    UnstackSizesInstruction(int opcode, bool differentiable) :
+        Instruction("unstack_sizes", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class BatchCatInstruction : public Instruction {
+public:
+    BatchCatInstruction(int opcode, bool differentiable) :
+        Instruction("batch_cat", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class BatchSplitInstruction : public Instruction {
+public:
+    BatchSplitInstruction(int opcode, bool differentiable) :
+        Instruction("batch_split", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class CatInstruction : public Instruction {
+public:
+    CatInstruction(int opcode, bool differentiable) :
+        Instruction("cat", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class BatchSizeInstruction : public Instruction {
+public:
+    BatchSizeInstruction(int opcode, bool differentiable) :
+        Instruction("batch_size", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class OffsetIndicesInstruction : public Instruction {
+public:
+    OffsetIndicesInstruction(int opcode, bool differentiable) :
+        Instruction("offset_indices", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class FullInstruction : public Instruction {
+public:
+    FullInstruction(int opcode, bool differentiable) :
+        Instruction("full", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class SqueezeInstruction : public Instruction {
+public:
+    SqueezeInstruction(int opcode, bool differentiable) :
+        Instruction("squeeze", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class UnsqueezeInstruction : public Instruction {
+public:
+    UnsqueezeInstruction(int opcode, bool differentiable) :
+        Instruction("unsqueeze", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class RqsReshapeInstruction : public Instruction {
+public:
+    RqsReshapeInstruction(int opcode, bool differentiable) :
+        Instruction("rqs_reshape", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class NonzeroInstruction : public Instruction {
+public:
+    NonzeroInstruction(int opcode, bool differentiable) :
+        Instruction("nonzero", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class BatchGatherInstruction : public Instruction {
+public:
+    BatchGatherInstruction(int opcode, bool differentiable) :
+        Instruction("batch_gather", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class BatchScatterInstruction : public Instruction {
+public:
+    BatchScatterInstruction(int opcode, bool differentiable) :
+        Instruction("batch_scatter", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class RandomInstruction : public Instruction {
+public:
+    RandomInstruction(int opcode, bool differentiable) :
+        Instruction("random", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class UnweightInstruction : public Instruction {
+public:
+    UnweightInstruction(int opcode, bool differentiable) :
+        Instruction("unweight", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+class MatrixElementInstruction : public Instruction {
+public:
+    MatrixElementInstruction(int opcode, bool differentiable) :
+        Instruction("matrix_element", opcode, differentiable) {}
+    TypeVec signature(const ValueVec& args) const override;
+};
+
+using InstructionOwner = std::unique_ptr<const Instruction>;
+using InstructionPtr = Instruction const*;
+const std::unordered_map<std::string, InstructionOwner> build_instruction_set();
+const std::unordered_map<std::string, InstructionOwner> instruction_set =
+    build_instruction_set();
+
+} // namespace madspace

madspace/include/madspace/madcode/opcode_mixin.h

@@ -0,0 +1,134 @@
+stack = 0,
+unstack = 1,
+unstack_sizes = 2,
+pop = 3,
+batch_cat = 4,
+batch_split = 5,
+cat = 6,
+batch_size = 7,
+offset_indices = 8,
+full = 9,
+squeeze = 10,
+unsqueeze = 11,
+add = 12,
+add_int = 13,
+sub = 14,
+mul = 15,
+reduce_sum = 16,
+reduce_sum_vector = 17,
+reduce_product = 18,
+sqrt = 19,
+square = 20,
+min = 21,
+max = 22,
+obs_sqrt_s = 23,
+obs_e = 24,
+obs_px = 25,
+obs_py = 26,
+obs_pz = 27,
+obs_mass = 28,
+obs_pt = 29,
+obs_p_mag = 30,
+obs_phi = 31,
+obs_theta = 32,
+obs_y = 33,
+obs_y_abs = 34,
+obs_eta = 35,
+obs_eta_abs = 36,
+obs_delta_eta = 37,
+obs_delta_phi = 38,
+obs_delta_r = 39,
+boost_beam = 40,
+boost_beam_inverse = 41,
+com_p_in = 42,
+r_to_x1x2 = 43,
+x1x2_to_r = 44,
+diff_cross_section = 45,
+two_body_decay_com = 46,
+two_body_decay_com_inverse = 47,
+two_body_decay = 48,
+two_body_decay_inverse = 49,
+two_to_two_particle_scattering_com = 50,
+two_to_two_particle_scattering_com_inverse = 51,
+two_to_two_particle_scattering = 52,
+two_to_two_particle_scattering_inverse = 53,
+two_to_three_particle_scattering = 54,
+two_to_three_particle_scattering_inverse = 55,
+three_body_decay_com = 56,
+three_body_decay_com_inverse = 57,
+three_body_decay = 58,
+three_body_decay_inverse = 59,
+t_inv_min_max = 60,
+t_inv_value_and_min_max = 61,
+s23_min_max = 62,
+s23_value_and_min_max = 63,
+invariants_from_momenta = 64,
+sde2_channel_weights = 65,
+subchannel_weights = 66,
+apply_subchannel_weights = 67,
+pt_eta_phi_x = 68,
+mirror_momenta = 69,
+momenta_to_x1x2 = 70,
+uniform_invariant = 71,
+uniform_invariant_inverse = 72,
+breit_wigner_invariant = 73,
+breit_wigner_invariant_inverse = 74,
+stable_invariant = 75,
+stable_invariant_inverse = 76,
+stable_invariant_nu = 77,
+stable_invariant_nu_inverse = 78,
+fast_rambo_massless = 79,
+fast_rambo_massless_inverse = 80,
+fast_rambo_massless_com = 81,
+fast_rambo_massive = 82,
+fast_rambo_massive_inverse = 83,
+fast_rambo_massive_com = 84,
+cut_unphysical = 85,
+cut_one = 86,
+cut_all = 87,
+cut_any = 88,
+scale_transverse_energy = 89,
+scale_transverse_mass = 90,
+scale_half_transverse_mass = 91,
+scale_partonic_energy = 92,
+chili_forward = 93,
+chili_inverse = 94,
+matrix_element = 95,
+collect_channel_weights = 96,
+interpolate_pdf = 97,
+interpolate_alpha_s = 98,
+matmul = 99,
+relu = 100,
+leaky_relu = 101,
+elu = 102,
+gelu = 103,
+sigmoid = 104,
+softplus = 105,
+rqs_reshape = 106,
+rqs_find_bin = 107,
+rqs_forward = 108,
+rqs_inverse = 109,
+softmax = 110,
+softmax_prior = 111,
+sample_discrete = 112,
+sample_discrete_inverse = 113,
+sample_discrete_probs = 114,
+sample_discrete_probs_inverse = 115,
+discrete_histogram = 116,
+permute_momenta = 117,
+gather = 118,
+gather_int = 119,
+select_int = 120,
+select = 121,
+select_vector = 122,
+argsort = 123,
+one_hot = 124,
+nonzero = 125,
+batch_gather = 126,
+batch_scatter = 127,
+random = 128,
+unweight = 129,
+vegas_forward = 130,
+vegas_inverse = 131,
+vegas_histogram = 132,
+histogram = 133

madspace/include/madspace/madcode/optimizer.h

@@ -0,0 +1,31 @@
+#pragma once
+
+#include "madspace/madcode/function.h"
+
+#include <vector>
+
+namespace madspace {
+
+class InstructionDependencies {
+public:
+    InstructionDependencies(const Function& function);
+    bool depends(std::size_t test_index, std::size_t dependency_index) {
+        return matrix[test_index * size + dependency_index];
+    }
+
+private:
+    std::size_t size;
+    std::vector<bool> matrix;
+    std::vector<int> ranks;
+};
+
+class LastUseOfLocals {
+public:
+    LastUseOfLocals(const Function& function);
+    std::vector<int>& local_indices(std::size_t index) { return last_used[index]; }
+
+private:
+    std::vector<std::vector<int>> last_used;
+};
+
+} // namespace madspace

madspace/include/madspace/madcode/type.h

@@ -0,0 +1,203 @@
+#pragma once
+
+#include <iostream>
+#include <string>
+#include <unordered_map>
+#include <variant>
+#include <vector>
+
+#include <nlohmann/json.hpp>
+
+namespace madspace {
+
+enum class DataType { dt_int, dt_float, batch_sizes };
+
+using me_int_t = int;
+
+template <typename T>
+concept ScalarType = std::same_as<T, me_int_t> || std::same_as<T, double>;
+
+class BatchSize {
+public:
+    using Named = std::string;
+    class UnnamedBody {
+    public:
+        UnnamedBody() : id(counter++) {}
+        friend std::ostream& operator<<(std::ostream& out, const BatchSize& batch_size);
+        friend void to_json(nlohmann::json& j, const BatchSize& batch_size);
+        bool operator==(const UnnamedBody& other) const { return id == other.id; }
+        bool operator!=(const UnnamedBody& other) const { return id != other.id; }
+
+    private:
+        static std::size_t counter;
+        std::size_t id;
+    };
+    using Unnamed = std::shared_ptr<UnnamedBody>;
+    using One = std::monostate;
+    using Compound = std::unordered_map<std::variant<Named, Unnamed, One>, int>;
+
+    static const BatchSize zero;
+    static const BatchSize one;
+
+    BatchSize(const std::string& name) : value(name) {}
+    BatchSize(One value) : value(value) {}
+    BatchSize() : value(std::make_shared<UnnamedBody>()) {}
+    BatchSize operator+(const BatchSize& other) const { return add(other, 1); }
+    BatchSize operator-(const BatchSize& other) const { return add(other, -1); }
+    bool operator==(const BatchSize& other) const { return value == other.value; }
+    bool operator!=(const BatchSize& other) const { return value != other.value; }
+
+    friend std::ostream& operator<<(std::ostream& out, const BatchSize& batch_size);
+    friend void to_json(nlohmann::json& j, const BatchSize& batch_size);
+    friend void from_json(const nlohmann::json& j, BatchSize& batch_size);
+
+private:
+    BatchSize(Compound value) : value(value) {}
+    BatchSize(Unnamed value) : value(value) {}
+    BatchSize add(const BatchSize& other, int factor) const;
+
+    std::variant<Named, Unnamed, One, Compound> value;
+};
+
+void to_json(nlohmann::json& j, const BatchSize& batch_size);
+void to_json(nlohmann::json& j, const BatchSize& batch_size);
+void from_json(const nlohmann::json& j, BatchSize& batch_size);
+
+struct Type {
+    DataType dtype;
+    BatchSize batch_size;
+    std::vector<int> shape;
+    std::vector<BatchSize> batch_size_list;
+
+    Type(DataType dtype, BatchSize batch_size, const std::vector<int>& shape) :
+        dtype(dtype), batch_size(batch_size), shape(shape) {}
+    Type(const std::vector<BatchSize>& batch_size_list) :
+        dtype(DataType::batch_sizes),
+        batch_size(BatchSize::one),
+        batch_size_list(batch_size_list) {}
+};
+
+std::ostream& operator<<(std::ostream& out, const BatchSize& batch_size);
+std::ostream& operator<<(std::ostream& out, const DataType& dtype);
+std::ostream& operator<<(std::ostream& out, const Type& type);
+
+inline bool operator==(const Type& lhs, const Type& rhs) {
+    return lhs.dtype == rhs.dtype && lhs.batch_size == rhs.batch_size &&
+        lhs.shape == rhs.shape;
+}
+
+inline bool operator!=(const Type& lhs, const Type& rhs) {
+    return lhs.dtype != rhs.dtype || lhs.batch_size != rhs.batch_size ||
+        lhs.shape != rhs.shape;
+}
+
+using TypeVec = std::vector<Type>;
+
+const Type single_float{DataType::dt_float, BatchSize::One{}, {}};
+const Type single_int{DataType::dt_int, BatchSize::One{}, {}};
+inline Type single_float_array(int count) {
+    return {DataType::dt_float, BatchSize::one, {count}};
+}
+inline Type single_int_array(int count) {
+    return {DataType::dt_int, BatchSize::one, {count}};
+}
+inline Type single_float_array_2d(int count1, int count2) {
+    return {DataType::dt_float, BatchSize::one, {count1, count2}};
+}
+inline Type single_int_array_2d(int count1, int count2) {
+    return {DataType::dt_int, BatchSize::one, {count1, count2}};
+}
+
+const BatchSize batch_size = BatchSize("batch_size");
+Type multichannel_batch_size(int count);
+const Type batch_float{DataType::dt_float, batch_size, {}};
+const Type batch_int{DataType::dt_int, batch_size, {}};
+const Type batch_four_vec{DataType::dt_float, batch_size, {4}};
+inline Type batch_float_array(int count) {
+    return {DataType::dt_float, batch_size, {count}};
+}
+inline Type batch_four_vec_array(int count) {
+    return {DataType::dt_float, batch_size, {count, 4}};
+}
+
+using TensorValue = std::tuple<
+    std::vector<int>,
+    std::variant<std::vector<me_int_t>, std::vector<double>>>; // TODO: make this a
+                                                               // class
+
+using LiteralValue = std::variant<me_int_t, double, TensorValue, std::monostate>;
+
+struct Value {
+    Type type;
+    LiteralValue literal_value;
+    int local_index = -1;
+
+    Value() : type(single_float), literal_value(std::monostate{}) {}
+
+    Value(me_int_t value) : type(single_int), literal_value(value) {}
+    Value(double value) : type(single_float), literal_value(value) {}
+
+    template <ScalarType T>
+    Value(const std::vector<std::vector<T>>& values) :
+        Value(
+            [&] {
+                std::size_t outer_size = values.size();
+                std::size_t inner_size = values.at(0).size();
+                std::vector<T> flat_values;
+                for (auto& vec : values) {
+                    if (vec.size() != inner_size) {
+                        throw std::invalid_argument(
+                            "All inner vectors must have the same size"
+                        );
+                    }
+                }
+                for (std::size_t j = 0; j < inner_size; ++j) {
+                    for (std::size_t i = 0; i < outer_size; ++i) {
+                        flat_values.push_back(values.at(i).at(j));
+                    }
+                }
+                return flat_values;
+            }(),
+            {static_cast<int>(values.size()), static_cast<int>(values.at(0).size())}
+        ) {}
+
+    template <ScalarType T>
+    Value(const std::vector<T>& values, const std::vector<int>& shape = {}) :
+        type{
+            std::is_same_v<T, me_int_t> ? DataType::dt_int : DataType::dt_float,
+            BatchSize::one,
+            shape.size() == 0 ? std::vector<int>{static_cast<int>(values.size())}
+                              : shape
+        },
+        literal_value(TensorValue(type.shape, values)) {
+        std::size_t prod = 1;
+        for (auto size : type.shape) {
+            prod *= size;
+        }
+        if (prod != values.size()) {
+            throw std::invalid_argument(
+                "size of value vector not compatible with given shape"
+            );
+        }
+    }
+
+    Value(Type _type, int _local_index) :
+        type(_type), literal_value(std::monostate{}), local_index(_local_index) {}
+    Value(Type _type, LiteralValue _literal_value, int _local_index = -1) :
+        type(_type), literal_value(_literal_value), local_index(_local_index) {}
+
+    operator bool() {
+        return !(
+            local_index == -1 && std::holds_alternative<std::monostate>(literal_value)
+        );
+    }
+};
+
+using ValueVec = std::vector<Value>;
+
+void to_json(nlohmann::json& j, const DataType& dtype);
+void to_json(nlohmann::json& j, const Value& value);
+void from_json(const nlohmann::json& j, DataType& dtype);
+void from_json(const nlohmann::json& j, Value& dtype);
+
+} // namespace madspace

madspace/include/madspace/madcode.h

@@ -0,0 +1,6 @@
+#pragma once
+
+#include "madcode/function.h"
+#include "madcode/instruction.h"
+#include "madcode/optimizer.h"
+#include "madcode/type.h"

madspace/include/madspace/phasespace/base.h

@@ -0,0 +1,74 @@
+#pragma once
+
+#include "madspace/madcode.h"
+
+namespace madspace {
+
+class Mapping {
+public:
+    using Result = std::tuple<ValueVec, Value>;
+
+    Mapping(
+        const std::string& name,
+        const TypeVec& input_types,
+        const TypeVec& output_types,
+        const TypeVec& condition_types
+    ) :
+        _name(name),
+        _input_types(input_types),
+        _output_types(output_types),
+        _condition_types(condition_types) {}
+    virtual ~Mapping() = default;
+    Result build_forward(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions = {}
+    ) const;
+    Result build_inverse(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions = {}
+    ) const;
+    Function forward_function() const;
+    Function inverse_function() const;
+    const TypeVec& input_types() const { return _input_types; }
+    const TypeVec& output_types() const { return _output_types; }
+    const TypeVec& condition_types() const { return _condition_types; }
+    const std::string& name() const { return _name; }
+
+protected:
+    // TODO: make parameters const ref
+    virtual Result build_forward_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const = 0;
+    virtual Result build_inverse_impl(
+        FunctionBuilder& fb, const ValueVec& inputs, const ValueVec& conditions
+    ) const = 0;
+
+private:
+    std::string _name;
+    TypeVec _input_types;
+    TypeVec _output_types;
+    TypeVec _condition_types;
+};
+
+class FunctionGenerator {
+public:
+    FunctionGenerator(
+        const std::string& name, const TypeVec& arg_types, const TypeVec& return_types
+    ) :
+        _name(name), _arg_types(arg_types), _return_types(return_types) {}
+    virtual ~FunctionGenerator() = default;
+    ValueVec build_function(FunctionBuilder& fb, const ValueVec& args) const;
+    Function function() const;
+    const TypeVec& arg_types() const { return _arg_types; }
+    const TypeVec& return_types() const { return _return_types; }
+    const std::string& name() const { return _name; }
+
+protected:
+    virtual ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const = 0;
+
+private:
+    std::string _name;
+    TypeVec _arg_types;
+    TypeVec _return_types;
+};
+
+} // namespace madspace

madspace/include/madspace/phasespace/channel_weight_network.h

@@ -0,0 +1,46 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/phasespace/mlp.h"
+
+namespace madspace {
+
+class MomentumPreprocessing : public FunctionGenerator {
+public:
+    MomentumPreprocessing(std::size_t particle_count);
+    std::size_t output_dim() const { return _output_dim; };
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::size_t _output_dim;
+};
+
+class ChannelWeightNetwork : public FunctionGenerator {
+public:
+    ChannelWeightNetwork(
+        std::size_t channel_count,
+        std::size_t particle_count,
+        std::size_t hidden_dim = 32,
+        std::size_t layers = 3,
+        MLP::Activation activation = MLP::leaky_relu,
+        const std::string& prefix = ""
+    );
+
+    const MLP& mlp() const { return _mlp; }
+    const MomentumPreprocessing& preprocessing() const { return _preprocessing; }
+    void initialize_globals(ContextPtr context) const;
+    const std::string& mask_name() const { return _mask_name; }
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    MomentumPreprocessing _preprocessing;
+    MLP _mlp;
+    std::size_t _channel_count;
+    std::string _mask_name;
+};
+
+} // namespace madspace