Esempio n. 1
0
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)
Esempio n. 2
0
 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
Esempio n. 3
0
 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]))
Esempio n. 4
0
 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)
Esempio n. 5
0
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)
Esempio n. 6
0
 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)
Esempio n. 7
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)
     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]))
Esempio n. 8
0
 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))
Esempio n. 9
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))
Esempio n. 10
0
 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))
Esempio n. 11
0
 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))