"""Parameter samplers for Monte Carlo simulations."""
import abc
from typing import Any, Sequence, cast
import numpy as np
import torch
from scimba_torch.domain.meshless_domain.domain_1d import Segment1D
from scimba_torch.domain.meshless_domain.domain_2d import Circle2D, Square2D
from scimba_torch.domain.meshless_domain.domain_3d import Cube3D
from scimba_torch.utils.scimba_tensors import LabelTensor
PARAM_TYPE = int | float
def _is_param_type(arg: Any) -> bool:
return isinstance(arg, int) or isinstance(arg, float)
def _is_param_domain_type(arg: Any) -> bool:
return isinstance(arg, Sequence) and all(
isinstance(a, Sequence)
and (len(a) == 2)
and _is_param_type(a[0])
and _is_param_type(a[0])
for a in arg
)
def _check_and_cast_argument(bounds: Any) -> tuple[bool, list[tuple[float, float]]]:
resT = []
resB = _is_param_domain_type(bounds)
if resB:
arg = cast(Sequence, bounds)
for a in arg:
resT.append((float(a[0]), float(a[1])))
return resB, resT
[docs]
class ParametricSampler(abc.ABC):
"""Abstract base class for parametric samplers.
Args:
bounds: A list of tuples, containing the lower and upper bounds.
**kwargs: Keyword arguments including:
- pre-sampling: whether to pre-sample points before training;
- n_pre_sampled: the number of pre-sampled points to generate if so.
Raises:
TypeError:
- If parameters domain is not a list of tuples of two floats.
- If n_pre_sampled is not an integer.
ValueError: If any bound has lower value greater than upper value.
"""
DEFAULT_N_PRE_SAMPLED_POINTS = 100_000
def __init__(self, bounds: list[tuple[float, float]], **kwargs):
check, nbounds = _check_and_cast_argument(bounds)
if not check:
raise TypeError("parameters domain must be a list of tuples of two floats")
if not all(bound[0] <= bound[1] for bound in nbounds):
raise ValueError(
f"cannot create a sampler for the empty parameter domain {str(nbounds)}"
)
self.set_new_bounds(nbounds)
# handle pre-sampling options
self.pre_sampling = kwargs.get("pre_sampling", False)
self.n_pre_sampled_points = kwargs.get(
"n_pre_sampled_points", self.DEFAULT_N_PRE_SAMPLED_POINTS
)
# check whether to do random sampling with numpy (for debug)
self.numpy = kwargs.get("numpy", False)
# Perform pre-sampling if enabled
if self.pre_sampling:
self.pts = self.sample_new_points(self.n_pre_sampled_points)
self.sample = self.sample_from_pre_sampled_points
else:
self.sample = self.sample_new_points
[docs]
def set_new_bounds(self, nbounds: list[tuple[float, float]]) -> None:
"""Updates the bounds of the parameters in the approximation space.
Args:
nbounds: A list of tuples containing the new bounds for each parameter.
"""
# bounds is a list of tuples, where
# each tuple contains the lower and upper bounds for each dimension.
self.bounds = nbounds
# dim is the number of dimensions, inferred from the length of bounds.
self.dim = len(self.bounds)
# lower and upper are tensors containing the lower and upper bounds.
if self.dim > 0:
self.lower = torch.tensor([self.bounds[i][0:1] for i in range(self.dim)]).T
self.upper = torch.tensor([self.bounds[i][1:2] for i in range(self.dim)]).T
[docs]
@abc.abstractmethod
def sample_new_points(self, n: int) -> LabelTensor:
"""Generates samples within the specified bounds for each dimension.
Args:
n: The number of samples to generate.
Returns:
A tensor containing the generated samples and corresponding labels.
"""
[docs]
def sample_from_pre_sampled_points(self, n: int) -> LabelTensor:
"""Samples from the pre-sampled points.
Args:
n: The number of samples to generate.
Returns:
A tensor containing the generated samples and corresponding labels.
"""
self.check_sample_size(n)
limit = self.n_pre_sampled_points
if n <= limit:
indices = torch.randperm(limit)[:n]
else:
indices = torch.randint(0, limit, (n,))
return self.pts[indices]
[docs]
def check_sample_size(self, n: int):
"""Checks if the sample size is a non-negative integer.
Args:
n: The number of samples to generate.
Raises:
TypeError: If argument is not an integer.
ValueError: If argument is negative.
"""
if not isinstance(n, int):
raise TypeError("argument to sample method must be an integer")
if n < 0:
raise ValueError("argument to sample method must be non-negative")
[docs]
def rescale_samples(self, samples: torch.Tensor) -> LabelTensor:
r"""Rescales the samples to fit within the specified bounds.
This function assumes that the samples are generated in the unit hypercube
:math:`[0, 1]^{\text{dim}}`
and rescales them to fit within the bounds specified for each dimension.
It also:
- make the samples require gradients;
- converts the samples to a LabelTensor with zero labels.
Args:
samples: A tensor containing the generated samples.
Returns:
A tensor containing the rescaled samples.
"""
if self.dim > 0:
samples = samples * (self.upper - self.lower) + self.lower
samples.requires_grad_()
return LabelTensor(samples)
[docs]
class SobolParametricSampler(ParametricSampler):
"""Sample from a Sobol sequence within given bounds for each dimension.
Args:
bounds: A list of tuples, containing the lower and upper bounds.
**kwargs: Keyword arguments including:
- pre-sampling: whether to pre-sample points before training;
- n_pre_sampled: the number of pre-sampled points to generate if so.
Raises:
TypeError: If bounds is not a list of tuples of two floats.
"""
def __init__(self, bounds: list[tuple[float, float]], **kwargs):
check, nbounds = _check_and_cast_argument(bounds)
if not check:
raise TypeError("parameters domain must be a list of tuples of two floats")
dim = len(nbounds)
# first, define the sampler
if dim == 0:
class ZeroDimensionalSampler:
"""Sampler for zero-dimensional parameter space."""
def draw(self, n: int) -> torch.Tensor:
"""Generates samples for zero-dimensional parameter space.
Args:
n: The number of samples to generate.
Returns:
A tensor of shape (n, 0) containing the generated samples.
"""
return torch.rand(n, 0)
self.sampler = ZeroDimensionalSampler()
else:
self.sampler = torch.quasirandom.SobolEngine(dimension=dim, scramble=True)
self.device = torch.get_default_device()
# then, call the parent constructor to handle bounds and pre-sampling options
super().__init__(bounds, **kwargs)
[docs]
def sample_new_points(self, n: int) -> LabelTensor:
"""Generates samples from a Sobol sequence within the specified bounds.
Args:
n: The number of samples to generate.
Returns:
A tensor containing the generated samples and corresponding labels.
"""
self.check_sample_size(n)
if n == 0:
samples = torch.zeros((0, self.dim))
else:
samples = self.sampler.draw(n).to(self.device)
return self.rescale_samples(samples)