diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py index a321bac..250a86c 100644 --- a/src/scgenerator/noise.py +++ b/src/scgenerator/noise.py @@ -1,6 +1,5 @@ from __future__ import annotations -from dataclasses import dataclass, field from typing import Self, Sequence import numpy as np @@ -10,11 +9,28 @@ from scipy.integrate import cumulative_trapezoid from scgenerator import math, units -@dataclass class NoiseMeasurement: freq: np.ndarray psd: np.ndarray - rng: np.random.Generator = field(default_factory=np.random.default_rng) + rng: np.random.Generator + + def __init__( + self, freq: np.ndarray, psd: np.ndarray, rng: np.random.Generator | int | None = None + ): + if freq.ndim > 1: + raise TypeError(f"freq must be 1 dim, got {freq.shape}") + elif psd.shape[-1] != freq.shape[-1]: + raise TypeError( + "last axis of psd must be of same len as freq, " + f"got {psd.shape[-1]} and {freq.shape[-1]}" + ) + + if rng is None or isinstance(rng, int): + rng = np.random.default_rng(rng) + + self.freq = freq + self.psd = psd + self.rng = rng @classmethod def from_dBc(cls, freq: np.ndarray, psd_dBc: np.ndarray) -> NoiseMeasurement: @@ -194,12 +210,14 @@ class NoiseMeasurement: interp = np.interp(f, self.freq, self.psd, left=left, right=self.psd[-1]) return f, interp - def resampled(self, nf: int, dt: float | None = None, log_mode: bool = False) -> Self: - return self.__class__(*self.sample_spectrum(nf, dt, log_mode)) + def resampled( + self, nf: int, dt: float | None = None, log_mode: bool = False, left: float | None = None + ) -> Self: + return self.__class__(*self.sample_spectrum(nf, dt, log_mode, left)) def time_series( self, - nf: int, + nf: int | None = None, nseg: int = 1, dt: float | None = None, log_mode: bool = False, @@ -209,8 +227,9 @@ class NoiseMeasurement: Parameters ---------- - nf : int - number of frequency points to sample. Recommended is a power of 2 + 1 (129, 513, ...) + nf : int | None, optional + number of frequency points to sample. if None (default), will sample len(self.freq) + Recommended is a power of 2 + 1 (129, 513, ...) nseg : int | None, number of times to sample the spectrum, each time with a different random phase dt : float | None, optional @@ -218,6 +237,7 @@ class NoiseMeasurement: log_mode : bool, optional sample on a log-log scale rather than on a linear scale, by default False """ + nf = nf or len(self.freq) freq, amp = self.sample_spectrum(nf, dt, log_mode, left=0) fs = freq[-1] * 2