From dfbe47cf523b930ec5d2f8c88ba6ddac4223aaf8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= Date: Fri, 27 Oct 2023 11:14:18 +0200 Subject: [PATCH] improved noise --- src/scgenerator/math.py | 2 +- src/scgenerator/noise.py | 40 ++++++++++++++++++++++------------ tests/test_noise.py | 46 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 73 insertions(+), 15 deletions(-) diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index f91338c..3c5e14f 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -774,4 +774,4 @@ def mean_angle(values: np.ndarray, axis: int = 0): values, np.abs(values), out=np.zeros(values.shape, dtype=values.dtype), where=values != 0 ) total_phase = np.sum(values, axis=axis) - return (total_phase) / np.abs(total_phase) + return np.divide(total_phase, np.abs(total_phase), where=total_phase != 0, out=total_phase) diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py index 0e754b2..8da2eba 100644 --- a/src/scgenerator/noise.py +++ b/src/scgenerator/noise.py @@ -70,9 +70,11 @@ class NoiseMeasurement: take out the DC component (0-frequency) of each segement after segmentation """ signal = np.asanyarray(signal) + if signal.ndim > 1: + raise ValueError( + f"got signal of shape {signal.shape}. Only one 1D signals are supported" + ) signal_segments = segments(signal, num_segments) - if force_no_dc: - signal_segments = (signal_segments.T - signal_segments.mean(axis=1)).T n = signal_segments.shape[-1] try: window_arr = cls._window_functions[window](n) @@ -84,10 +86,14 @@ class NoiseMeasurement: window_correction = np.sum(window_arr**2) / n signal_segments = signal_segments * window_arr + if force_no_dc: + signal_segments = (signal_segments.T - signal_segments.mean(axis=1)).T + freq = np.fft.rfftfreq(n, dt) - psd = np.fft.rfft(signal_segments) * np.sqrt(2 * dt / n) + psd = np.fft.rfft(signal_segments) * np.sqrt(dt / n) phase = math.mean_angle(psd) psd = psd.real**2 + psd.imag**2 + psd[..., 1:-1] *= 2 return cls(freq, psd.mean(axis=0) / window_correction, phase=phase) def plottable( @@ -138,7 +144,7 @@ class NoiseMeasurement: return self.freq[crop:], psd[crop:] def sample_spectrum( - self, nt: int, dt: float | None = None, log_mode: bool = False + self, nf: int, dt: float | None = None, log_mode: bool = False ) -> tuple[np.ndarray, np.ndarray]: """ sample an amplitude spectrum with nt points. The corresponding sample spacing in the time @@ -154,8 +160,6 @@ class NoiseMeasurement: log_mode : bool, optional sample on a log-log scale rather than on a linear scale, by default False """ - if nt % 2: - raise ValueError(f"nt must be an even number, got {nt!r}") fmax = 0.5 / dt if dt is not None else self.freq.max() if fmax > self.freq.max(): @@ -164,7 +168,7 @@ class NoiseMeasurement: f" goes up to {self.freq.max():g}Hz" ) - f = np.linspace(0, fmax, nt // 2 + 1) + f = np.linspace(0, fmax, nf) if log_mode: interp = np.zeros_like(f) ind = self.freq > 0 @@ -177,9 +181,10 @@ class NoiseMeasurement: right=np.log(self.psd[ind][-1]), ) ) + if self.freq[0] == 0: + interp[0] = self.psd[0] else: interp = np.interp(f, self.freq, self.psd, left=0, right=self.psd[-1]) - interp[0] = 0 return f, interp def time_series( @@ -194,7 +199,7 @@ class NoiseMeasurement: Parameters ---------- - nt_or_phase : int + nt : int number of points to sample. must be even. phase : np.ndarray, shape (nt//2)+1 | None, optional phase to apply to the amplitude spectrum before taking the inverse Fourier transform. @@ -205,15 +210,22 @@ class NoiseMeasurement: sample on a log-log scale rather than on a linear scale, by default False """ - freq, spec = self.sample_spectrum(nt, dt, log_mode) + freq, spec = self.sample_spectrum(nt // 2 + 1, dt, log_mode) + spec[1:-1] *= 0.5 if phase is None: - phase = 2 * np.pi * self.rng.random(len(freq)) + phase = 2 * np.pi * self.rng.random(len(freq) - 2) time, dt = math.irfftfreq(freq, True) - amp = np.sqrt(spec) * np.exp(1j * phase) - signal = np.fft.irfft(amp) * np.sqrt(0.5 * nt / dt) + amp = np.asarray(np.sqrt(spec), dtype=complex) - return time, signal + # DC is assumed to be always positive + # Nyquist frequency has an unknowable phase + amp[-1] *= self.rng.choice((1, -1)) + amp[1:-1] *= np.exp(1j * phase) + + signal = np.fft.irfft(amp) * np.sqrt(nt / dt) + + return time, np.fft.fftshift(signal) def root_integrated(self) -> np.ndarray: """ diff --git a/tests/test_noise.py b/tests/test_noise.py index ece6b8a..85b8019 100644 --- a/tests/test_noise.py +++ b/tests/test_noise.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import scgenerator as sc @@ -17,3 +18,48 @@ def test_segmentation(): assert np.all( sc.noise.segments(t, 7) == np.vstack([r, r + 4, r + 8, r + 12, r + 16, r + 20, r + 24]) ) + + +def test_normalisation(): + rng = np.random.default_rng(56) + t = np.linspace(-10, 10, 512) + s = np.exp(-((t / 2.568) ** 2)) + rng.random(len(t)) / 15 + + target = np.sum(sc.abs2(np.fft.fft(s))) / 512 + noise = sc.noise.NoiseMeasurement.from_time_series(s, 1, "Square", force_no_dc=False) + assert np.sum(noise.psd) == pytest.approx(target) + + +def test_no_dc(): + rng = np.random.default_rng(56) + t = np.linspace(-10, 10, 512) + s = rng.normal(0, 1, len(t)).cumsum() + noise = sc.noise.NoiseMeasurement.from_time_series(s, 1) + assert noise.psd[0] == pytest.approx(0) + + +def test_time_and_back(): + """ + sampling a time series from a spectrum and transforming + it back should yield the same spectrum + """ + rng = np.random.default_rng(57) + t = np.linspace(-10, 10, 512) + signal = np.exp(-((t / 2.568) ** 2)) + rng.random(len(t)) / 15 + + noise = sc.noise.NoiseMeasurement.from_time_series(signal, 1, "Square", force_no_dc=False) + _, new_signal = noise.time_series(len(t)) + new_noise = sc.noise.NoiseMeasurement.from_time_series( + new_signal, 1, "Square", force_no_dc=False + ) + assert new_noise.psd == pytest.approx(noise.psd) + + +def test_sampling(): + f = np.geomspace(10, 2e6, 138) + spec = 1 / (f + 1) + + noise = sc.noise.NoiseMeasurement(f, spec) + + assert noise.sample_spectrum(257)[0][0] == 0 + assert noise.sample_spectrum(257, log_mode=True)[0][0] == 0