improved noise

src/scgenerator/math.py

@@ -774,4 +774,4 @@ def mean_angle(values: np.ndarray, axis: int = 0):
         values, np.abs(values), out=np.zeros(values.shape, dtype=values.dtype), where=values != 0
     )
     total_phase = np.sum(values, axis=axis)
-    return (total_phase) / np.abs(total_phase)
+    return np.divide(total_phase, np.abs(total_phase), where=total_phase != 0, out=total_phase)

src/scgenerator/noise.py

@@ -70,9 +70,11 @@ class NoiseMeasurement:
             take out the DC component (0-frequency) of each segement after segmentation
         """
         signal = np.asanyarray(signal)
+        if signal.ndim > 1:
+            raise ValueError(
+                f"got signal of shape {signal.shape}. Only one 1D signals are supported"
+            )
         signal_segments = segments(signal, num_segments)
-        if force_no_dc:
-            signal_segments = (signal_segments.T - signal_segments.mean(axis=1)).T
         n = signal_segments.shape[-1]
         try:
             window_arr = cls._window_functions[window](n)
@@ -84,10 +86,14 @@ class NoiseMeasurement:
         window_correction = np.sum(window_arr**2) / n
         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(2 * dt / n)
+        psd = np.fft.rfft(signal_segments) * np.sqrt(dt / n)
         phase = math.mean_angle(psd)
         psd = psd.real**2 + psd.imag**2
+        psd[..., 1:-1] *= 2
         return cls(freq, psd.mean(axis=0) / window_correction, phase=phase)
 
     def plottable(
@@ -138,7 +144,7 @@ class NoiseMeasurement:
         return self.freq[crop:], psd[crop:]
 
     def sample_spectrum(
-        self, nt: int, dt: float | None = None, log_mode: bool = False
+        self, nf: int, dt: float | None = None, log_mode: bool = False
     ) -> tuple[np.ndarray, np.ndarray]:
         """
         sample an amplitude spectrum with nt points. The corresponding sample spacing in the time
@@ -154,8 +160,6 @@ class NoiseMeasurement:
         log_mode : bool, optional
             sample on a log-log scale rather than on a linear scale, by default False
         """
-        if nt % 2:
-            raise ValueError(f"nt must be an even number, got {nt!r}")
 
         fmax = 0.5 / dt if dt is not None else self.freq.max()
         if fmax > self.freq.max():
@@ -164,7 +168,7 @@ class NoiseMeasurement:
                 f" goes up to {self.freq.max():g}Hz"
             )
 
-        f = np.linspace(0, fmax, nt // 2 + 1)
+        f = np.linspace(0, fmax, nf)
         if log_mode:
             interp = np.zeros_like(f)
             ind = self.freq > 0
@@ -177,9 +181,10 @@ class NoiseMeasurement:
                     right=np.log(self.psd[ind][-1]),
                 )
             )
+            if self.freq[0] == 0:
+                interp[0] = self.psd[0]
         else:
             interp = np.interp(f, self.freq, self.psd, left=0, right=self.psd[-1])
-            interp[0] = 0
         return f, interp
 
     def time_series(
@@ -194,7 +199,7 @@ class NoiseMeasurement:
 
         Parameters
         ----------
-        nt_or_phase : int
+        nt : int
             number of points to sample. must be even.
         phase : np.ndarray, shape (nt//2)+1 | None, optional
             phase to apply to the amplitude spectrum before taking the inverse Fourier transform.
@@ -205,15 +210,22 @@ class NoiseMeasurement:
             sample on a log-log scale rather than on a linear scale, by default False
         """
 
-        freq, spec = self.sample_spectrum(nt, dt, log_mode)
+        freq, spec = self.sample_spectrum(nt // 2 + 1, dt, log_mode)
+        spec[1:-1] *= 0.5
         if phase is None:
-            phase = 2 * np.pi * self.rng.random(len(freq))
+            phase = 2 * np.pi * self.rng.random(len(freq) - 2)
         time, dt = math.irfftfreq(freq, True)
 
-        amp = np.sqrt(spec) * np.exp(1j * phase)
+        amp = np.asarray(np.sqrt(spec), dtype=complex)
-        signal = np.fft.irfft(amp) * np.sqrt(0.5 * nt / dt)
 
-        return time, signal
+        # DC is assumed to be always positive
+        # Nyquist frequency has an unknowable phase
+        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:
         """

tests/test_noise.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 
 import scgenerator as sc
 
@@ -17,3 +18,48 @@ def test_segmentation():
     assert np.all(
         sc.noise.segments(t, 7) == np.vstack([r, r + 4, r + 8, r + 12, r + 16, r + 20, r + 24])
     )
+
+
+def test_normalisation():
+    rng = np.random.default_rng(56)
+    t = np.linspace(-10, 10, 512)
+    s = np.exp(-((t / 2.568) ** 2)) + rng.random(len(t)) / 15
+
+    target = np.sum(sc.abs2(np.fft.fft(s))) / 512
+    noise = sc.noise.NoiseMeasurement.from_time_series(s, 1, "Square", force_no_dc=False)
+    assert np.sum(noise.psd) == pytest.approx(target)
+
+
+def test_no_dc():
+    rng = np.random.default_rng(56)
+    t = np.linspace(-10, 10, 512)
+    s = rng.normal(0, 1, len(t)).cumsum()
+    noise = sc.noise.NoiseMeasurement.from_time_series(s, 1)
+    assert noise.psd[0] == pytest.approx(0)
+
+
+def test_time_and_back():
+    """
+    sampling a time series from a spectrum and transforming
+    it back should yield the same spectrum
+    """
+    rng = np.random.default_rng(57)
+    t = np.linspace(-10, 10, 512)
+    signal = np.exp(-((t / 2.568) ** 2)) + rng.random(len(t)) / 15
+
+    noise = sc.noise.NoiseMeasurement.from_time_series(signal, 1, "Square", force_no_dc=False)
+    _, new_signal = noise.time_series(len(t))
+    new_noise = sc.noise.NoiseMeasurement.from_time_series(
+        new_signal, 1, "Square", force_no_dc=False
+    )
+    assert new_noise.psd == pytest.approx(noise.psd)
+
+
+def test_sampling():
+    f = np.geomspace(10, 2e6, 138)
+    spec = 1 / (f + 1)
+
+    noise = sc.noise.NoiseMeasurement(f, spec)
+
+    assert noise.sample_spectrum(257)[0][0] == 0
+    assert noise.sample_spectrum(257, log_mode=True)[0][0] == 0