"""
Noise models for stochastic variables.
"""
from __future__ import annotations
from typing import Protocol, Optional, Union
import numpy as np
from dataclasses import dataclass, field
class NoiseModel(Protocol):
"""
Protocol for noise models in the CausaLoop system.
Noise models add stochasticity to variable values while ensuring
reproducibility through explicit random number generators.
"""
def __call__(
self, rng: np.random.Generator, size: int = 1
) -> np.ndarray:
"""
Generate noise sample(s).
Parameters
----------
rng : np.random.Generator
NumPy random number generator for reproducibility.
Returns
-------
np.ndarray
Noise sample(s).
"""
...
[docs]
@dataclass
class GaussianNoise:
"""
Gaussian (normal) noise model.
Parameters
----------
mean : float, default=0.0
Mean of the Gaussian distribution.
std : float, default=1.0
Standard deviation of the Gaussian distribution.
Examples
--------
>>> from causaloop import GaussianNoise
>>> import numpy as np
>>> rng = np.random.default_rng(seed=42)
>>> noise = GaussianNoise(mean=0, std=0.1)
>>> noise(rng)
array([0.030...])
"""
mean: float = 0.0
std: float = 1.0
def __call__(
self, rng: np.random.Generator, size: int = 1
) -> np.ndarray:
"""
Generate Gaussian noise sample.
Parameters
----------
rng : np.random.Generator
Random number generator.
size : int
Sample size
Returns
-------
np.ndarray
Noise sample from N(mean, std^2).
"""
return rng.normal(self.mean, self.std, size)
[docs]
@dataclass
class LogNormalNoise:
"""
Log-normal noise model for positive-valued variables.
Parameters
----------
mean : float, default=0.0
Mean of the underlying normal distribution.
sigma : float, default=1.0
Standard deviation of the underlying normal distribution.
Examples
--------
>>> from causaloop import LogNormalNoise
>>> import numpy as np
>>> rng = np.random.default_rng(seed=42)
>>> noise = LogNormalNoise(mean=0, sigma=0.1)
>>> noise(rng)
array([1.03...])
"""
mean: float = 0.0
sigma: float = 1.0
def __call__(
self, rng: np.random.Generator, size: int = 1
) -> np.ndarray:
"""
Generate log-normal noise sample.
Parameters
----------
rng : np.random.Generator
Random number generator.
size : int
Sample size
Returns
-------
np.ndarray
Noise sample from Log-normal distribution.
"""
return rng.lognormal(self.mean, self.sigma, size)
[docs]
@dataclass
class PoissonNoise:
"""
Poisson noise for count variables.
Parameters
----------
lam : float, default=1.0
Expected value (lambda) of Poisson distribution.
Examples
--------
>>> from causaloop import PoissonNoise
>>> import numpy as np
>>> rng = np.random.default_rng(seed=42)
>>> noise = PoissonNoise(lam=5.0)
>>> noise(rng)
array([8])
"""
lam: float = 1.0
def __call__(
self, rng: np.random.Generator, size: int = 1
) -> np.ndarray:
"""
Generate Poisson noise sample.
Parameters
----------
rng : np.random.Generator
Random number generator.
size : int
Sample size
Returns
-------
np.ndarray
Noise sample from Poisson distribution.
"""
return rng.poisson(self.lam, size)
[docs]
@dataclass
class CompositeNoise:
"""
Composite noise model combining multiple noise sources.
Parameters
----------
models : list
List of noise models to combine.
weights : list, optional
Relative weights for each model. If None, equal weights.
Examples
--------
>>> from causaloop import GaussianNoise, UniformNoise, CompositeNoise
>>> import numpy as np
>>> rng = np.random.default_rng(seed=42)
>>> models = [
... GaussianNoise(mean=0, std=0.1),
... UniformNoise(low=-0.05, high=0.05)
... ]
>>> noise = CompositeNoise(models=models, weights=[0.7, 0.3])
>>> noise(rng, 3)
array([ 0.02..., -0.08..., 0.06...])
"""
models: list = field(default_factory=list)
weights: Optional[list] = None
def __call__(
self, rng: np.random.Generator, size: int = 1
) -> np.ndarray:
"""
Generate composite noise sample.
Parameters
----------
rng : np.random.Generator
Random number generator.
size : int
Sample size
Returns
-------
np.ndarray
Noise sample from Composite distribution.
"""
if not self.models:
return np.zeros(size)
if self.weights is None:
weights = [1.0] * len(self.models)
else:
weights = self.weights
# Normalize weights
total = sum(weights)
if total == 0:
return np.zeros(size)
samples = [model(rng, size) for model in self.models]
# Weighted sum
result = np.zeros(size)
for sample, weight in zip(samples, weights):
result += sample * (weight / total)
return result
[docs]
@dataclass
class ARNoise:
"""
Autoregressive noise model for temporal correlations.
Parameters
----------
alpha : float
Autoregressive coefficient (0 < alpha < 1).
sigma : float, default=1.0
Innovation standard deviation.
Examples
--------
>>> from causaloop import ARNoise
>>> import numpy as np
>>> rng = np.random.default_rng(seed=42)
>>> noise = ARNoise(alpha=0.8, sigma=0.1)
>>> # First sample
>>> noise(rng)
array([0.03...])
>>> # Second sample (correlated with first)
>>> noise(rng)
array([-0.07...])
"""
alpha: float
sigma: float = 1.0
_state: float = field(default=0.0, init=False)
_initialized: bool = field(default=False, init=False)
def __call__(
self, rng: np.random.Generator, size: int = 1
) -> np.ndarray:
"""
Generate autoregressive noise sample.
Parameters
----------
rng : np.random.Generator
Random number generator.
size : int
Sample size
Returns
-------
np.ndarray
Noise sample from AutoRegressive distribution.
"""
innovation = rng.normal(0, self.sigma, size)
if not self._initialized:
self._state = innovation
self._initialized = True
else:
self._state = self.alpha * self._state + innovation
return self._state
[docs]
def reset(self) -> None:
"""Reset AR process state."""
self._state = 0.0
self._initialized = False
[docs]
def create_noise(
noise_type: str = "gaussian",
**params
) -> NoiseModel:
"""
Create noise function.
Parameters
----------
noise_type : str, default="gaussian"
Type of noise: "gaussian", "uniform"
**params
Parameters for noise model.
Returns
-------
NoiseModel
Configured noise model.
Example
-------
>>> from causaloop import create_noise
>>> import numpy as np
>>> noise = create_noise(noise_type="gaussian", **{'mean': 0.0, 'std': 1.0})
>>> rng = np.random.default_rng(seed=42)
>>> noise(rng, 100).shape
(100,)
"""
if noise_type == "gaussian":
noise = GaussianNoise(**params)
elif noise_type == "uniform":
noise = UniformNoise(**params)
else:
raise ValueError(f"Unknown noise type: {noise_type}")
return noise