diff --git a/src/scgenerator/noise.py b/src/scgenerator/noise.py index f5bfa56..0b19191 100644 --- a/src/scgenerator/noise.py +++ b/src/scgenerator/noise.py @@ -80,7 +80,7 @@ class NoiseMeasurement: dB: bool = True, wavelength: float | None = None, power: float | None = None, - crop: int = 1, + crop: int | float = 1, ) -> tuple[np.ndarray, np.ndarray]: """ Transforms the PSD in a way that makes it easy to plot @@ -117,6 +117,9 @@ class NoiseMeasurement: if dB: psd = math.to_dB(psd, ref=1.0) + if isinstance(crop, (float, np.floating)): + crop = math.argclosest(self.freq, crop) + return self.freq[crop:], psd[crop:] def sample_spectrum(self, nt: int, dt: float | None = None) -> tuple[np.ndarray, np.ndarray]: @@ -221,11 +224,24 @@ def segments(signal: np.ndarray, num_segments: int) -> np.ndarray: if num_segments == 1: return signal[None] n_init = len(signal) - seg_size = 1 << int(np.log2(n_init / (num_segments + 1))) + 1 + seg_size = segement_size(n_init, num_segments) 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 segement_size(nt: int, num_segments: int) -> int: + return 1 << int(np.log2(nt / (num_segments + 1))) + 1 + + +def get_frequencies(nt: int, num_segments: int, dt: float) -> np.ndarray: + """ + returns the frequency array that would be associated to a NoiseMeasurement.from_time_series + call, where `nt` is the size of the signal. + """ + + return np.fft.rfftfreq(segement_size(nt, num_segments), dt) + + def quantum_noise_limit(wavelength: float, power: float) -> float: return units.m_rads(wavelength) * units.hbar * 2 / power