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, _: float, x: ArrayLike) -> ArrayLike:
        if np.isscalar(x):
            noise = np.random.normal(0, 1)
            x += self.__sigma * noise + self.__mu
        else:
            dim = len(x)

            if np.isscalar(self.__sigma):
                sigma = np.identity(dim) * self.__sigma
            else:
                sigma = self.__sigma
                
            if np.isscalar(self.__mu):
                mu = np.ones(dim)
            else:
                mu = self.__mu

            noise = np.random.normal(0, 1, np.shape(x))
            x += sigma @ noise + mu

        return x


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

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