dedebenui/scgenerator
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

noise improvement

Browse Source 
removed two counteracting `n`. These simplify but should still be present in the theory
This commit is contained in:
Benoît Sierro
2023-10-30 10:23:33 +01:00
parent dfbe47cf52
commit 15e8f30514
1 changed files with 12 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/scgenerator/noise.py
View File

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