Spaces:
Sleeping
Sleeping
import torch | |
from typing import Optional | |
class SobolEngine: | |
r""" | |
The :class:`torch.quasirandom.SobolEngine` is an engine for generating | |
(scrambled) Sobol sequences. Sobol sequences are an example of low | |
discrepancy quasi-random sequences. | |
This implementation of an engine for Sobol sequences is capable of | |
sampling sequences up to a maximum dimension of 21201. It uses direction | |
numbers from https://web.maths.unsw.edu.au/~fkuo/sobol/ obtained using the | |
search criterion D(6) up to the dimension 21201. This is the recommended | |
choice by the authors. | |
References: | |
- Art B. Owen. Scrambling Sobol and Niederreiter-Xing points. | |
Journal of Complexity, 14(4):466-489, December 1998. | |
- I. M. Sobol. The distribution of points in a cube and the accurate | |
evaluation of integrals. | |
Zh. Vychisl. Mat. i Mat. Phys., 7:784-802, 1967. | |
Args: | |
dimension (Int): The dimensionality of the sequence to be drawn | |
scramble (bool, optional): Setting this to ``True`` will produce | |
scrambled Sobol sequences. Scrambling is | |
capable of producing better Sobol | |
sequences. Default: ``False``. | |
seed (Int, optional): This is the seed for the scrambling. The seed | |
of the random number generator is set to this, | |
if specified. Otherwise, it uses a random seed. | |
Default: ``None`` | |
Examples:: | |
>>> # xdoctest: +SKIP("unseeded random state") | |
>>> soboleng = torch.quasirandom.SobolEngine(dimension=5) | |
>>> soboleng.draw(3) | |
tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000], | |
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000], | |
[0.7500, 0.2500, 0.2500, 0.2500, 0.7500]]) | |
""" | |
MAXBIT = 30 | |
MAXDIM = 21201 | |
def __init__(self, dimension, scramble=False, seed=None): | |
if dimension > self.MAXDIM or dimension < 1: | |
raise ValueError("Supported range of dimensionality " | |
f"for SobolEngine is [1, {self.MAXDIM}]") | |
self.seed = seed | |
self.scramble = scramble | |
self.dimension = dimension | |
cpu = torch.device("cpu") | |
self.sobolstate = torch.zeros(dimension, self.MAXBIT, device=cpu, dtype=torch.long) | |
torch._sobol_engine_initialize_state_(self.sobolstate, self.dimension) | |
if not self.scramble: | |
self.shift = torch.zeros(self.dimension, device=cpu, dtype=torch.long) | |
else: | |
self._scramble() | |
self.quasi = self.shift.clone(memory_format=torch.contiguous_format) | |
self._first_point = (self.quasi / 2 ** self.MAXBIT).reshape(1, -1) | |
self.num_generated = 0 | |
def draw(self, n: int = 1, out: Optional[torch.Tensor] = None, | |
dtype: torch.dtype = torch.float32) -> torch.Tensor: | |
r""" | |
Function to draw a sequence of :attr:`n` points from a Sobol sequence. | |
Note that the samples are dependent on the previous samples. The size | |
of the result is :math:`(n, dimension)`. | |
Args: | |
n (Int, optional): The length of sequence of points to draw. | |
Default: 1 | |
out (Tensor, optional): The output tensor | |
dtype (:class:`torch.dtype`, optional): the desired data type of the | |
returned tensor. | |
Default: ``torch.float32`` | |
""" | |
if self.num_generated == 0: | |
if n == 1: | |
result = self._first_point.to(dtype) | |
else: | |
result, self.quasi = torch._sobol_engine_draw( | |
self.quasi, n - 1, self.sobolstate, self.dimension, self.num_generated, dtype=dtype, | |
) | |
result = torch.cat((self._first_point, result), dim=-2) | |
else: | |
result, self.quasi = torch._sobol_engine_draw( | |
self.quasi, n, self.sobolstate, self.dimension, self.num_generated - 1, dtype=dtype, | |
) | |
self.num_generated += n | |
if out is not None: | |
out.resize_as_(result).copy_(result) | |
return out | |
return result | |
def draw_base2(self, m: int, out: Optional[torch.Tensor] = None, | |
dtype: torch.dtype = torch.float32) -> torch.Tensor: | |
r""" | |
Function to draw a sequence of :attr:`2**m` points from a Sobol sequence. | |
Note that the samples are dependent on the previous samples. The size | |
of the result is :math:`(2**m, dimension)`. | |
Args: | |
m (Int): The (base2) exponent of the number of points to draw. | |
out (Tensor, optional): The output tensor | |
dtype (:class:`torch.dtype`, optional): the desired data type of the | |
returned tensor. | |
Default: ``torch.float32`` | |
""" | |
n = 2 ** m | |
total_n = self.num_generated + n | |
if not (total_n & (total_n - 1) == 0): | |
raise ValueError("The balance properties of Sobol' points require " | |
f"n to be a power of 2. {self.num_generated} points have been " | |
f"previously generated, then: n={self.num_generated}+2**{m}={total_n}. " | |
"If you still want to do this, please use " | |
"'SobolEngine.draw()' instead." | |
) | |
return self.draw(n=n, out=out, dtype=dtype) | |
def reset(self): | |
r""" | |
Function to reset the ``SobolEngine`` to base state. | |
""" | |
self.quasi.copy_(self.shift) | |
self.num_generated = 0 | |
return self | |
def fast_forward(self, n): | |
r""" | |
Function to fast-forward the state of the ``SobolEngine`` by | |
:attr:`n` steps. This is equivalent to drawing :attr:`n` samples | |
without using the samples. | |
Args: | |
n (Int): The number of steps to fast-forward by. | |
""" | |
if self.num_generated == 0: | |
torch._sobol_engine_ff_(self.quasi, n - 1, self.sobolstate, self.dimension, self.num_generated) | |
else: | |
torch._sobol_engine_ff_(self.quasi, n, self.sobolstate, self.dimension, self.num_generated - 1) | |
self.num_generated += n | |
return self | |
def _scramble(self): | |
g: Optional[torch.Generator] = None | |
if self.seed is not None: | |
g = torch.Generator() | |
g.manual_seed(self.seed) | |
cpu = torch.device("cpu") | |
# Generate shift vector | |
shift_ints = torch.randint(2, (self.dimension, self.MAXBIT), device=cpu, generator=g) | |
self.shift = torch.mv(shift_ints, torch.pow(2, torch.arange(0, self.MAXBIT, device=cpu))) | |
# Generate lower triangular matrices (stacked across dimensions) | |
ltm_dims = (self.dimension, self.MAXBIT, self.MAXBIT) | |
ltm = torch.randint(2, ltm_dims, device=cpu, generator=g).tril() | |
torch._sobol_engine_scramble_(self.sobolstate, ltm, self.dimension) | |
def __repr__(self): | |
fmt_string = [f'dimension={self.dimension}'] | |
if self.scramble: | |
fmt_string += ['scramble=True'] | |
if self.seed is not None: | |
fmt_string += [f'seed={self.seed}'] | |
return self.__class__.__name__ + '(' + ', '.join(fmt_string) + ')' | |