From e13192259ef0ad7968832b673e9356dfa1e1d872 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <dedebenui@protonmail.com>
Date: Tue, 6 Feb 2024 16:28:28 +0100
Subject: [PATCH] quick fixes

---
 src/scgenerator/noise.py   | 11 ++++++++---
 src/scgenerator/spectra.py |  2 +-
 tests/test_integrator.py   |  2 ++
 tests/test_noise.py        |  4 ++--
 4 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py
index 109676a..a75c49d 100644
--- a/src/scgenerator/noise.py
+++ b/src/scgenerator/noise.py
@@ -211,18 +211,23 @@ class NoiseMeasurement:
         freq, amp = self.sample_spectrum(nf, dt, log_mode)
         fs = freq[-1] * 2
 
-        right = -1 if nf % 2 else None
+        if nf % 2:
+            right = -1
+            phase_n = nf - 2
+        else:
+            right = None
+            phase_n = nf - 1
 
         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)))
+            amp[1:right] *= np.exp(2j * np.pi * self.rng.random((phase_n, 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))
+            amp[1:right] *= np.exp(2j * np.pi * self.rng.random(phase_n))
             signal = np.fft.irfft(amp) * np.sqrt((amp.size - 1) * 2 * fs)
             time = np.arange(len(signal)) / fs
         return time, signal
diff --git a/src/scgenerator/spectra.py b/src/scgenerator/spectra.py
index 4c07de1..9130b1d 100644
--- a/src/scgenerator/spectra.py
+++ b/src/scgenerator/spectra.py
@@ -54,7 +54,7 @@ class Spectrum(np.ndarray):
             b"real": np.fft.irfft,
             b"scic": sfft.ifft,
             b"scir": sfft.irfft,
-        }
+        }[buf[:4]]
         shape_n = buf[4]
         nt, *shape = (
             int.from_bytes(buf[5 + i * 8 : 5 + (i + 1) * 8], "big") for i in range(shape_n)
diff --git a/tests/test_integrator.py b/tests/test_integrator.py
index 45688f5..f5c722d 100644
--- a/tests/test_integrator.py
+++ b/tests/test_integrator.py
@@ -56,6 +56,8 @@ def test_rk43_soliton(plot=False):
         op.constant_quantity(np.zeros(n)),
         op.envelope_spm(0),
         no_op,
+        np.fft.fft,
+        np.fft.ifft,
     )
 
     res = sc.integrate(spec0, end, lin, non_lin, targets=targets, atol=1e-10, rtol=1e-9)
diff --git a/tests/test_noise.py b/tests/test_noise.py
index 3d36111..3dd956e 100644
--- a/tests/test_noise.py
+++ b/tests/test_noise.py
@@ -17,14 +17,14 @@ def test_normalisation():
 def test_time_and_back():
     """
     sampling a time series from a spectrum and transforming
-    it back should yield the same spectrum
+    it back yields 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, "boxcar", detrend=False)
-    _, new_signal = noise.time_series(len(t))
+    _, new_signal = noise.time_series(len(noise.freq))
     new_noise = sc.noise.NoiseMeasurement.from_time_series(new_signal, 1, "boxcar", detrend=False)
     assert new_noise.psd == pytest.approx(noise.psd)