def draw_sobol_normal_samples(
d: int,
n: int,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
seed: Optional[int] = None,
) -> Tensor:
r"""Draw qMC samples from a multi-variate standard normal N(0, I_d)
A primary use-case for this functionality is to compute an QMC average
of f(X) over X where each element of X is drawn N(0, 1).
Args:
d: The dimension of the normal distribution.
n: The number of samples to return.
device: The torch device.
dtype: The torch dtype.
seed: The seed used for initializing Owen scrambling. If None (default),
use a random seed.
Returns:
A tensor of qMC standard normal samples with dimension `n x d` with device
and dtype specified by the input.
Example:
>>> samples = draw_sobol_normal_samples(2, 10)
"""
normal_qmc_engine = NormalQMCEngine(d=d, seed=seed, inv_transform=True)
samples = normal_qmc_engine.draw(n, dtype=torch.float if dtype is None else dtype)
return samples.to(device=device)
def __init__(self, ndim, use_sobol=False, use_inv=True, cache=False):
self.ndim = ndim
self.cache = cache
if use_sobol:
self.sampler = NormalQMCEngine(d=ndim, inv_transform=use_inv)
self.cached_points = torch.tensor([])
else:
self.sampler = None
def test_NormalQMCEngineInvTransform(self):
for d in (1, 2):
engine = NormalQMCEngine(d=d, inv_transform=True)
samples = engine.draw()
self.assertEqual(samples.dtype, torch.float)
self.assertEqual(samples.shape, torch.Size([1, d]))
samples = engine.draw(n=5)
self.assertEqual(samples.shape, torch.Size([5, d]))
# test double dtype
samples = engine.draw(dtype=torch.double)
self.assertEqual(samples.dtype, torch.double)
self.assertEqual(samples.shape, torch.Size([1, d]))
def test_NormalQMCEngineShapiroInvTransform(self):
engine = NormalQMCEngine(d=2, seed=12345, inv_transform=True)
samples = engine.draw(n=256)
self.assertEqual(samples.dtype, torch.float)
self.assertTrue(torch.all(torch.abs(samples.mean(dim=0)) < 1e-2))
self.assertTrue(torch.all(torch.abs(samples.std(dim=0) - 1) < 1e-2))
# perform Shapiro-Wilk test for normality
for i in (0, 1):
_, pval = shapiro(samples[:, i])
self.assertGreater(pval, 0.9)
# make sure samples are uncorrelated
cov = np.cov(samples.numpy().transpose())
self.assertLess(np.abs(cov[0, 1]), 1e-2)
class randn_sampler():
"""
Generates z~N(0,1) using random sampling or scrambled Sobol sequences.
Args:
ndim: (int)
The dimension of z.
use_sobol: (bool)
If True, sample z from scrambled Sobol sequence. Else, sample
from standard normal distribution.
Default: False
use_inv: (bool)
If True, use inverse CDF to transform z from U[0,1] to N(0,1).
Else, use Box-Muller transformation.
Default: True
cache: (bool)
If True, we cache some amount of Sobol points and reorder them.
This is mainly used for training GANs when we use two separate
Sobol generators which helps stabilize the training.
Default: False
Examples::
>>> sampler = randn_sampler(128, True)
>>> z = sampler.draw(10) # Generates [10, 128] vector
"""
def __init__(self, ndim, use_sobol=False, use_inv=True, cache=False):
self.ndim = ndim
self.cache = cache
if use_sobol:
self.sampler = NormalQMCEngine(d=ndim, inv_transform=use_inv)
self.cached_points = torch.tensor([])
else:
self.sampler = None
def draw(self, batch_size):
if self.sampler is None:
return torch.randn([batch_size, self.ndim])
else:
if self.cache:
if len(self.cached_points) < batch_size:
# sample from sampler and reorder the points
self.cached_points = self.sampler.draw(int(1e6))[torch.randperm(int(1e6))]
# Sample without replacement from cached points
samples = self.cached_points[:batch_size]
self.cached_points = self.cached_points[batch_size:]
return samples
else:
return self.sampler.draw(batch_size)
def test_NormalQMCEngineInvTransform(self):
# d = 1
engine = NormalQMCEngine(d=1, inv_transform=True)
samples = engine.draw()
self.assertEqual(samples.dtype, torch.float)
self.assertEqual(samples.shape, torch.Size([1, 1]))
samples = engine.draw(n=5)
self.assertEqual(samples.shape, torch.Size([5, 1]))
# d = 2
engine = NormalQMCEngine(d=2, inv_transform=True)
samples = engine.draw()
self.assertEqual(samples.shape, torch.Size([1, 2]))
samples = engine.draw(n=5)
self.assertEqual(samples.shape, torch.Size([5, 2]))
# test double dtype
samples = engine.draw(dtype=torch.double)
self.assertEqual(samples.dtype, torch.double)
def test_NormalQMCEngineSeededInvTransform(self):
# test even dimension
engine = NormalQMCEngine(d=2, seed=12345, inv_transform=True)
samples = engine.draw(n=2)
self.assertEqual(samples.dtype, torch.float)
self.assertEqual(samples.shape, torch.Size([2, 2]))
# test odd dimension
engine = NormalQMCEngine(d=3, seed=12345, inv_transform=True)
samples = engine.draw(n=2)
self.assertEqual(samples.shape, torch.Size([2, 3]))
def test_NormalQMCEngineSeededOut(self):
# test even dimension
engine = NormalQMCEngine(d=2, seed=12345)
out = torch.zeros(2, 2)
self.assertIsNone(engine.draw(n=2, out=out))
self.assertTrue(torch.all(out != 0))
# test odd dimension
engine = NormalQMCEngine(d=3, seed=12345)
out = torch.empty(2, 3)
self.assertIsNone(engine.draw(n=2, out=out))
self.assertTrue(torch.all(out != 0))
def test_NormalQMCEngineSeededInvTransform(self):
# test even dimension
engine = NormalQMCEngine(d=2, seed=12345, inv_transform=True)
samples = engine.draw(n=2)
self.assertEqual(samples.dtype, torch.float)
samples_expected = torch.tensor([[-0.41622922, 0.46622792],
[-0.96063897, -0.75568963]])
self.assertTrue(torch.allclose(samples, samples_expected))
# test odd dimension
engine = NormalQMCEngine(d=3, seed=12345, inv_transform=True)
samples = engine.draw(n=2)
samples_expected = torch.tensor([
[-1.40525266, 1.37652443, -0.8519666],
[-0.166497, -2.3153681, -0.15975676],
])
self.assertTrue(torch.allclose(samples, samples_expected))
def test_NormalQMCEngineSeeded(self):
# test even dimension
engine = NormalQMCEngine(d=2, seed=12345)
samples = engine.draw(n=2)
self.assertEqual(samples.dtype, torch.float)
samples_expected = torch.tensor([[-0.63099602, -1.32950772],
[0.29625805, 1.86425618]])
self.assertTrue(torch.allclose(samples, samples_expected))
# test odd dimension
engine = NormalQMCEngine(d=3, seed=12345)
samples = engine.draw(n=2)
samples_expected = torch.tensor([
[1.83169884, -1.40473647, 0.24334828],
[0.36596099, 1.2987395, -1.47556275],
])
self.assertTrue(torch.allclose(samples, samples_expected))
def test_NormalQMCEngineSeededOut(self):
# test even dimension
engine = NormalQMCEngine(d=2, seed=12345)
out = torch.empty(2, 2)
self.assertIsNone(engine.draw(n=2, out=out))
samples_expected = torch.tensor([[-0.63099602, -1.32950772],
[0.29625805, 1.86425618]])
self.assertTrue(torch.allclose(out, samples_expected))
# test odd dimension
engine = NormalQMCEngine(d=3, seed=12345)
out = torch.empty(2, 3)
self.assertIsNone(engine.draw(n=2, out=out))
samples_expected = torch.tensor([
[1.83169884, -1.40473647, 0.24334828],
[0.36596099, 1.2987395, -1.47556275],
])
self.assertTrue(torch.allclose(out, samples_expected))