distributions.py

import numpy as np


class Distribution:

    def sample(self, size: int | tuple) -> np.ndarray:
        raise NotImplementedError


class Gaussian(Distribution):

    def __init__(self, mean: np.ndarray, cov: np.ndarray, rng: np.random.Generator):
        self.mean = mean
        self.cov = cov
        self.rng = rng

    def sample(self, size: int) -> np.ndarray:
        return self.rng.multivariate_normal(self.mean, self.cov, size=size)


class EasyDistribution(Distribution):

    def __init__(self, dim) -> None:
        self.dim = dim

    def sample(self, size: int | tuple) -> np.ndarray:
        if type(size) == tuple:
            samples = []
            for _ in range(size[0]):
                sample = []
                for j in range(3):
                    if j == 0:
                        val = np.random.uniform(0, 1, size=self.dim)
                    elif j == 1:
                        val = np.random.uniform(1, 2, size=self.dim)
                    else:
                        val = np.random.uniform(10, 11, size=self.dim)
                    sample.append(val)
                samples.append(sample)
        else:
            raise NotImplementedError

        samples = np.array(samples)
        return samples


class ThreeGaussian(Distribution):

    def __init__(self, rng) -> None:
        self.rng = rng
        self.weights = [0.3, 0.4, 0.3]
        self.means = np.array([[0, 0], [1, 1], [2, 0]])
        self.covs = [np.eye(2) for _ in range(3)]

    def sample(self, size) -> np.ndarray:

        n_components = self.means.shape[0]
        samples = []

        for _ in range(size[0]):
            sample = []
            for j in range(3):
                component = self.rng.choice(n_components, p=self.weights)
                sample.append(
                    self.rng.multivariate_normal(
                        self.means[component], self.covs[component]
                    )
                )
            samples.append(sample)

        return np.array(samples)