SPATZ/spatz/transforms/noise.py

import numpy as np

from numpy.typing import ArrayLike
from typing import Any, Tuple

from spatz.transforms import Transform


class GaussianNoise(Transform):
    def __init__(self, mu: ArrayLike, sigma: ArrayLike) -> None:
        super().__init__()

        self.__mu = mu
        self.__sigma = sigma

    def __call__(self, t: float, x: ArrayLike) -> ArrayLike:
        assert np.shape(self.__mu) == np.shape(x), "Mu and x have to match in shape."

        if np.isscalar(x):
            noise = np.random.normal(0, 1)
            x += self.__sigma * noise + self.__mu
        else:
            noise = np.random.normal(0, 1, np.shape(x))
            x += self.__sigma @ noise + self.__mu

        return x


class PinkNoise(Transform):
    def __init__(self) -> None:
        super().__init__()

    def __call__(self, t: float, x: ArrayLike) -> Any:
        pass