diff --git a/examples/noise_study.py b/examples/noise_study.py index 59a7523..3d3d2eb 100644 --- a/examples/noise_study.py +++ b/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() diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py index 6bd4121..109676a 100644 --- a/src/scgenerator/noise.py +++ b/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: """