diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py index 6b86ce0..6e3bc74 100644 --- a/src/scgenerator/noise.py +++ b/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