diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py index 250a86c..0878254 100644 --- a/src/scgenerator/noise.py +++ b/src/scgenerator/noise.py @@ -12,11 +12,8 @@ from scgenerator import math, units class NoiseMeasurement: freq: np.ndarray psd: np.ndarray - rng: np.random.Generator - def __init__( - self, freq: np.ndarray, psd: np.ndarray, rng: np.random.Generator | int | None = None - ): + def __init__(self, freq: np.ndarray, psd: np.ndarray): if freq.ndim > 1: raise TypeError(f"freq must be 1 dim, got {freq.shape}") elif psd.shape[-1] != freq.shape[-1]: @@ -25,12 +22,8 @@ class NoiseMeasurement: 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: @@ -221,6 +214,7 @@ class NoiseMeasurement: nseg: int = 1, dt: float | None = None, log_mode: bool = False, + rng: np.random.Generator | int | None = None, ) -> tuple[np.ndarray, np.ndarray]: """ computes a pulse train whose psd matches the measurement @@ -238,6 +232,8 @@ class NoiseMeasurement: sample on a log-log scale rather than on a linear scale, by default False """ nf = nf or len(self.freq) + if rng is None or isinstance(rng, int): + rng = np.random.default_rng(rng) freq, amp = self.sample_spectrum(nf, dt, log_mode, left=0) fs = freq[-1] * 2 @@ -254,11 +250,11 @@ class NoiseMeasurement: if nseg > 1: amp = np.tile(amp, (nseg, 1)).T - amp[1:right] *= np.exp(2j * np.pi * self.rng.random((phase_n, nseg))) + amp[1:right] *= np.exp(2j * np.pi * rng.random((phase_n, nseg))) time, signal = ss.istft(amp, fs, nperseg=(nf - 1) * 2, noverlap=nf - 1, scaling="psd") signal *= np.sqrt(2) else: - amp[1:right] *= np.exp(2j * np.pi * self.rng.random(phase_n)) + amp[1:right] *= np.exp(2j * np.pi * rng.random(phase_n)) signal = np.fft.irfft(amp) * np.sqrt((amp.size - 1) * 2 * fs) time = np.arange(len(signal)) / fs return time, signal