diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py index 8da2eba..0382502 100644 --- a/src/scgenerator/noise.py +++ b/src/scgenerator/noise.py @@ -19,9 +19,10 @@ class NoiseMeasurement: @classmethod def window_function(cls, name: str): + if name in cls._window_functions: + raise ValueError(f"a function labeled {name!r} has already been registered") + def wrapper(func: Callable[[int], np.ndarray]): - if name in cls._window_functions: - raise ValueError(f"a function labeled {name!r} has already been registered") cls._window_functions[name] = func return func @@ -45,7 +46,7 @@ class NoiseMeasurement: cls, signal: Sequence[float], dt: float, - window: str = "Hann", + window: str | None = "Hann", num_segments: int = 1, force_no_dc: bool = True, ) -> NoiseMeasurement: @@ -60,9 +61,10 @@ class NoiseMeasurement: the 0 frequency bin of the PSD window : str | None, optional window to use on the input data to avoid leakage. Possible values are - 'Square' (default), 'Bartlett', 'Welch' and 'Hann'. You may register your own window + 'Square', 'Bartlett', 'Welch' and 'Hann' (default). You may register your own window function by using `NoiseMeasurement.window_function` decorator, and use the name you gave it as this argument. + `None` is an alias for square, since in that case, no windowing is performed. num_segments : int, optional number of segments to cut the signal into. This will trade lower frequency information for better variance of the estimated PSD. The default 1 means no cutting. @@ -83,14 +85,14 @@ class NoiseMeasurement: f"window function {window!r} not found. " f"Possible values are {set(cls._window_functions)}" ) from None - window_correction = np.sum(window_arr**2) / n + window_correction = np.sum(window_arr**2) 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(dt / n) + psd = np.fft.rfft(signal_segments) * np.sqrt(dt) phase = math.mean_angle(psd) psd = psd.real**2 + psd.imag**2 psd[..., 1:-1] *= 2 @@ -220,7 +222,7 @@ class NoiseMeasurement: # DC is assumed to be always positive # Nyquist frequency has an unknowable phase - amp[-1] *= self.rng.choice((1, -1)) + # amp[-1] *= self.rng.choice((1, -1)) amp[1:-1] *= np.exp(1j * phase) signal = np.fft.irfft(amp) * np.sqrt(nt / dt) @@ -248,6 +250,9 @@ def square_window(n: int): return np.ones(n) +NoiseMeasurement._window_functions[None] = square_window + + @NoiseMeasurement.window_function("Bartlett") def bartlett_window(n: int): hn = 0.5 * n