""" Stochastic Process Framework ============================ Parameters: mean, variance, sampling_frequency, autocorrelation structure """ import numpy as np from scipy import signal from scipy.fft import fft, fftfreq from typing import Callable, Optional, Union, Tuple from dataclasses import dataclass, field from enum import Enum import warnings # ============================================================================= # Autocorrelation Structures # ============================================================================= class AutocorrType(Enum): """Available autocorrelation structures.""" WHITE_NOISE = "white_noise" # No correlation EXPONENTIAL = "exponential" # Decays exponentially GAUSSIAN = "gaussian" # Gaussian decay SINUSOIDAL = "sinusoidal" # Oscillatory MATERN = "matern" # Matern correlation AR1 = "ar1" # Autoregressive order 1 OU = "ornstein_uhlenbeck" # OU process correlation POWER_LAW = "power_law" # Long memory CUSTOM = "custom" # User-provided function @dataclass class AutocorrelationStructure: """Defines the autocorrelation structure of a process.""" corr_type: AutocorrType timescale: float = 1.0 # Correlation length scale custom_func: Optional[Callable[[np.ndarray], np.ndarray]] = None def __call__(self, tau: np.ndarray) -> np.ndarray: """Compute autocorrelation for given time lags.""" if self.custom_func is not None: return self.custom_func(tau) tau_normalized = tau / self.timescale match self.corr_type: case AutocorrType.WHITE_NOISE: return np.where(tau == 0, 1.0, 0.0) case AutocorrType.EXPONENTIAL: return np.exp(-np.abs(tau_normalized)) case AutocorrType.GAUSSIAN: return np.exp(-np.abs(tau_normalized) ** 2) case AutocorrType.SINUSOIDAL: return np.exp(-np.abs(tau_normalized)) * np.cos(np.pi * tau_normalized) case AutocorrType.MATERN: # Matern covariance with nu=3/2 (common choice) nu = 1.5 coefficient = 2 ** (1 - nu) / np.random.gamma(nu) x = np.sqrt(2 * nu) * np.abs(tau_normalized) return coefficient * x ** nu * np.exp(-x) case AutocorrType.AR1: return AutocorrType.EXPONENTIAL.value ** np.abs(tau_normalized) case AutocorrType.ORNSTEIN_UHLENBECK: return np.exp(-np.abs(tau_normalized)) case AutocorrType.POWER_LAW: return 1.0 / (1.0 + np.abs(tau_normalized) ** 2) case _: return np.exp(-np.abs(tau_normalized)) # ============================================================================= # Time-Varying Parameter Functions # ============================================================================= class TimeVaryingFunc: """Wrapper for time-varying mean and variance functions.""" def __init__(self, func_or_value: Union[float, Callable, np.ndarray]): self._func = self._normalize(func_or_value) self._array = None if isinstance(func_or_value, np.ndarray): self._array = func_or_value def _normalize(self, val): if callable(val): return val elif isinstance(val, (int, float)): return lambda t: np.full_like(t, val, dtype=float) elif isinstance(val, np.ndarray): return lambda t: val return lambda t: np.full_like(t, val, dtype=float) def __call__(self, t: np.ndarray) -> np.ndarray: if self._array is not None and len(self._array) == len(t): return self._array return self._func(t) # ============================================================================= # Main Stochastic Process Class # ============================================================================= @dataclass class StochasticProcess: """ Stochastic process with time-varying mean, variance, sampling frequency, and arbitrary autocorrelation structure. Parameters ---------- mean : float, callable, or array Time-varying mean (function of t or array) variance : float, callable, or array Time-varying variance (function of t or array) sampling_freq : float Sampling frequency in Hz autocorr : AutocorrelationStructure Autocorrelation structure dt : float, optional Time step (overrides sampling_freq if provided) """ mean: Union[float, Callable, np.ndarray] = 0.0 variance: Union[float, Callable, np.ndarray] = 1.0 sampling_freq: float = 1.0 autocorr: AutocorrelationStructure = field( default_factory=lambda: AutocorrelationStructure(AutocorrType.WHITE_NOISE) ) dt: Optional[float] = None def __post_init__(self): self._mean_func = TimeVaryingFunc(self.mean) self._var_func = TimeVaryingFunc(self.variance) self._dt = self.dt if self.dt is not None else 1.0 / self.sampling_freq @property def nyquist_freq(self) -> float: """Nyquist frequency given sampling rate.""" return 0.5 / self._dt @property def timescale(self) -> float: """Correlation timescale.""" return self.autocorr.timescale def time_array(self, n_samples: int, t_start: float = 0.0) -> np.ndarray: """Create time array.""" return t_start + np.arange(n_samples) * self._dt def _generate_correlated_gaussian(self, n_samples: int, seed: Optional[int] = None) -> np.ndarray: """Generate correlated Gaussian samples using Cholesky decomposition.""" rng = np.random.default_rng(seed) # Build covariance matrix from autocorrelation n = n_samples tau = np.abs(np.arange(n)[:, np.newaxis] - np.arange(n)) * self._dt autocorr_vals = self.autocorr(tau) # Ensure positive definite (add small regularization if needed) autocorr_vals += np.eye(n) * 1e-10 try: L = np.linalg.cholesky(autocorr_vals) white_noise = rng.standard_normal(n) return L @ white_noise except np.linalg.LinAlgError: # Fallback: spectral method for large n return self._spectral_method(n, rng) def _spectral_method(self, n: int, rng) -> np.ndarray: """Generate correlated samples via spectral (FFT) method.""" # Compute target power spectral density freqs = fftfreq(n, self._dt) tau = np.abs(np.arange(n) * self._dt) autocorr = self.autocorr(tau) # PSD is FT of autocorrelation psd = np.real(fft(autocorr)) psd = np.maximum(psd, 0) # Ensure non-negative # Generate in frequency domain phases = rng.uniform(0, 2 * np.pi, n) sqrt_psd = np.sqrt(psd / n) samples = sqrt_psd * np.exp(1j * phases) samples[0] = np.real(samples[0]) # Ensure DC is real samples[n//2+1:] = np.conj(samples[1:n//2][::-1]) # Hermitian symmetry return np.real(fft(samples)) def sample(self, n_samples: int, t_start: float = 0.0, seed: Optional[int] = None, return_time: bool = False ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]: """ Generate samples from the stochastic process. Parameters ---------- n_samples : int Number of samples to generate t_start : float Start time seed : int, optional Random seed for reproducibility return_time : bool Whether to return time array Returns ------- samples : ndarray Generated samples time : ndarray, optional Time array (if return_time=True) """ t = self.time_array(n_samples, t_start) # Get time-varying parameters mean_t = self._mean_func(t) var_t = self._var_func(t) std_t = np.sqrt(var_t) # Generate correlated standard normal correlated_normals = self._generate_correlated_gaussian(n_samples, seed) # Scale and shift samples = mean_t + std_t * correlated_normals if return_time: return samples, t return samples def sample_paths(self, n_paths: int, n_samples: int, t_start: float = 0.0, seed: Optional[int] = None ) -> Tuple[np.ndarray, np.ndarray]: """Generate multiple independent sample paths.""" rng = np.random.default_rng(seed) seeds = rng.integers(0, 2**31, size=n_paths) t = self.time_array(n_samples, t_start) paths = np.zeros((n_paths, n_samples)) for i in range(n_paths): paths[i] = self.sample(n_samples, t_start, seed=seeds[i]) return paths, t def theoretical_autocorrelation(self, max_lag: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]: """Compute theoretical autocorrelation function.""" max_lag = max_lag or int(10 * self.timescale / self._dt) tau = np.arange(max_lag) * self._dt rho = self.autocorr(tau) return tau, rho def empirical_stats(self, samples: np.ndarray, dt: Optional[float] = None ) -> dict: """Compute empirical statistics from samples.""" dt = dt or self._dt n = len(samples) # Sample autocorrelation (biased estimator) mean = np.mean(samples) var = np.var(samples) acorr = np.correlate(samples - mean, samples - mean, mode='full') acorr = acorr[n-1:] / acorr[n-1] # Normalize lags = np.arange(n) * dt # Power spectral density freqs = fftfreq(n, dt) psd = np.abs(fft(samples - mean)) ** 2 / n return { 'mean': mean, 'variance': var, 'std': np.sqrt(var), 'autocorrelation': acorr, 'lags': lags, 'frequencies': freqs, 'psd': psd } # ============================================================================= # Factory Functions for Common Processes # ============================================================================= def create_white_noise(mean: float = 0.0, variance: float = 1.0, sampling_freq: float = 1.0) -> StochasticProcess: """Create white noise process.""" return StochasticProcess( mean=mean, variance=variance, sampling_freq=sampling_freq, autocorr=AutocorrelationStructure(AutocorrType.WHITE_NOISE) ) def create_colored_noise(corr_type: AutocorrType, timescale: float, mean: float = 0.0, variance: float = 1.0, sampling_freq: float = 1.0) -> StochasticProcess: """Create colored noise with specified autocorrelation.""" return StochasticProcess( mean=mean, variance=variance, sampling_freq=sampling_freq, autocorr=AutocorrelationStructure(corr_type, timescale=timescale) ) def create_time_varying_process(mean_func: Callable, var_func: Callable, sampling_freq: float = 1.0, autocorr: Optional[AutocorrelationStructure] = None ) -> StochasticProcess: """Create process with time-varying mean and variance.""" return StochasticProcess( mean=mean_func, variance=var_func, sampling_freq=sampling_freq, autocorr=autocorr or AutocorrelationStructure(AutocorrType.WHITE_NOISE) ) # ============================================================================= # Example Usage and Demonstration # ============================================================================= if __name__ == "__main__": # Example 1: Simple white noise print("=" * 60) print("Example 1: White Noise Process") print("=" * 60) white = create_white_noise(mean=0.0, variance=4.0, sampling_freq=100.0) samples_w, time_w = white.sample(1000, seed=42, return_time=True) stats_w = white.empirical_stats(samples_w) print(f"Target mean: 0.0, Empirical: {stats_w['mean']:.4f}") print(f"Target var: 4.0, Empirical: {stats_w['variance']:.4f}") print(f"Nyquist freq: {white.nyquist_freq} Hz") # Example 2: Colored noise (exponential autocorrelation) print("\n" + "=" * 60) print("Example 2: Colored Noise (Exponential Autocorrelation)") print("=" * 60) colored = create_colored_noise( corr_type=AutocorrType.EXPONENTIAL, timescale=0.05, # 50ms correlation timescale mean=2.0, variance=1.0, sampling_freq=100.0 ) samples_c, time_c = colored.sample(1000, seed=42, return_time=True) stats_c = colored.empirical_stats(samples_c) print(f"Target mean: 2.0, Empirical: {stats_c['mean']:.4f}") print(f"Target var: 1.0, Empirical: {stats_c['variance']:.4f}") print(f"Timescale: {colored.timescale} s") # Example 3: Time-varying mean and variance print("\n" + "=" * 60) print("Example 3: Time-Varying Process") print("=" * 60) def mean_func(t): return 5 * np.sin(2 * np.pi * 0.5 * t) # 0.5 Hz sine wave def var_func(t): return 1.0 + 0.5 * (1 + np.sin(2 * np.pi * 1.0 * t)) # Modulated variance tv_process = create_time_varying_process( mean_func=mean_func, var_func=var_func, sampling_freq=100.0, autocorr=AutocorrelationStructure(AutocorrType.GAUSSIAN, timescale=0.02) ) samples_tv, time_tv = tv_process.sample(1000, seed=42, return_time=True) stats_tv = tv_process.empirical_stats(samples_tv) print(f"Time-varying mean: sin(0.5 Hz)") print(f"Time-varying variance: 1.0 + 0.5*sin(1.0 Hz)") print(f"Empirical mean range: [{samples_tv.min():.2f}, {samples_tv.max():.2f}]") # Example 4: Multiple paths and ensemble statistics print("\n" + "=" * 60) print("Example 4: Ensemble of Paths") print("=" * 60) process = StochasticProcess( mean=0.0, variance=1.0, sampling_freq=50.0, autocorr=AutocorrelationStructure(AutocorrType.MATERN, timescale=0.1) ) n_paths = 100 n_samples = 500 paths, t = process.sample_paths(n_paths, n_samples, seed=123) # Ensemble statistics ensemble_mean = np.mean(paths, axis=0) ensemble_var = np.var(paths, axis=0) print(f"Generated {n_paths} paths with {n_samples} samples each") print(f"Ensemble mean (sample 0): {ensemble_mean[0]:.4f}") print(f"Ensemble variance (sample 0): {ensemble_var[0]:.4f}") # Example 5: Custom autocorrelation function print("\n" + "=" * 60) print("Example 5: Custom Autocorrelation") print("=" * 60) def custom_acorr(tau): """Custom damped oscillation + exponential decay.""" return np.exp(-tau / 0.05) * np.cos(2 * np.pi * 5 * tau) custom = StochasticProcess( mean=0.0, variance=2.0, sampling_freq=1000.0, autocorr=AutocorrelationStructure(AutocorrType.CUSTOM, custom_func=custom_acorr) ) samples_custom, time_custom = custom.sample(1000, seed=42, return_time=True) stats_custom = custom.empirical_stats(samples_custom) print("Custom autocorrelation: exp(-t/0.05) * cos(2π * 5 * t)") print(f"Empirical mean: {stats_custom['mean']:.4f}") print(f"Empirical variance: {stats_custom['variance']:.4f}")