Source code for romtools.workflows.sampling_methods

# pylint: disable=invalid-name

from typing import Protocol
import numpy as np
from scipy.stats import qmc




class Sampler(Protocol):
    '''
    Generate UIID samples

    Returns np.ndarray of shape (number_of_samples, dimensionality)
    '''

    def __call__(self, number_of_samples: int, dimensionality: int = 1, seed=None) -> np.ndarray:
        pass



##########################################
# Sampling Methods
##########################################




def MonteCarloSampler(number_of_samples: int, dimensionality: int = 1, seed=None) -> np.ndarray:
    '''
    Generate UIID Monte Carlo samples

    Conforms to the Sampler protocol
    '''
    if seed is not None:
        np.random.seed(seed)
    return np.random.uniform(size=(number_of_samples, dimensionality))





def LatinHypercubeSampler(number_of_samples: int, dimensionality: int = 1, seed=None) -> np.ndarray:
    '''
    Generate UIID LHS samples

    Conforms to the Sampler protocol
    '''
    sampler = qmc.LatinHypercube(dimensionality, seed=seed)
    return sampler.random(n=number_of_samples)





def RandomizedQuasiMonteCarloSampler(number_of_samples: int,
                                     dimensionality: int = 1,
                                     seed=None) -> np.ndarray:
    '''
    Generate randomized quasi-Monte Carlo samples using Owen-scrambled Sobol.

    Randomization preserves low-discrepancy structure while ensuring unbiased
    integral estimation in expectation over the scrambling.

    Conforms to the Sampler protocol.
    '''
    sampler = qmc.Sobol(d=dimensionality, scramble=True, seed=seed)
    return sampler.random(n=number_of_samples)