madspace 0.3.1__cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

madspace/include/madspace/phasespace/vegas.h

@@ -0,0 +1,53 @@
+#pragma once
+
+#include "madspace/phasespace/base.h"
+#include "madspace/runtime/context.h"
+
+namespace madspace {
+
+class VegasHistogram : public FunctionGenerator {
+public:
+    VegasHistogram(std::size_t dimension, std::size_t bin_count);
+
+private:
+    ValueVec
+    build_function_impl(FunctionBuilder& fb, const ValueVec& args) const override;
+
+    std::size_t _bin_count;
+};
+
+class VegasMapping : public Mapping {
+public:
+    VegasMapping(
+        std::size_t dimension, std::size_t bin_count, const std::string& prefix = ""
+    ) :
+        Mapping(
+            "VegasMapping",
+            {batch_float_array(dimension)},
+            {batch_float_array(dimension)},
+            {}
+        ),
+        _dimension(dimension),
+        _bin_count(bin_count),
+        _grid_name(prefixed_name(prefix, "vegas_grid")) {}
+    const std::string& grid_name() const { return _grid_name; }
+    void initialize_globals(ContextPtr context) const;
+    std::size_t dimension() const { return _dimension; }
+    std::size_t bin_count() const { return _bin_count; }
+
+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 _dimension;
+    std::size_t _bin_count;
+    std::string _grid_name;
+};
+
+void initialize_vegas_grid(ContextPtr context, const std::string& grid_name);
+
+} // namespace madspace

madspace/include/madspace/phasespace.h

@@ -0,0 +1,27 @@
+#pragma once
+
+#include "phasespace/base.h"
+#include "phasespace/channel_weight_network.h"
+#include "phasespace/channel_weights.h"
+#include "phasespace/cross_section.h"
+#include "phasespace/cuts.h"
+#include "phasespace/discrete_flow.h"
+#include "phasespace/discrete_sampler.h"
+#include "phasespace/flow.h"
+#include "phasespace/histograms.h"
+#include "phasespace/integrand.h"
+#include "phasespace/invariants.h"
+#include "phasespace/luminosity.h"
+#include "phasespace/matrix_element.h"
+#include "phasespace/mlp.h"
+#include "phasespace/multichannel.h"
+#include "phasespace/observable.h"
+#include "phasespace/pdf.h"
+#include "phasespace/phasespace.h"
+#include "phasespace/rambo.h"
+#include "phasespace/scale.h"
+#include "phasespace/t_propagator_mapping.h"
+#include "phasespace/three_particle.h"
+#include "phasespace/topology.h"
+#include "phasespace/two_particle.h"
+#include "phasespace/vegas.h"

madspace/include/madspace/runtime/context.h

@@ -0,0 +1,147 @@
+#pragma once
+
+#include <stdint.h>
+#include <unordered_map>
+
+#include "madspace/madcode.h"
+#include "madspace/runtime/tensor.h"
+#include "madspace/runtime/thread_pool.h"
+#include "madspace/umami.h"
+
+namespace madspace {
+
+class MatrixElementApi {
+public:
+    MatrixElementApi(
+        const std::string& file, const std::string& param_card, std::size_t index = 0
+    );
+    MatrixElementApi(MatrixElementApi&&) noexcept = default;
+    MatrixElementApi& operator=(MatrixElementApi&&) noexcept = default;
+    MatrixElementApi(const MatrixElementApi&) = delete;
+    MatrixElementApi& operator=(const MatrixElementApi&) = delete;
+    DevicePtr device() const {
+        UmamiDevice dev;
+        check_umami_status(_get_meta(UMAMI_META_DEVICE, &dev));
+        switch (dev) {
+        case UMAMI_DEVICE_CPU:
+            return cpu_device();
+        case UMAMI_DEVICE_CUDA:
+            return cuda_device();
+        case UMAMI_DEVICE_HIP:
+            return hip_device();
+        default:
+            throw_error("matrix element device not known");
+        }
+    }
+    std::size_t particle_count() const {
+        int count;
+        check_umami_status(_get_meta(UMAMI_META_PARTICLE_COUNT, &count));
+        return count;
+    }
+    std::size_t diagram_count() const {
+        int count;
+        check_umami_status(_get_meta(UMAMI_META_DIAGRAM_COUNT, &count));
+        return count;
+    }
+    std::size_t helicity_count() const {
+        int count;
+        check_umami_status(_get_meta(UMAMI_META_HELICITY_COUNT, &count));
+        return count;
+    }
+    std::size_t index() const { return _index; }
+    const std::string& file_name() const { return _file_name; }
+
+    void call(
+        UmamiHandle handle,
+        size_t count,
+        size_t stride,
+        size_t offset,
+        size_t input_count,
+        UmamiInputKey const* input_keys,
+        void const* const* inputs,
+        size_t output_count,
+        UmamiOutputKey const* output_keys,
+        void* const* outputs
+    ) const {
+        check_umami_status(_matrix_element(
+            handle,
+            count,
+            stride,
+            offset,
+            input_count,
+            input_keys,
+            inputs,
+            output_count,
+            output_keys,
+            outputs
+        ));
+    }
+
+    void* process_instance(std::size_t index) const {
+        return _instances.get(index).get();
+    }
+
+private:
+    void check_umami_status(UmamiStatus status) const;
+    [[noreturn]] void throw_error(const std::string& message) const;
+    std::unique_ptr<void, std::function<void(void*)>> _shared_lib;
+    decltype(&umami_get_meta) _get_meta;
+    decltype(&umami_initialize) _initialize;
+    decltype(&umami_matrix_element) _matrix_element;
+    decltype(&umami_free) _free;
+    using InstanceType = std::unique_ptr<void, std::function<void(void*)>>;
+    ThreadResource<InstanceType> _instances;
+    std::string _file_name;
+    std::size_t _index;
+};
+
+class Context {
+    /**
+     * Contains global variables and matrix elements
+     */
+public:
+    Context() : _device(cpu_device()) {}
+    Context(DevicePtr device) : _device(device) {}
+    Context(Context&&) = default;
+    Context& operator=(Context&&) = default;
+    Context(const Context&) = delete;
+    Context& operator=(const Context&) = delete;
+    const MatrixElementApi&
+    load_matrix_element(const std::string& file, const std::string& param_card);
+    Tensor define_global(
+        const std::string& name,
+        DataType dtype,
+        const SizeVec& shape,
+        bool requires_grad = false
+    );
+    Tensor global(const std::string& name);
+    bool global_requires_grad(const std::string& name);
+    bool global_exists(const std::string& name);
+    std::vector<std::string> global_names() const;
+    void delete_global(const std::string& name);
+    const MatrixElementApi& matrix_element(std::size_t index) const;
+    void save(const std::string& file) const;
+    void load(const std::string& file);
+    DevicePtr device() { return _device; }
+
+private:
+    // make sure that ThreadPool outlives any Context instance
+    inline static ThreadPool& thread_pool_ref = default_thread_pool();
+    DevicePtr _device;
+    std::unordered_map<std::string, std::tuple<Tensor, bool>> _globals;
+    std::vector<MatrixElementApi> _matrix_elements;
+    std::vector<std::string> _param_card_paths;
+};
+
+using ContextPtr = std::shared_ptr<Context>;
+
+ContextPtr default_context();
+ContextPtr default_cuda_context();
+ContextPtr default_hip_context();
+ContextPtr default_device_context(DevicePtr device);
+
+inline std::string prefixed_name(const std::string& prefix, const std::string& name) {
+    return prefix == "" ? name : std::format("{}.{}", prefix, name);
+}
+
+} // namespace madspace

madspace/include/madspace/runtime/discrete_optimizer.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "madspace/madcode.h"
+#include "madspace/runtime/context.h"
+#include "madspace/runtime/tensor.h"
+
+namespace madspace {
+
+class DiscreteOptimizer {
+public:
+    DiscreteOptimizer(ContextPtr context, const std::vector<std::string>& prob_names) :
+        _context(context), _prob_names(prob_names), _sample_count(7000) {}
+    void add_data(const std::vector<Tensor>& values_and_counts);
+    void optimize();
+
+private:
+    ContextPtr _context;
+    std::vector<std::string> _prob_names;
+    double _damping;
+    std::size_t _sample_count;
+    std::vector<std::tuple<std::vector<std::size_t>, std::vector<double>>> _data;
+};
+
+} // namespace madspace

madspace/include/madspace/runtime/event_generator.h

@@ -0,0 +1,257 @@
+#pragma once
+
+#include <chrono>
+#include <optional>
+#include <random>
+#include <vector>
+
+#include <nlohmann/json.hpp>
+
+#include "madspace/madcode.h"
+#include "madspace/phasespace.h"
+#include "madspace/runtime/discrete_optimizer.h"
+#include "madspace/runtime/format.h"
+#include "madspace/runtime/io.h"
+#include "madspace/runtime/runtime_base.h"
+#include "madspace/runtime/vegas_optimizer.h"
+
+namespace madspace {
+
+class RunningIntegral {
+public:
+    RunningIntegral() : _mean(0), _var_sum(0), _count(0) {}
+    double mean() const { return _mean; }
+    double variance() const { return _count > 1 ? _var_sum / (_count - 1) : 0; }
+    double error() const { return std::sqrt(variance() / _count); }
+    double rel_error() const { return error() / mean(); }
+    double rel_std_dev() const { return std::sqrt(variance()) / _mean; }
+    std::size_t count() const { return _count; }
+    void reset() {
+        _mean = 0;
+        _var_sum = 0;
+        _count = 0;
+    }
+    void push(double value) {
+        ++_count;
+        if (_count == 1) {
+            _mean = value;
+            _var_sum = 0;
+        } else {
+            double mean_diff = value - _mean;
+            _mean += mean_diff / _count;
+            _var_sum += mean_diff * (value - _mean);
+        }
+    }
+
+private:
+    double _mean;
+    double _var_sum;
+    std::size_t _count;
+};
+
+class EventGenerator {
+public:
+    static inline const int integrand_flags = Integrand::sample |
+        Integrand::return_momenta | Integrand::return_indices |
+        Integrand::return_random | Integrand::return_discrete;
+    enum Verbosity { silent, log, pretty };
+    struct Config {
+        std::size_t target_count = 10000;
+        double vegas_damping = 0.2;
+        double max_overweight_truncation = 0.01;
+        std::size_t freeze_max_weight_after = 10000;
+        std::size_t start_batch_size = 1000;
+        std::size_t max_batch_size = 64000;
+        std::size_t survey_min_iters = 3;
+        std::size_t survey_max_iters = 4;
+        double survey_target_precision = 0.1;
+        std::size_t optimization_patience = 3;
+        double optimization_threshold = 0.99;
+        std::size_t batch_size = 1000;
+        Verbosity verbosity = silent;
+        bool write_live_data = false;
+    };
+    static const Config default_config;
+    struct Status {
+        std::size_t index;
+        std::size_t subprocess;
+        std::string name;
+        double mean;
+        double error;
+        double rel_std_dev;
+        std::size_t count;
+        std::size_t count_opt;
+        std::size_t count_after_cuts;
+        std::size_t count_after_cuts_opt;
+        double count_unweighted;
+        double count_target;
+        std::size_t iterations;
+        bool optimized;
+        bool done;
+    };
+    struct Histogram {
+        std::string name;
+        double min;
+        double max;
+        std::vector<double> bin_values;
+        std::vector<double> bin_errors;
+    };
+    static void set_abort_check_function(std::function<void(void)> func) {
+        _abort_check_function = func;
+    }
+
+    EventGenerator(
+        ContextPtr context,
+        const std::vector<Integrand>& channels,
+        const std::string& temp_file_prefix,
+        const std::string& status_file = "",
+        const Config& config = default_config,
+        const std::vector<std::size_t>& channel_subprocesses = {},
+        const std::vector<std::string>& channel_names = {},
+        const std::vector<ObservableHistograms>& channel_histograms = {}
+    );
+    void survey();
+    void generate();
+    void combine_to_compact_npy(const std::string& file_name);
+    void combine_to_lhe_npy(const std::string& file_name, LHECompleter& lhe_completer);
+    void combine_to_lhe(const std::string& file_name, LHECompleter& lhe_completer);
+    Status status() const { return _status_all; }
+    std::vector<Status> channel_status() const;
+    std::vector<Histogram> histograms() const;
+
+private:
+    struct ChannelState {
+        std::size_t index;
+        RuntimePtr runtime;
+        EventFile event_file;
+        EventFile weight_file;
+        std::optional<VegasGridOptimizer> vegas_optimizer;
+        RuntimePtr vegas_histogram;
+        std::optional<DiscreteOptimizer> discrete_optimizer;
+        RuntimePtr discrete_histogram;
+        RuntimePtr observable_histograms;
+        std::size_t batch_size;
+        std::string name;
+        std::size_t subprocess_index;
+        RunningIntegral cross_section;
+        bool needs_optimization = true;
+        double max_weight = 0.;
+        double eff_count = 0.;
+        double integral_fraction = 1.;
+        std::size_t total_sample_count = 0;
+        std::size_t total_sample_count_opt = 0;
+        std::size_t total_sample_count_after_cuts = 0;
+        std::size_t total_sample_count_after_cuts_opt = 0;
+        std::size_t iterations = 0;
+        std::size_t iters_without_improvement = 0;
+        double best_rsd = std::numeric_limits<double>::max();
+        std::vector<double> large_weights;
+        std::size_t job_count = 0;
+        nested_vector2<std::pair<double, double>> histograms;
+    };
+    struct RunningJob {
+        std::size_t channel_index;
+        TensorVec events;
+        std::size_t vegas_job_count;
+    };
+    struct CombineChannelData {
+        std::size_t cum_count;
+        EventBuffer event_buffer;
+        EventBuffer weight_buffer;
+        std::size_t buffer_index;
+    };
+    struct TimingData {
+        double wall_time_sec;
+        double cpu_time_sec;
+    };
+    inline static std::function<void(void)> _abort_check_function = [] {};
+
+    ContextPtr _context;
+    Config _config;
+    std::vector<ChannelState> _channels;
+    RuntimePtr _unweighter;
+    Status _status_all;
+    std::unordered_map<std::size_t, RunningJob> _running_jobs;
+    std::size_t _job_id;
+    std::chrono::time_point<std::chrono::steady_clock> _start_time;
+    std::size_t _start_cpu_microsec;
+    std::chrono::time_point<std::chrono::steady_clock> _last_print_time;
+    std::chrono::time_point<std::chrono::steady_clock> _last_status_time;
+    PrettyBox _pretty_box_upper;
+    PrettyBox _pretty_box_lower;
+    std::string _status_file;
+    std::unordered_map<std::string, TimingData> _timing_data;
+    std::vector<Histogram> _empty_histograms;
+
+    void reset_start_time();
+    void add_timing_data(const std::string& key);
+    std::string format_run_time(const std::string& key) const;
+    void unweight_all();
+    void unweight_channel(ChannelState& channel, std::mt19937 rand_gen);
+    std::tuple<Tensor, std::vector<Tensor>> integrate_and_optimize(
+        ChannelState& channel, TensorVec& events, bool always_optimize
+    );
+    double channel_weight_sum(ChannelState& channel, std::size_t event_count);
+    void start_job(
+        ChannelState& channel, std::size_t batch_size, std::size_t vegas_job_count = 0
+    );
+    void start_vegas_jobs(ChannelState& channel);
+    void clear_channel(ChannelState& channel);
+    void update_max_weight(ChannelState& channel, Tensor weights);
+    void unweight_and_write(ChannelState& channel, const std::vector<Tensor>& momenta);
+    std::size_t max_particle_count();
+    std::tuple<std::vector<CombineChannelData>, std::size_t, double> init_combine();
+    void read_and_combine(
+        std::vector<CombineChannelData>& channel_data,
+        EventBuffer& buffer,
+        double norm_factor
+    );
+    void fill_lhe_event(
+        LHECompleter& lhe_completer,
+        LHEEvent& lhe_event,
+        EventBuffer& buffer,
+        std::size_t event_index
+    );
+
+    void init_status(const std::string& status);
+    void write_status(const std::string& status, bool force_write);
+
+    void print_survey_init();
+    void print_survey_update(
+        bool done,
+        std::size_t done_job_count,
+        std::size_t total_job_count,
+        std::size_t iter
+    );
+    void print_survey_update_pretty(
+        bool done,
+        std::size_t done_job_count,
+        std::size_t total_job_count,
+        std::size_t iter
+    );
+    void print_survey_update_log(
+        bool done,
+        std::size_t done_job_count,
+        std::size_t total_job_count,
+        std::size_t iter
+    );
+
+    void print_gen_init();
+    void print_gen_update(bool done);
+    void print_gen_update_pretty(bool done);
+    void print_gen_update_log(bool done);
+
+    void print_combine_init();
+    void print_combine_update(std::size_t count);
+    void print_combine_update_pretty(std::size_t count);
+    void print_combine_update_log(std::size_t count);
+
+    friend void
+    to_json(nlohmann::json& j, const EventGenerator::TimingData& timing_data);
+};
+
+void to_json(nlohmann::json& j, const EventGenerator::TimingData& timing_data);
+void to_json(nlohmann::json& j, const EventGenerator::Status& status);
+void to_json(nlohmann::json& j, const EventGenerator::Histogram& hist);
+
+} // namespace madspace

madspace/include/madspace/runtime/format.h

@@ -0,0 +1,68 @@
+#pragma once
+
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace madspace {
+
+std::string format_si_prefix(double value);
+std::string format_with_error(double value, double error);
+std::string format_progress(double progress, int width);
+
+class PrettyBox {
+public:
+    PrettyBox() = default;
+    PrettyBox(
+        const std::string& title,
+        std::size_t rows,
+        const std::vector<std::size_t>& column_sizes,
+        std::size_t offset = 0,
+        std::size_t box_width = 91
+    );
+
+    void print_first() const;
+    void print_update() const;
+    std::size_t line_count() const { return _rows + 3; }
+
+    void set_row(std::size_t row, const std::vector<std::string>& values) {
+        if (row >= _rows) {
+            throw std::out_of_range("row index out of range");
+        }
+        for (std::size_t column = 0; auto& value : values) {
+            _content.at(row * _columns + column) = value;
+            ++column;
+        }
+    }
+
+    void set_column(std::size_t column, const std::vector<std::string>& values) {
+        if (column >= _columns) {
+            throw std::out_of_range("column index out of range");
+        }
+        for (std::size_t row = 0; auto& value : values) {
+            _content.at(row * _columns + column) = value;
+            ++row;
+        }
+    }
+
+    void set_cell(std::size_t row, std::size_t column, std::string value) {
+        if (row >= _rows) {
+            throw std::out_of_range("row index out of range");
+        }
+        if (column >= _columns) {
+            throw std::out_of_range("column index out of range");
+        }
+        _content.at(row * _columns + column) = value;
+    }
+
+private:
+    std::string _header;
+    std::string _footer;
+    std::size_t _rows;
+    std::size_t _columns;
+    std::size_t _offset;
+    std::vector<std::size_t> _column_ends;
+    std::vector<std::string> _content;
+};
+
+} // namespace madspace