noise improvement

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

View File

@@ -19,9 +19,10 @@ class NoiseMeasurement:
@classmethod
def window_function(cls, name: str):
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")
def wrapper(func: Callable[[int], np.ndarray]):
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