From 0e1eb502a4d57e70ef80d1d1858b836a3fa86244 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <benoit.sierro@bluewin.ch>
Date: Tue, 17 Oct 2023 10:46:15 +0200
Subject: [PATCH] noise update (log_mode)

---
 src/scgenerator/noise.py | 33 ++++++++++++++++++++++++++++-----
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py
index 0b19191..6667d7b 100644
--- a/src/scgenerator/noise.py
+++ b/src/scgenerator/noise.py
@@ -122,7 +122,9 @@ class NoiseMeasurement:
 
         return self.freq[crop:], psd[crop:]
 
-    def sample_spectrum(self, nt: int, dt: float | None = None) -> tuple[np.ndarray, np.ndarray]:
+    def sample_spectrum(
+        self, nt: 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
         is 1/freq.max().
@@ -134,6 +136,8 @@ class NoiseMeasurement:
         dt : float | None, optional
             if given, freq will only be sampled up to 0.5/dt. if that value is higher than the
             max of freq, an exception is raised.
+        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}")
@@ -146,12 +150,29 @@ class NoiseMeasurement:
             )
 
         f = np.linspace(0, fmax, nt // 2 + 1)
-        interp = np.interp(f, self.freq, self.psd, left=0, right=self.psd[-1])
-        interp[0] = 0
+        if log_mode:
+            interp = np.zeros_like(f)
+            ind = self.freq > 0
+            interp[1:] = np.exp(
+                np.interp(
+                    np.log(f[1:]),
+                    np.log(self.freq[ind]),
+                    np.log(self.psd[ind]),
+                    left=np.log(self.psd[ind][0]),
+                    right=np.log(self.psd[ind][-1]),
+                )
+            )
+        else:
+            interp = np.interp(f, self.freq, self.psd, left=0, right=self.psd[-1])
+            interp[0] = 0
         return f, interp
 
     def time_series(
-        self, nt: int | np.ndarray, phase: np.ndarray | None = None, dt: float | None = None
+        self,
+        nt: int | np.ndarray,
+        phase: np.ndarray | None = None,
+        dt: float | None = None,
+        log_mode: bool = False,
     ) -> tuple[np.ndarray, np.ndarray]:
         """
         computes a pulse train whose psd matches the measurement
@@ -165,9 +186,11 @@ class NoiseMeasurement:
             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, dt)
+        freq, spec = self.sample_spectrum(nt, dt, log_mode)
         if phase is None:
             phase = 2 * np.pi * self.rng.random(len(freq))
         time, dt = math.irfftfreq(freq, True)