reverted some noise changes

examples/noise_study.py

@@ -5,34 +5,54 @@ import scgenerator as sc
 
 
 def main():
-    nperseg = 513
+    nf = 513
     nseg = 240
-    ntot = (nperseg - 1) * (nseg - 1)
+    ntot = (nf - 1) * (nseg - 1)
 
     fs = 44100
-    nf = 45611
-    f = np.linspace(0, fs // 2, nf)
+    nf_full = 45611
+    f = np.linspace(0, fs // 2, nf_full)
 
     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()
+    spec += np.exp(-(((f - (0.1 * fs)) / (0.5 * fs)) ** 2))
+    spec += np.exp(-(((f - (0.45 * fs)) / (0.02 * fs)) ** 2))
+    # spec = np.ones_like(f) * 1e-5
+    # spec += np.exp(-(((f - 15000) / 300) ** 2))
 
     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)
+    t, y = noise.time_series(nf=nf, nseg=nseg)
+    tf, yf = noise.time_series(nf=ntot // 2 + 1)
+    print(y.std(), yf.std())
+
     plt.plot(t, y, ls="", marker=".")
+    plt.plot(tf, yf, ls="", marker=".")
     plt.show()
 
-    newnoise = noise.from_time_series(y, t[1] - t[0], nperseg=(nperseg - 1) * 2)
+    newnoise = noise.from_time_series(y, t[1] - t[0], nf=nf)
+    newnoise_full = noise.from_time_series(yf, tf[1] - tf[0], nperseg=(nf - 1) * 2)
+
+    print(newnoise.psd[1:-1].var())
+    print(newnoise_full.psd[1:-1].var())
+
+    fig, ax = plt.subplots()
+    ax.plot(*noise.plottable())
+    ax.plot(*newnoise.plottable())
+    ax.plot(*newnoise_full.plottable(discard_nyquist=True))
+    ax.set_xscale("log")
+
+    axin = plt.gca().inset_axes(
+        [0.1, 0.1, 0.3, 0.3],
+        xlim=(1.5e4, fs / 1.9),
+        ylim=(-3.1, 2.4),
+        yticklabels=[],
+    )
+    # axin.set_xscale("log")
+    axin.plot(*newnoise.plottable())
+    axin.plot(*newnoise_full.plottable())
+    axin.plot(*noise.plottable())
+    ax.indicate_inset_zoom(axin)
+    axin.tick_params(labelbottom=False)
 
-    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()
 
 

src/scgenerator/noise.py

@@ -36,6 +36,7 @@ class NoiseMeasurement:
         dt: float = 1.0,
         window: str | None = "hann",
         nperseg: int | None = None,
+        nf: int | None = None,
         detrend: bool | str = "constant",
     ) -> NoiseMeasurement:
         """
@@ -53,6 +54,9 @@ class NoiseMeasurement:
             number of points per segment. The PSD of each segment is computed and then averaged
             to reduce variange. By default None, which means only one segment (i.e. the full signal
             at once) is computed.
+        nf : int, optional
+            specify the number of frequency points rather than the size of each segments. if both
+            are specified. Takes precedence over nperseg when both are specified.
         detrend : bool, optional
             remove DC and optionally linear trend, by default only removes DC. See
             scipy.signal.welch for more details.
@@ -62,12 +66,18 @@ class NoiseMeasurement:
             raise ValueError(
                 f"got signal of shape {signal.shape}. Only one 1D signals are supported"
             )
-        if nperseg is None:
+
+        if nf is not None:
+            nperseg = (nf - 1) * 2
+        elif nperseg is None:
             nperseg = len(signal)
+
         if detrend is True:
             detrend = "constant"
+
         if window is None:
             window = "boxcar"
+
         freq, psd = ss.welch(
             signal, fs=1 / dt, window=window, nperseg=nperseg, detrend=detrend, scaling="density"
         )
@@ -79,7 +89,8 @@ class NoiseMeasurement:
         dB: bool = True,
         wavelength: float | None = None,
         power: float | None = None,
-        left: int = 1,
+        left: int | None = None,
+        discard_nyquist: bool = False,
     ) -> tuple[np.ndarray, np.ndarray]:
         """
         Transforms the PSD in a way that makes it easy to plot
@@ -92,11 +103,13 @@ class NoiseMeasurement:
             if both are provided, will be used to compute the quantum noise limit and the PSD will
             be output relative to that amount
         left : int, optional
-            Index of the first frequency point. The default (1) is to drop only the 0 frequency
-            point, as it is outside the plot range when plotting the PSD on a log-log plot.
-        right : int, optional
-            index of the last frequency point, non inclusive (i.e. like slices). The default (-1) is
-            to drop only the Nyquist frequency, as segmentation artificially reduces its value.
+            Index of the first frequency point. The default is to drop only the 0 frequency
+            point if it exits to simplify plotting in log-log scales. choose `left=0` to return
+            everything
+        discard_nyquist : bool, optional
+            crop the psd before the Nyquist frequency, as this is a special point that is not
+            guaranteed to be present and/or accurate depending on how the signal was generated, by
+            default False
 
         Returns
         -------
@@ -113,13 +126,21 @@ class NoiseMeasurement:
         >>> plt.show()
         """
         psd = self.psd
+        freq = self.freq
         if wavelength is not None and power is not None:
             psd = psd / quantum_noise_limit(wavelength, power)
 
         if dB:
             psd = math.to_dB(psd, ref=1.0)
 
-        return self.freq[left:], psd[left:]
+        if freq.size % 2 and discard_nyquist:
+            freq = freq[:-1]
+            psd = psd[:-1]
+
+        if left is None and freq[0] == 0:
+            left = 1
+
+        return freq[left:], psd[left:]
 
     def sample_spectrum(
         self, nf: int, dt: float | None = None, log_mode: bool = False
@@ -168,7 +189,7 @@ class NoiseMeasurement:
     def time_series(
         self,
         nf: int,
-        nseg: int,
+        nseg: int = 1,
         dt: float | None = None,
         log_mode: bool = False,
     ) -> tuple[np.ndarray, np.ndarray]:
@@ -179,7 +200,7 @@ class NoiseMeasurement:
         ----------
         nf : int
             number of frequency points to sample. Recommended is a power of 2 + 1 (129, 513, ...)
-        nseg : int
+        nseg : int | None,
             number of times to sample the spectrum, each time with a different random phase
         dt : float | None, optional
             if given, choose a sample spacing of dt instead of 1/f_max
@@ -190,14 +211,21 @@ class NoiseMeasurement:
         freq, amp = self.sample_spectrum(nf, dt, log_mode)
         fs = freq[-1] * 2
 
-        # 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)))
+        right = -1 if nf % 2 else None
 
-        time, signal = ss.istft(amp, fs, nperseg=(nf - 1) * 2, noverlap=nf - 1, scaling="psd")
-        return time, signal * np.sqrt(2)
+        amp[1:right] *= 0.5
+        amp = np.sqrt(amp) + 0j
+
+        if nseg > 1:
+            amp = np.tile(amp, (nseg, 1)).T
+            amp[1:right] *= np.exp(2j * np.pi * self.rng.random((nf - 2, nseg)))
+            time, signal = ss.istft(amp, fs, nperseg=(nf - 1) * 2, noverlap=nf - 1, scaling="psd")
+            signal *= np.sqrt(2)
+        else:
+            amp[1:right] *= np.exp(2j * np.pi * self.rng.random(nf - 2))
+            signal = np.fft.irfft(amp) * np.sqrt((amp.size - 1) * 2 * fs)
+            time = np.arange(len(signal)) / fs
+        return time, signal
 
     def root_integrated(self) -> np.ndarray:
         """