# Source code for botorch.sampling.qmc

#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# LICENSE file in the root directory of this source tree.

r"""
Quasi Monte-Carlo sampling from Normal distributions.

References:

.. [Pages2018numprob]
G. Pages. Numerical Probability: An Introduction with Applications to
Finance. Universitext. Springer International Publishing, 2018.
"""

from __future__ import annotations

import math
from typing import Optional

import torch
from torch import Tensor
from torch.quasirandom import SobolEngine

[docs]class NormalQMCEngine:
r"""Engine for qMC sampling from a Multivariate Normal N(0, I_d).

By default, this implementation uses Box-Muller transformed Sobol samples
following pg. 123 in [Pages2018numprob]_. To use the inverse transform
instead, set inv_transform=True.

Example:
>>> engine = NormalQMCEngine(3)
>>> samples = engine.draw(16)
"""

def __init__(
self, d: int, seed: Optional[int] = None, inv_transform: bool = False
) -> None:
r"""Engine for drawing qMC samples from a multivariate normal N(0, I_d).

Args:
d: The dimension of the samples.
seed: The seed with which to seed the random number generator of the
underlying SobolEngine.
inv_transform: If True, use inverse transform instead of Box-Muller.
"""
self._d = d
self._seed = seed
self._inv_transform = inv_transform
if inv_transform:
sobol_dim = d
else:
# to apply Box-Muller, we need an even number of dimensions
sobol_dim = 2 * math.ceil(d / 2)
self._sobol_engine = SobolEngine(dimension=sobol_dim, scramble=True, seed=seed)

[docs]    def draw(
self, n: int = 1, out: Optional[Tensor] = None, dtype: torch.dtype = torch.float
) -> Optional[Tensor]:
r"""Draw n qMC samples from the standard Normal.

Args:
n: The number of samples to draw. As a best practice, use powers of 2.
out: An option output tensor. If provided, draws are put into this
tensor, and the function returns None.
dtype: The desired torch data type (ignored if out is provided).

Returns:
A n x d tensor of samples if out=None and None otherwise.
"""
# get base samples
samples = self._sobol_engine.draw(n, dtype=dtype)
if self._inv_transform:
# apply inverse transform (values to close to 0/1 result in inf values)
v = 0.5 + (1 - torch.finfo(samples.dtype).eps) * (samples - 0.5)
samples_tf = torch.erfinv(2 * v - 1) * math.sqrt(2)
else:
# apply Box-Muller transform (note:  indexes starting from 1)
even = torch.arange(0, samples.shape[-1], 2)
Rs = (-2 * torch.log(samples[:, even])).sqrt()
thetas = 2 * math.pi * samples[:, 1 + even]
cos = torch.cos(thetas)
sin = torch.sin(thetas)
samples_tf = torch.stack([Rs * cos, Rs * sin], -1).reshape(n, -1)
# make sure we only return the number of dimension requested
samples_tf = samples_tf[:, : self._d]
if out is None:
return samples_tf
else:
out.copy_(samples_tf)

[docs]class MultivariateNormalQMCEngine:
r"""Engine for qMC sampling from a multivariate Normal N(\mu, \Sigma).

By default, this implementation uses Box-Muller transformed Sobol samples
following pg. 123 in [Pages2018numprob]_. To use the inverse transform
instead, set inv_transform=True.

Example:
>>> mean = torch.tensor([1.0, 2.0])
>>> cov = torch.tensor([[1.0, 0.25], [0.25, 2.0]])
>>> engine = MultivariateNormalQMCEngine(mean, cov)
>>> samples = engine.draw(16)
"""

def __init__(
self,
mean: Tensor,
cov: Tensor,
seed: Optional[int] = None,
inv_transform: bool = False,
) -> None:
r"""Engine for qMC sampling from a multivariate Normal N(\mu, \Sigma).

Args:
mean: The mean vector.
cov: The covariance matrix.
seed: The seed with which to seed the random number generator of the
underlying SobolEngine.
inv_transform: If True, use inverse transform instead of Box-Muller.
"""
# validate inputs
if not cov.shape == cov.shape:
raise ValueError("Covariance matrix is not square.")
if not mean.shape == cov.shape:
raise ValueError("Dimension mismatch between mean and covariance.")
if not torch.allclose(cov, cov.transpose(-1, -2)):
raise ValueError("Covariance matrix is not symmetric.")
self._mean = mean
self._normal_engine = NormalQMCEngine(
d=mean.shape, seed=seed, inv_transform=inv_transform
)
# compute Cholesky decomp; if it fails, do the eigendecomposition
try:
self._corr_matrix = torch.linalg.cholesky(cov).transpose(-1, -2)
except RuntimeError:
eigval, eigvec = torch.linalg.eigh(cov)
tol = 1e-8 if eigval.dtype == torch.double else 1e-6
if torch.any(eigval < -tol):
raise ValueError("Covariance matrix not PSD.")
eigval_root = eigval.clamp_min(0.0).sqrt()
self._corr_matrix = (eigvec * eigval_root).transpose(-1, -2)

[docs]    def draw(self, n: int = 1, out: Optional[Tensor] = None) -> Optional[Tensor]:
r"""Draw n qMC samples from the multivariate Normal.

Args:
n: The number of samples to draw. As a best practice, use powers of 2.
out: An option output tensor. If provided, draws are put into this
tensor, and the function returns None.

Returns:
A n x d tensor of samples if out=None and None otherwise.
"""
dtype = out.dtype if out is not None else self._mean.dtype
device = out.device if out is not None else self._mean.device
base_samples = self._normal_engine.draw(n, dtype=dtype).to(device=device)
corr_mat = self._corr_matrix.to(dtype=dtype, device=device)
mean = self._mean.to(dtype=dtype, device=device)
qmc_samples = base_samples @ corr_mat + mean
if out is None:
return qmc_samples
else:
out.copy_(qmc_samples)