madspace 0.3.1__cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

madspace/madnis/distribution.py

@@ -0,0 +1,103 @@
+import math
+from collections.abc import Callable
+from typing import Literal, Protocol
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+L2PI = -0.5 * math.log(2 * math.pi)
+
+Mapping = Callable[[torch.Tensor, bool], tuple[torch.Tensor, torch.Tensor]]
+
+
+class Distribution(Protocol):
+    """
+    Protocol for a (potentially learnable) distribution that can be used for sampling and
+    density estimation, like a normalizing flow.
+    """
+
+    def sample(
+        self,
+        n: int,
+        c: torch.Tensor | None = None,
+        channel: torch.Tensor | list[int] | int | None = None,
+        return_log_prob: bool = False,
+        return_prob: bool = False,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        """
+        Draws samples following the distribution
+
+        Args:
+            n: number of samples
+            c: condition, shape (n, dims_c) or None for an unconditional flow
+            channel: encodes the channel of the samples. It must have one of the following types:
+
+                - ``Tensor``: integer tensor of shape (n, ), containing the channel index for every
+                  input sample;
+                - ``list``: list of integers, specifying the number of samples in each channel;
+                - ``int``: integer specifying a single channel containing all the samples;
+                - ``None``: used in the single-channel case or to indicate that all channels contain
+                  the same number of samples in the multi-channel case.
+            return_log_prob: if True, also return the log-probabilities
+            return_prob: if True, also return the probabilities
+            device: device of the returned tensor. Only required if no condition is given.
+            dtype: dtype of the returned tensor. Only required if no condition is given.
+        Returns:
+            samples with shape (n, dims_in). Depending on the arguments ``return_log_prob``,
+            ``return_prob``, this function should also return the log-probabilities with shape (n, ),
+            the probabilities with shape (n, ).
+        """
+        ...
+
+    def log_prob(
+        self,
+        x: torch.Tensor,
+        c: torch.Tensor | None = None,
+        channel: torch.Tensor | list[int] | int | None = None,
+    ) -> torch.Tensor:
+        """
+        Computes the log-probabilities of the input data.
+
+        Args:
+            x: input data, shape (n, dims_in)
+            c: condition, shape (n, dims_c) or None for an unconditional flow
+            channel: encodes the channel of the samples. It must have one of the following types:
+
+                - ``Tensor``: integer tensor of shape (n, ), containing the channel index for every
+                  input sample;
+                - ``list``: list of integers, specifying the number of samples in each channel;
+                - ``int``: integer specifying a single channel containing all the samples;
+                - ``None``: used in the single-channel case or to indicate that all channels contain
+                  the same number of samples in the multi-channel case.
+        Returns:
+            log-probabilities with shape (n, )
+        """
+        return self.prob(x, c, channel).log()
+
+    def prob(
+        self,
+        x: torch.Tensor,
+        c: torch.Tensor | None = None,
+        channel: torch.Tensor | list[int] | int | None = None,
+    ) -> torch.Tensor:
+        """
+        Computes the probabilities of the input data.
+
+        Args:
+            x: input data, shape (n, dims_in)
+            c: condition, shape (n, dims_c) or None for an unconditional flow
+            channel: encodes the channel of the samples. It must have one of the following types:
+
+                - ``Tensor``: integer tensor of shape (n, ), containing the channel index for every
+                  input sample;
+                - ``list``: list of integers, specifying the number of samples in each channel;
+                - ``int``: integer specifying a single channel containing all the samples;
+                - ``None``: used in the single-channel case or to indicate that all channels contain
+                  the same number of samples in the multi-channel case.
+        Returns:
+            probabilities with shape (n, )
+        """
+        return self.log_prob(x, c, channel).exp()

madspace/madnis/integrand.py

@@ -0,0 +1,175 @@
+from collections.abc import Callable
+from typing import Literal
+
+import torch
+import torch.nn as nn
+
+from .channel_grouping import ChannelGrouping
+
+
+class Integrand(nn.Module):
+    """
+    Class that wraps an integrand function and meta-data necessary to use advanced MadNIS
+    features like learnable multi-channel weights, grouped channels and channel weight priors.
+    """
+
+    def __init__(
+        self,
+        function: Callable,
+        input_dim: int,
+        bounds: list[list[float]] | None = None,
+        channel_count: int | None = None,
+        remapped_dim: int | None = None,
+        has_channel_weight_prior: bool = False,
+        channel_grouping: ChannelGrouping | None = None,
+        function_includes_sampling: bool = False,
+        update_active_channels_mask: Callable[[torch.Tensor], None] | None = None,
+        discrete_dims: list[int] = [],
+        discrete_dims_position: Literal["first", "last"] = "first",
+        discrete_prior_prob_function: (
+            Callable[[torch.Tensor, int], torch.Tensor] | None
+        ) = None,
+    ):
+        """
+        Args:
+            function: integrand function.
+                The signature depends on the other arguments:
+
+                  - single-channel integration, ``channel_count=None``: ``x -> f``
+                  - basic multi-channel integration, ``remapped_dim=None``,
+                    ``has_channel_weight_prior=False``: ``(x, c) -> f``
+                  - with channel weights, ``remapped_dim=None``, ``has_channel_weight_prior=True``:
+                    ``(x, c) -> (f, alpha)`` (no trainable channel weights possible)
+                  - with channel-dependent mapping, ``remapped_dim: int``,
+                    ``has_channel_weight_prior=False``: ``(x, c) -> (f, y)``
+                  - all features, ``remapped_dim: int``, ``has_channel_weight_prior=True``:
+                    ``(x, c) -> (f, y, alpha)``
+
+                with the following tensors:
+
+                  - ``x`` is a point generated by the importance sampling, shape (n, input_dim),
+                  - ``c`` is the channel index, shape (n, ),
+                  - ``f`` is the integrand value, shape (n, ),
+                  - ``y`` is the point after applying a channel-dependent mapping, shape
+                    (n, remapped_dim)
+                  - ``alpha`` is the prior channel weight, shape (n, channel_count).
+            input_dim: dimension of the integration space
+            bounds: List of pairs ``[lower bound, upper bound]`` of the integration interval for
+                all dimensions. The integrand is rescaled so that the MadNIS training can be
+                performed on the unit hypercube. If None, the unit hypercube is used as integration
+                domain.
+            channel_count: None in the single-channel case, specifies the number of channels
+                otherwise.
+            remapped_dim: If different from None, it gives the dimension of a remapped space,
+                with a channel-dependent mapping computed as part of the integrand function.
+            has_channel_weight_prior: If True, the integrand returns channel weights
+            channel_grouping: ChannelGrouping object or None if all channels are independent
+        """
+        # TODO: update documentation
+        super().__init__()
+        self.input_dim = input_dim
+        self.remapped_dim = input_dim if remapped_dim is None else remapped_dim
+        self.channel_count = channel_count
+        self.has_channel_weight_prior = has_channel_weight_prior
+        self.channel_grouping = channel_grouping
+        self.function_includes_sampling = function_includes_sampling
+        self.update_active_channels_mask_func = update_active_channels_mask
+
+        self.discrete_dims = discrete_dims
+        self.discrete_dims_position = discrete_dims_position
+        self.discrete_prior_prob_function = discrete_prior_prob_function
+
+        if function_includes_sampling:
+            self.function = function
+        elif channel_count is None:
+            self.function = lambda x, channels: (function(x), None, None)
+        elif remapped_dim is None:
+            if has_channel_weight_prior:
+
+                def func(x, channels):
+                    w, prior = function(x, channels)
+                    return w, None, prior
+
+                self.function = func
+            else:
+                self.function = lambda x, channels: (function(x, channels), None, None)
+
+        elif has_channel_weight_prior:
+            self.function = function
+        else:
+
+            def func(x, channels):
+                w, y = function(x, channels)
+                return w, y, None
+
+            self.function = func
+
+        if bounds is not None:
+            bounds = torch.tensor(bounds)
+            self.register_buffer("scale", bounds[:, 1] - bounds[:, 0])
+            self.register_buffer("offset", bounds[:, 0])
+            self.register_buffer("scale_det", self.scale.prod())
+            old_func = self.function
+
+            def rescaled_func(x, channels):
+                w, y, prior = old_func(self.scale * x + self.offset, channels)
+                return self.scale_det * w, y, prior
+
+            self.function = rescaled_func
+
+        self.register_buffer(
+            "channel_id_map",
+            (
+                None
+                if self.channel_grouping is None
+                else torch.tensor(
+                    [
+                        channel.group.group_index
+                        for channel in self.channel_grouping.channels
+                    ]
+                )
+            ),
+        )
+
+    def unique_channel_count(self) -> int:
+        """
+        Returns the number of channels, or, if some channels are grouped together, the number of
+        channel groups
+        """
+        if self.channel_grouping is None:
+            return self.channel_count
+        else:
+            return len(self.channel_grouping.groups)
+
+    def remap_channels(self, channels: torch.Tensor | int) -> torch.Tensor | int:
+        """
+        Remaps channel indices to the indices of their respective channel groups if a
+        ``ChannelGrouping`` object was provided, otherwise returns the indices unchanged.
+
+        Args:
+            channels: channel indices, tensor with shape (n, ) or integer
+        Returns:
+            remapped channel indices, tensor with shape (n, ) or integer
+        """
+        if self.channel_grouping is None:
+            return channels
+        elif isinstance(channels, int):
+            return self.channel_id_map[channels].item()
+        else:
+            return self.channel_id_map[channels]
+
+    def update_active_channels_mask(self, mask: torch.Tensor) -> None:
+        if self.update_active_channels_mask_func is None:
+            return
+
+        full_mask = mask[
+            self.remap_channels(
+                torch.arange(len(self.channel_grouping.channels), device=mask.device)
+            )
+        ]
+        self.update_active_channels_mask_func(full_mask)
+
+    def forward(
+        self, x: torch.Tensor, channels: torch.Tensor | None
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        return self.function(x, channels)