noise improvements with (i)stft

examples/noise_study.py

@@ -0,0 +1,40 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+import scgenerator as sc
+
+
+def main():
+    nperseg = 513
+    nseg = 240
+    ntot = (nperseg - 1) * (nseg - 1)
+
+    fs = 44100
+    nf = 45611
+    f = np.linspace(0, fs // 2, nf)
+
+    spec = 1 / (f + 1)
+    spec += np.exp(-(((f - (0.2 * fs)) / (0.5 * fs)) ** 2))
+    plt.plot(f, spec)
+    plt.xscale("log")
+    plt.yscale("log")
+    plt.show()
+
+    noise = sc.noise.NoiseMeasurement(f, spec)
+    t, y = noise.time_series(nf=nperseg, nseg=nseg)
+    # tf, yf = noise.time_series(nperseg=(ntot // 2) + 1, nseg=1)
+    plt.plot(t, y, ls="", marker=".")
+    plt.show()
+
+    newnoise = noise.from_time_series(y, t[1] - t[0], nperseg=(nperseg - 1) * 2)
+
+    plt.plot(*noise.plottable())
+    f, sxx = noise.sample_spectrum(nperseg)
+    plt.plot(f, 10 * np.log10(sxx))
+    plt.plot(*newnoise.plottable())
+    plt.xscale("log")
+    plt.show()
+
+
+if __name__ == "__main__":
+    main()

src/scgenerator/noise.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Callable, ClassVar, Sequence
+from typing import Sequence
 
 import numpy as np
 import scipy.signal as ss
@@ -68,7 +68,9 @@ class NoiseMeasurement:
             detrend = "constant"
         if window is None:
             window = "boxcar"
-        freq, psd = ss.welch(signal, fs=1 / dt, window=window, nperseg=nperseg, detrend=detrend)
+        freq, psd = ss.welch(
+            signal, fs=1 / dt, window=window, nperseg=nperseg, detrend=detrend, scaling="density"
+        )
 
         return cls(freq, psd)
 
@@ -78,7 +80,6 @@ class NoiseMeasurement:
         wavelength: float | None = None,
         power: float | None = None,
         left: int = 1,
-        right: int = -1,
     ) -> tuple[np.ndarray, np.ndarray]:
         """
         Transforms the PSD in a way that makes it easy to plot
@@ -118,7 +119,7 @@ class NoiseMeasurement:
         if dB:
             psd = math.to_dB(psd, ref=1.0)
 
-        return self.freq[left:right], psd[left:right]
+        return self.freq[left:], psd[left:]
 
     def sample_spectrum(
         self, nf: int, dt: float | None = None, log_mode: bool = False
@@ -166,8 +167,8 @@ class NoiseMeasurement:
 
     def time_series(
         self,
-        nt: int | np.ndarray,
+        nf: int,
-        phase: np.ndarray | None = None,
+        nseg: int,
         dt: float | None = None,
         log_mode: bool = False,
     ) -> tuple[np.ndarray, np.ndarray]:
@@ -176,35 +177,27 @@ class NoiseMeasurement:
 
         Parameters
         ----------
-        nt : int
+        nf : int
-            number of points to sample. must be even.
+            number of frequency points to sample. Recommended is a power of 2 + 1 (129, 513, ...)
-        phase : np.ndarray, shape (nt//2)+1 | None, optional
+        nseg : int
-            phase to apply to the amplitude spectrum before taking the inverse Fourier transform.
+            number of times to sample the spectrum, each time with a different random phase
-            if None, A random phase is then applied.
         dt : float | None, optional
             if given, choose a sample spacing of dt instead of 1/f_max
         log_mode : bool, optional
             sample on a log-log scale rather than on a linear scale, by default False
         """
 
-        freq, spec = self.sample_spectrum(nt // 2 + 1, dt, log_mode)
+        freq, amp = self.sample_spectrum(nf, dt, log_mode)
-        spec[1:-1] *= 0.5
+        fs = freq[-1] * 2
-        if phase is None:
-            phase = 2 * np.pi * self.rng.random(len(freq) - 2)
-        elif phase.shape[-1] != len(freq) - 2:
-            phase = np.interp(freq, self.freq, phase)[1:-1]
-        time, dt = math.irfftfreq(freq, True)
 
-        amp = np.asarray(np.sqrt(spec), dtype=complex)
+        # istft expects a stft, not a spectrogram
+        amp[1:-1] *= 0.5
+        amp = np.sqrt(amp)
+        amp = np.tile(amp, (nseg, 1)).T + 0j
+        amp[1:-1] *= np.exp(2j * np.pi * self.rng.random((nf - 2, nseg)))
 
-        # DC is assumed to be always positive
+        time, signal = ss.istft(amp, fs, nperseg=(nf - 1) * 2, noverlap=nf - 1, scaling="psd")
-        # Nyquist frequency has an unknowable phase
+        return time, signal * np.sqrt(2)
-        # 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:
         """