new plottable methode

src/scgenerator/noise.py

@@ -7,7 +7,7 @@ import numpy as np
 from scipy.integrate import cumulative_trapezoid
 from scipy.interpolate import interp1d
 
-from scgenerator import math
+from scgenerator import math, units
 
 
 @dataclass
@@ -99,6 +99,50 @@ class NoiseMeasurement:
         else:
             return self.freq, self.psd_dBc
 
+    def plottable(
+        self,
+        dB: bool = True,
+        wavelength: float | None = None,
+        power: float | None = None,
+        crop: int = 1,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Transforms the PSD in a way that makes it easy to plot
+
+        Parameters
+        ----------
+        dB : bool, optional
+            transform the PSD from a linear scale to dB, by default True
+        wavelength, power : float | None, optional
+            if both are provided, will be used to compute the quantum noise limit and the PSD will
+            be output relative to that amount
+        crop : int, optional
+            how many low frequency points to drop. 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.
+
+        Returns
+        -------
+        np.ndarray, shape (n,)
+            frequency array
+        np.ndarray, shape (n,)
+            psd, with the desired transformation applied
+
+        Example
+        -------
+        >>> noise = NoiseMeasurement(*np.load("my_measurement.npy"))
+        >>> plt.plot(*noise.plottable(wavelength=800e-9, power=0.3)) # dB above quantum limit
+        >>> plt.xscale("log")
+        >>> plt.show()
+        """
+        psd = self.psd
+        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[crop:], psd[crop:]
+
     def sample_spectrum(self, nt: int, dt: float | None = None) -> tuple[np.ndarray, np.ndarray]:
         """
         sample an amplitude spectrum with nt points. The corresponding sample spacing in the time
@@ -209,3 +253,7 @@ def segments(signal: np.ndarray, num_segments: int) -> np.ndarray:
     seg = np.arange(seg_size)
     off = int(n_init / (num_segments + 1))
     return np.array([signal[seg + i * off] for i in range(num_segments)])
+
+
+def quantum_noise_limit(wavelength: float, power: float) -> float:
+    return units.m(wavelength) * units.hbar * 2 / power