madspace 0.3.1__cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

madspace/include/madspace/runtime/vegas_optimizer.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "madspace/madcode.h"
+#include "madspace/runtime/context.h"
+#include "madspace/runtime/tensor.h"
+
+namespace madspace {
+
+class VegasGridOptimizer {
+public:
+    VegasGridOptimizer(
+        ContextPtr context, const std::string& grid_name, double damping
+    ) :
+        _context(context), _grid_name(grid_name), _damping(damping) {}
+    void add_data(Tensor weights, Tensor inputs);
+    void optimize();
+    std::size_t input_dim() const;
+
+private:
+    ContextPtr _context;
+    std::string _grid_name;
+    double _damping;
+    std::vector<std::tuple<std::vector<std::size_t>, std::vector<double>>> _data;
+};
+
+} // namespace madspace

madspace/include/madspace/runtime.h

@@ -0,0 +1,12 @@
+#pragma once
+
+#include "runtime/context.h"
+#include "runtime/discrete_optimizer.h"
+#include "runtime/event_generator.h"
+#include "runtime/format.h"
+#include "runtime/io.h"
+#include "runtime/lhe_output.h"
+#include "runtime/logger.h"
+#include "runtime/tensor.h"
+#include "runtime/thread_pool.h"
+#include "runtime/vegas_optimizer.h"

madspace/include/madspace/umami.h

@@ -0,0 +1,202 @@
+/*
+ *                                   _
+ *                                  (_)
+ *   _   _ _ __ ___   __ _ _ __ ___  _
+ *  | | | | '_ ` _ \ / _` | '_ ` _ \| |
+ *  | |_| | | | | | | (_| | | | | | | |
+ *   \__,_|_| |_| |_|\__,_|_| |_| |_|_|
+ *
+ *  Unified  MAtrix  eleMent  Interface
+ *
+ *
+ */
+
+#ifndef UMAMI_HEADER
+#define UMAMI_HEADER 1
+
+#include <stddef.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * Major version number of the UMAMI interface. If the major version is the same
+ * between caller and implementation, binary compatibility is ensured.
+ */
+const inline int UMAMI_MAJOR_VERSION = 1;
+/**
+ * Minor version number of the UMAMI interface. Between minor versions, new keys for
+ * errors, devices, metadata, inputs and outputs can be added.
+ */
+const inline int UMAMI_MINOR_VERSION = 0;
+
+typedef enum {
+    UMAMI_SUCCESS,
+    UMAMI_ERROR,
+    UMAMI_ERROR_NOT_IMPLEMENTED,
+    UMAMI_ERROR_UNSUPPORTED_INPUT,
+    UMAMI_ERROR_UNSUPPORTED_OUTPUT,
+    UMAMI_ERROR_UNSUPPORTED_META,
+    UMAMI_ERROR_MISSING_INPUT,
+} UmamiStatus;
+
+typedef enum {
+    UMAMI_DEVICE_CPU,
+    UMAMI_DEVICE_CUDA,
+    UMAMI_DEVICE_HIP,
+} UmamiDevice;
+
+typedef enum {
+    UMAMI_META_DEVICE,
+    UMAMI_META_PARTICLE_COUNT,
+    UMAMI_META_DIAGRAM_COUNT,
+    UMAMI_META_HELICITY_COUNT,
+    UMAMI_META_COLOR_COUNT,
+} UmamiMetaKey;
+
+typedef enum {
+    UMAMI_IN_MOMENTA,
+    UMAMI_IN_ALPHA_S,
+    UMAMI_IN_FLAVOR_INDEX,
+    UMAMI_IN_RANDOM_COLOR,
+    UMAMI_IN_RANDOM_HELICITY,
+    UMAMI_IN_RANDOM_DIAGRAM,
+    UMAMI_IN_HELICITY_INDEX,
+    UMAMI_IN_DIAGRAM_INDEX,
+    UMAMI_IN_GPU_STREAM,
+} UmamiInputKey;
+
+typedef enum {
+    UMAMI_OUT_MATRIX_ELEMENT,
+    UMAMI_OUT_DIAGRAM_AMP2,
+    UMAMI_OUT_COLOR_INDEX,
+    UMAMI_OUT_HELICITY_INDEX,
+    UMAMI_OUT_DIAGRAM_INDEX,
+    // NLO: born, virtual, poles, counterterms
+    // color: LC-ME, FC-ME
+} UmamiOutputKey;
+
+typedef void* UmamiHandle;
+
+/**
+ * Creates an instance of the matrix element. Each instance is independent, so thread
+ * safety can be achieved by creating a separate one for every thread.
+ *
+ * @param meta_key
+ *     path to the parameter file
+ * @param handle
+ *     pointer to an instance of the subprocess. Has to be cleaned up by
+ *     the caller with `free_subprocess`.
+ * @return
+ *     UMAMI_SUCCESS on success, error code otherwise
+ */
+UmamiStatus umami_get_meta(UmamiMetaKey meta_key, void* result);
+
+/**
+ * Creates an instance of the matrix element. Each instance is independent, so thread
+ * safety can be achieved by creating a separate one for every thread.
+ *
+ * @param param_card_path
+ *     path to the parameter file
+ * @param handle
+ *     pointer to an instance of the subprocess. Has to be cleaned up by
+ *     the caller with `free_subprocess`.
+ * @return
+ *     UMAMI_SUCCESS on success, error code otherwise
+ */
+UmamiStatus umami_initialize(UmamiHandle* handle, char const* param_card_path);
+
+/**
+ * Sets the value of a model parameter
+ *
+ * @param handle
+ *     handle of a matrix element instance
+ * @param name
+ *     name of the parameter
+ * @param parameter_real
+ *     real part of the parameter value
+ * @param parameter_imag
+ *     imaginary part of the parameter value. Ignored for real valued parameters.
+ * @return
+ *     UMAMI_SUCCESS on success, error code otherwise
+ */
+UmamiStatus umami_set_parameter(
+    UmamiHandle handle, char const* name, double parameter_real, double parameter_imag
+);
+
+/**
+ * Retrieves the value of a model parameter
+ *
+ * @param handle
+ *     handle of a matrix element instance
+ * @param name
+ *     name of the parameter
+ * @param parameter_real
+ *     pointer to double to return real part of the parameter value
+ * @param parameter_imag
+ *     pointer to double to return imaginary part of the parameter value. Ignored
+ *     for real-valued parameters (i.e. you may pass a null pointer)
+ * @return
+ *     UMAMI_SUCCESS on success, error code otherwise
+ */
+UmamiStatus umami_get_parameter(
+    UmamiHandle handle, char const* name, double* parameter_real, double* parameter_imag
+);
+
+/**
+ * Evaluates the matrix element as a function of the given inputs, filling the
+ * requested outputs.
+ *
+ * @param handle
+ *     handle of a matrix element instance
+ * @param count
+ *     number of events to evaluate the matrix element for
+ * @param stride
+ *     stride of the batch dimension of the input and output arrays, see memory layout
+ * @param offset
+ *     offset of the event index
+ * @param input_count
+ *     number of inputs to the matrix element
+ * @param input_keys
+ *     pointer to an array of input keys, length `input_count`
+ * @param inputs
+ *     pointer to an array of void pointers to the inputs. The type of the inputs
+ *     depends on the input key
+ * @param output_count
+ *     number of outputs to the matrix element
+ * @param output_keys
+ *     pointer to an array of output keys, length `output_count`
+ * @param outputs
+ *     pointer to an array of void pointers to the outputs. The type of the outputs
+ *     depends on the output key. The caller is responsible for allocating memory for
+ *     the outputs.
+ * @return
+ *     UMAMI_SUCCESS on success, error code otherwise
+ */
+UmamiStatus umami_matrix_element(
+    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
+);
+
+/**
+ * Frees matrix element instance
+ *
+ * @param handle
+ *     handle of a matrix element instance
+ */
+UmamiStatus umami_free(UmamiHandle handle);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // UMAMI_HEADER

madspace/include/madspace/util.h

@@ -0,0 +1,142 @@
+#pragma once
+
+#include <cstdio>
+#include <format>
+#include <ranges>
+#include <tuple>
+#include <vector>
+
+namespace madspace {
+
+template <class... Ts>
+struct Overloaded : Ts... {
+    using Ts::operator()...;
+};
+template <class... Ts>
+Overloaded(Ts...) -> Overloaded<Ts...>;
+
+template <typename T>
+using nested_vector2 = std::vector<std::vector<T>>;
+template <typename T>
+using nested_vector3 = std::vector<std::vector<std::vector<T>>>;
+template <typename T>
+using nested_vector4 = std::vector<std::vector<std::vector<std::vector<T>>>>;
+
+// Unfortunately nvcc does not support C++23 yet, so we implement our own zip function
+// here (based on https://github.com/alemuntoni/zip-views), otherwise use the standard
+// library function
+
+namespace detail {
+
+inline void print_impl(
+    std::FILE* stream, bool new_line, std::string_view fmt, std::format_args args
+) {
+    std::string str = std::vformat(fmt, args);
+    if (new_line) {
+        str.push_back('\n');
+    }
+    fwrite(str.data(), 1, str.size(), stream);
+}
+
+template <typename... Args, std::size_t... Index>
+bool any_match_impl(
+    const std::tuple<Args...>& lhs,
+    const std::tuple<Args...>& rhs,
+    std::index_sequence<Index...>
+) {
+    auto result = false;
+    result = (... || (std::get<Index>(lhs) == std::get<Index>(rhs)));
+    return result;
+}
+
+template <typename... Args>
+bool any_match(const std::tuple<Args...>& lhs, const std::tuple<Args...>& rhs) {
+    return any_match_impl(lhs, rhs, std::index_sequence_for<Args...>{});
+}
+
+template <std::ranges::viewable_range... Rng>
+class zip_iterator {
+public:
+    using value_type = std::tuple<std::ranges::range_reference_t<Rng>...>;
+
+    zip_iterator() = delete;
+    zip_iterator(std::ranges::iterator_t<Rng>&&... iters) :
+        _iters{std::forward<std::ranges::iterator_t<Rng>>(iters)...} {}
+
+    zip_iterator& operator++() {
+        std::apply([](auto&&... args) { ((++args), ...); }, _iters);
+        return *this;
+    }
+
+    zip_iterator operator++(int) {
+        auto tmp = *this;
+        ++*this;
+        return tmp;
+    }
+
+    bool operator!=(const zip_iterator& other) const { return !(*this == other); }
+
+    bool operator==(const zip_iterator& other) const {
+        return any_match(_iters, other._iters);
+    }
+
+    value_type operator*() {
+        return std::apply([](auto&&... args) { return value_type(*args...); }, _iters);
+    }
+
+private:
+    std::tuple<std::ranges::iterator_t<Rng>...> _iters;
+};
+
+template <std::ranges::viewable_range... T>
+class zipper {
+public:
+    using zip_type = zip_iterator<T...>;
+
+    template <typename... Args>
+    zipper(Args&&... args) : _args{std::forward<Args>(args)...} {}
+
+    zip_type begin() {
+        return std::apply(
+            [](auto&&... args) { return zip_type(std::ranges::begin(args)...); }, _args
+        );
+    }
+    zip_type end() {
+        return std::apply(
+            [](auto&&... args) { return zip_type(std::ranges::end(args)...); }, _args
+        );
+    }
+
+private:
+    std::tuple<T...> _args;
+};
+
+} // namespace detail
+
+template <std::ranges::viewable_range... T>
+auto zip(T&&... t) {
+    return detail::zipper<T...>{std::forward<T>(t)...};
+}
+
+template <typename... Args>
+inline void print(std::format_string<Args...> fmt, Args&&... args) {
+    detail::print_impl(stdout, false, fmt.get(), std::make_format_args(args...));
+}
+
+template <typename... Args>
+inline void print(std::FILE* stream, std::format_string<Args...> fmt, Args&&... args) {
+    detail::print_impl(stream, false, fmt.get(), std::make_format_args(args...));
+}
+
+template <typename... Args>
+inline void println(std::format_string<Args...> fmt, Args&&... args) {
+    detail::print_impl(stdout, true, fmt.get(), std::make_format_args(args...));
+}
+
+template <typename... Args>
+inline void
+println(std::FILE* stream, std::format_string<Args...> fmt, Args&&... args) {
+    detail::print_impl(stream, true, fmt.get(), std::make_format_args(args...));
+}
+
+} // namespace madspace

madspace/madnis/__init__.py

@@ -0,0 +1,44 @@
+"""
+This module contains functions and classes to train neural importance sampling networks and
+evaluate the integration and sampling performance.
+"""
+
+from .buffer import Buffer
+from .channel_grouping import ChannelData, ChannelGroup, ChannelGrouping
+from .distribution import Distribution
+from .integrand import Integrand
+from .integrator import Integrator, SampleBatch, TrainingStatus
+from .interface import (
+    MADNIS_INTEGRAND_FLAGS,
+    IntegrandDistribution,
+    IntegrandFunction,
+    build_madnis_integrand,
+)
+from .losses import (
+    kl_divergence,
+    multi_channel_loss,
+    rkl_divergence,
+    stratified_variance,
+    variance,
+)
+
+__all__ = [
+    "Integrator",
+    "TrainingStatus",
+    "SampleBatch",
+    "Integrand",
+    "Buffer",
+    "multi_channel_loss",
+    "stratified_variance",
+    "variance",
+    "kl_divergence",
+    "rkl_divergence",
+    "ChannelGroup",
+    "ChannelData",
+    "ChannelGrouping",
+    "Distribution",
+    "MADNIS_INTEGRAND_FLAGS",
+    "IntegrandDistribution",
+    "IntegrandFunction",
+    "build_madnis_integrand",
+]

madspace/madnis/buffer.py

@@ -0,0 +1,167 @@
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+
+
+class Buffer(nn.Module):
+    """
+    Circular buffer for multiple tensors with different shapes. The class is a torch.nn.Module to
+    allow for simple storage.
+    """
+
+    def __init__(
+        self,
+        capacity: int,
+        shapes: list[tuple[int, ...]],
+        persistent: bool = True,
+        dtypes: list[torch.dtype | None] | None = None,
+    ):
+        """
+        Args:
+            capacity: maximum number of samples stored in the buffer
+            shapes: shapes of the tensors to be stored, without batch dimension. If a shape is
+                None, no tensor is stored at that position. This allows for simpler handling of
+                optional stored fields.
+            persistent: if True, the content of the buffer is part of the module's state_dict
+            dtypes: if different from None, specifies the tensors which have a non-standard dtype
+        """
+        super().__init__()
+        self.keys = []
+        if dtypes is None:
+            dtypes = [None] * len(shapes)
+        for i, (shape, dtype) in enumerate(zip(shapes, dtypes)):
+            key = f"buffer{i}"
+            self.register_buffer(
+                key,
+                None if shape is None else torch.zeros((capacity, *shape), dtype=dtype),
+                persistent,
+            )
+            self.keys.append(key)
+        self.capacity = capacity
+        self.size = 0
+        self.store_index = 0
+
+    def _batch_slices(self, batch_size: int) -> slice:
+        """
+        Returns slices that split up the buffer into batches of at most ``batch_size``, respecting
+        the buffer size and periodic boundary.
+        """
+        start = self.store_index
+        while start < self.size:
+            stop = min(start + batch_size, self.size)
+            yield slice(start, stop)
+            start = stop
+        start = 0
+        while start < self.store_index:
+            stop = min(start + batch_size, self.store_index)
+            yield slice(start, stop)
+            start = stop
+
+    def _buffer_fields(self) -> torch.Tensor | None:
+        """
+        Iterates over the buffered tensors, without removing the padding if the buffer is not full.
+
+        Returns:
+            The buffered tensors, or None if a tensor was initialized with shape None
+        """
+        for key in self.keys:
+            yield getattr(self, key)
+
+    def __iter__(self) -> torch.Tensor | None:
+        """
+        Iterates over the buffered tensors
+
+        Returns:
+            The buffered tensors, or None if a tensor was initialized with shape None
+        """
+        for key in self.keys:
+            buffer = getattr(self, key)
+            yield None if buffer is None else buffer[: self.size]
+
+    def store(self, *tensors: torch.Tensor | None):
+        """
+        Adds the given tensors to the buffer. If the buffer is full, the oldest stored samples are
+        overwritten.
+
+        Args:
+            tensors: samples to be stored. The shapes of the tensors after the batch dimension must
+                match the shapes given during initialization. The argument can be None if the
+                corresponding shape was None during initialization.
+        """
+        store_slice1 = None
+        for buffer, data in zip(self._buffer_fields(), tensors):
+            if data is None:
+                continue
+            if store_slice1 is None:
+                size = min(data.shape[0], self.capacity)
+                end_index = self.store_index + size
+                if end_index < self.capacity:
+                    store_slice1 = slice(self.store_index, end_index)
+                    store_slice2 = slice(0, 0)
+                    load_slice1 = slice(0, size)
+                    load_slice2 = slice(0, 0)
+                else:
+                    store_slice1 = slice(self.store_index, self.capacity)
+                    store_slice2 = slice(0, end_index - self.capacity)
+                    load_slice1 = slice(0, self.capacity - self.store_index)
+                    load_slice2 = slice(self.capacity - self.store_index, size)
+                self.store_index = end_index % self.capacity
+                self.size = min(self.size + size, self.capacity)
+            buffer[store_slice1] = data[load_slice1]
+            buffer[store_slice2] = data[load_slice2]
+
+    def filter(
+        self,
+        predicate: Callable[[tuple[torch.Tensor | None, ...]], torch.Tensor],
+        batch_size: int = 100000,
+    ):
+        """
+        Removes samples from the buffer that do not fulfill the criterion given by the predicate
+        function.
+
+        Args:
+            predicate: function that returns a mask for a batch of samples, given a tuple with
+                all the buffered fields as argument
+            batch_size: maximal batch size to limit memory usage
+        """
+        masks = []
+        masked_size = 0
+        for batch_slice in self._batch_slices(batch_size):
+            mask = predicate(
+                tuple(
+                    None if t is None else t[batch_slice] for t in self._buffer_fields()
+                )
+            )
+            masked_size += torch.count_nonzero(mask)
+            masks.append(mask)
+        for buffer in self._buffer_fields():
+            if buffer is None:
+                continue
+            buffer[:masked_size] = torch.cat(
+                [
+                    buffer[batch_slice][mask]
+                    for batch_slice, mask in zip(self._batch_slices(batch_size), masks)
+                ],
+                dim=0,
+            )
+        self.size = masked_size
+        self.store_index = masked_size % self.capacity
+
+    def sample(self, count: int) -> list[torch.Tensor | None]:
+        """
+        Returns a batch of samples drawn from the buffer without replacement.
+
+        Args:
+            count: number of samples
+        Returns:
+            samples drawn from the buffer
+        """
+        weights = next(b for b in self._buffer_fields() if b is not None).new_ones(
+            self.size
+        )
+        indices = torch.multinomial(weights, min(count, self.size), replacement=False)
+        return [
+            None if buffer is None else buffer[indices]
+            for buffer in self._buffer_fields()
+        ]

madspace/madnis/channel_grouping.py

@@ -0,0 +1,85 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ChannelGroup:
+    """
+    A group of channels
+
+    Args:
+        group_index: index of the group in the list of groups
+        target_index: index of the channel that all other channels in the group are mapped to
+        channel_indices: indices of the channels in the group
+    """
+
+    group_index: int
+    target_index: int
+    channel_indices: list[int]
+
+
+@dataclass
+class ChannelData:
+    """
+    Information about a single channel
+
+    Args:
+        channel_index: index of the channel
+        target_index: index of the channel that it is mapped to
+        group: channel group that the channel belongs to
+        remapped: True if the channel is remapped to another channel
+        position_in_group: index of the channel within its group
+    """
+
+    channel_index: int
+    target_index: int
+    group: ChannelGroup
+    remapped: bool
+    position_in_group: int
+
+
+class ChannelGrouping:
+    """
+    Class that encodes how channels are grouped together for a multi-channel integrand
+    """
+
+    def __init__(self, channel_assignment: list[int | None]):
+        """
+        Args:
+            channel_assignment: list with an entry for each channel. If None, the channel is not
+                remapped. Otherwise, the index of the channel to which it is mapped.
+        """
+        group_dict = {}
+        for source_channel, target_channel in enumerate(channel_assignment):
+            if target_channel is None:
+                group_dict[source_channel] = ChannelGroup(
+                    group_index=len(group_dict),
+                    target_index=source_channel,
+                    channel_indices=[source_channel],
+                )
+
+        self.channels: list[ChannelData] = []
+        self.groups: list[ChannelGroup] = list(group_dict.values())
+
+        for source_channel, target_channel in enumerate(channel_assignment):
+            if target_channel is None:
+                self.channels.append(
+                    ChannelData(
+                        channel_index=source_channel,
+                        target_index=source_channel,
+                        group=group_dict[source_channel],
+                        remapped=False,
+                        position_in_group=0,
+                    )
+                )
+            else:
+                group = group_dict[target_channel]
+                self.channels.append(
+                    ChannelData(
+                        channel_index=source_channel,
+                        target_index=target_channel,
+                        group=group,
+                        remapped=True,
+                        position_in_group=len(group.channel_indices),
+                    )
+                )
+                group.channel_indices.append(source_channel)