Function get_samples

plib/sample_utils.py:46–89

Uniformly distribute `total_samples` within [0, 1)^{d}. Args: total_samples: total number of samples d: number of dimension method: 'random': 'LatinHypercube': fast_forward: number of samples alread

(
        total_samples: int,
        d: int,
        method: str = 'random',
        rng=None,
        fast_forward_n: int = None,
        shuffle: bool = False,
)

def get_samples(
        total_samples: int,
        d: int,
        method: str = 'random',
        rng=None,
        fast_forward_n: int = None,
        shuffle: bool = False,
) -> np.ndarray:
    """
    Uniformly distribute `total_samples` within [0, 1)^{d}.

    Args:
        total_samples:
            total number of samples
        d:
            number of dimension
        method:
            'random':
            'LatinHypercube':
        fast_forward:
            number of samples already generated
        shuffle:
            whether to shuffle the samples (along total_samples).
            It is needed when combining the samples with other samples.

    Returns:
        (total_samples, d)
    """

    if method == 'random':
        if rng is None:
            return np.random.rand(total_samples, d)
        else:
            return rng.rand(total_samples, d)
    elif method.lower() == 'LatinHypercube'.lower():
        sampler = qmc.LatinHypercube(d=d, seed=rng)
        if fast_forward_n is not None and fast_forward_n > 0:
            sampler = sampler.fast_forward(fast_forward_n)
        samples = sampler.random(n=total_samples)  # (total_samples, d)
        if shuffle:
            samples = shuffle_along_axis(arr=samples, axis=0, rng=rng)
        return samples
    else:
        raise NotImplementedError


def shuffle_along_axis(arr: np.ndarray, axis: int, rng=None) -> np.ndarray: