From 79b42021a52e3414ee7997551ace260cff38dca9 Mon Sep 17 00:00:00 2001 From: Valentin Boettcher Date: Wed, 21 Aug 2024 10:59:22 -0400 Subject: [PATCH] clean up calibration plot code --- scripts/experiments/013_step_through.py | 428 ++++++++++++------------ src/rabifun/analysis.py | 61 +++- src/rabifun/plots.py | 73 ++-- src/rabifun/system.py | 2 +- src/ringfit/utils.py | 56 ++++ 5 files changed, 365 insertions(+), 255 deletions(-) diff --git a/scripts/experiments/013_step_through.py b/scripts/experiments/013_step_through.py index d3d22a9..fb69dd6 100644 --- a/scripts/experiments/013_step_through.py +++ b/scripts/experiments/013_step_through.py @@ -1,89 +1,62 @@ +"""A demonstration of the ringdown spectroscopy protocol.""" + from rabifun.system import * from rabifun.plots import * from rabifun.utilities import * +from ringfit.utils import WelfordAggregator from rabifun.analysis import * -from ringfit.data import ScanData -import functools import multiprocessing import copy -from scipy.ndimage import rotate - -# %% interactive -class WelfordAggregator: - """A class to aggregate values using the Welford algorithm. +def solve_shot( + params: Params, t: np.ndarray, t_before: np.ndarray, t_after: np.ndarray +): + """A worker function to solve for the time evolution in separate processes. - The Welford algorithm is an online algorithm to calculate the mean - and variance of a series of values. - - The aggregator keeps track of the number of samples the mean and - the variance. Aggregation of identical values is prevented by - checking the sample index. Tracking can be disabled by setting - the initial index to ``None``. - - See also the `Wikipedia article - `_. - - :param first_value: The first value to aggregate. + :param params: The parameters of the system. + :param t: The time axis. + :param t_before: The time axis before the EOM is switched off. + :param t_after: The time axis after the EOM is switched off. """ - __slots__ = ["n", "mean", "_m_2"] - - def __init__(self, first_value: np.ndarray): - self.n = 1 - self.mean = first_value - self._m_2 = np.zeros_like(first_value) - - def update(self, new_value: np.ndarray): - """Updates the aggregator with a new value.""" - - self.n += 1 - delta = new_value - self.mean - self.mean += delta / self.n - delta2 = new_value - self.mean - self._m_2 += np.abs(delta) * np.abs(delta2) - - @property - def sample_variance(self) -> np.ndarray: - """ - The empirical sample variance. (:math:`\sqrt{N-1}` - normalization.) - """ - - if self.n == 1: - return np.zeros_like(self.mean) - - return self._m_2 / (self.n - 1) - - @property - def ensemble_variance(self) -> np.ndarray: - """The ensemble variance.""" - return self.sample_variance / self.n - - @property - def ensemble_std(self) -> np.ndarray: - """The ensemble standard deviation.""" - return np.sqrt(self.ensemble_variance) - - -def solve_shot(t, params, t_before, t_after): solution = solve(t, params) amps = solution.y[::, len(t_before) - 1 :] return t_after, amps -def make_shots(params, total_lifetimes, eom_range, eom_steps, num_freq): +def make_shots( + params: Params, + total_lifetimes: float, + eom_range: tuple[float, float], + eom_steps: int, + num_freq: int = 1, +): + """Generate a series of shots with varying EOM frequencies. + + The implementation here slightly varies the off time of the laser + so as to introduce some random relative phases of the modes. + + :param params: The parameters of the system. + :param total_lifetimes: The total time of the experiment in + lifetimes. + :param eom_range: The range of EOM frequencies in units of + :any:`params.Ω`. + :param eom_steps: The number of steps in the EOM frequency range. + :param num_freq: The number of frequencies to drive. If a number + greater than 1 is given, the EOM will be driven at multiple + frequencies where the highest frequency is the base frequency + plus an consecutive integer multiples of :any:`params.Ω`. + """ + solutions = [] - - analyze_time = params.lifetimes(total_lifetimes) - params.laser_off_time - t_after = time_axis(params, total_time=analyze_time, resolution=0.01) - - pool = multiprocessing.Pool() - shot_params = [] + rng = np.random.default_rng(seed=0) + off_time = params.laser_off_time or 0 + analyze_time = params.lifetimes(total_lifetimes) - off_time + t_after = time_axis(params, total_time=analyze_time, resolution=0.01) for step in range(eom_steps): base = params.Ω * ( @@ -97,43 +70,59 @@ def make_shots(params, total_lifetimes, eom_range, eom_steps, num_freq): ) current_params.drive_phases = rng.uniform(0, 2 * np.pi, size=num_freq) - off_time = rng.normal(params.laser_off_time, 0.1 * params.laser_off_time) + off_time = rng.normal(off_time, 0.1 * params.laser_off_time) - params.laser_off_time - current_params.laser_off_time = off_time + current_params.laser_off_time = None # off_time current_params.drive_off_time = off_time - current_params.total_lifetimes = (off_time + analyze_time) / params.lifetimes(1) t_before = time_axis(params, total_time=off_time, resolution=0.01) t = np.concatenate([t_before[:-1], t_after + t_before[-1]]) + shot_params.append((current_params, t, t_before, t_after)) - shot_params.append((t, current_params, t_before, t_after)) - + pool = multiprocessing.Pool() solutions = pool.starmap(solve_shot, shot_params) return solutions -def process_shots(solutions, noise_amplitude, params): +def process_shots( + params: Params, + solutions: list[tuple[np.ndarray, np.ndarray]], + noise_amplitude: float, + num_freq: int, +): + """ + Calculates the normalized average Fourier power spectrum of a + series of experimental (simulated) shots. + + :param params: The parameters of the system. + :param solutions: A list of solutions to process returned by + :any:`solve_shot`. + :param noise_amplitude: The amplitude of the noise to add to the + signal. + + The amplitude is normalized by 2/η which is roughly the steady + state signal amplitude if a bath mode is excited resonantly by + a unit-strength input. + + :param num_freq: The number of frequencies to drive. See + :any:`make_shots` for details. + """ + rng = np.random.default_rng(seed=0) - # let us get a measure calibrate the noise strength - signals = [] - for t, amps in solutions: - signal = output_signal(t, amps, params) - signals.append((t, signal)) - - noise_amplitude *= 2 * 2 * np.pi / params.η + noise_amplitude /= params.η * np.pi aggregate = None - for t, signal in signals: + for t, amps in solutions: + signal = output_signal(t, amps, params) signal += rng.normal(scale=noise_amplitude, size=len(signal)) window = (0, t[-1]) freq, fft = fourier_transform( t, signal, - low_cutoff=0.1 * params.Ω, - high_cutoff=params.Ω * 5, + low_cutoff=0.3 * params.Ω, + high_cutoff=params.Ω * (1 + num_freq), window=window, ) @@ -145,87 +134,128 @@ def process_shots(solutions, noise_amplitude, params): else: aggregate.update(power) + assert aggregate is not None # appease pyright + max_power = np.max(aggregate.mean) return (freq, aggregate.mean / max_power, aggregate.ensemble_std / max_power) -def plot_power_spectrum( - ax_spectrum, freq, average_power_spectrum, σ_power_spectrum, params, annotate=True +def process_and_plot_results( + params: Params, + ax: plt.Axes, + freq: np.ndarray, + average_power_spectrum: np.ndarray, + σ_power_spectrum: np.ndarray, + annotate: bool = True, ): - # ax_spectrum.plot(freq, average_power_spectrum) - runtime = RuntimeParams(params) + """ + Fits the ringdown spectrum and plots the results. + + :param params: The parameters of the system. + :param ax: The axis to plot on. + :param freq: The frequency array. + :param average_power_spectrum: The average power spectrum obtained from :any:`process_shots`. + :param σ_power_spectrum: The standard deviation of the power + spectrum. + :param annotate: Whether to annotate the plot with peak and mode positions. + """ ringdown_params = RingdownParams( fω_shift=params.measurement_detuning, - mode_window=(4, 4), + mode_window=(params.N, params.N), fΩ_guess=params.Ω, fδ_guess=params.Ω * params.δ, η_guess=0.5, - absolute_low_cutoff=0.1 * params.Ω, + absolute_low_cutoff=0.3 * params.Ω, ) peak_info = find_peaks( - freq, average_power_spectrum, ringdown_params, prominence=0.05, height=0.1 + freq, + average_power_spectrum, + ringdown_params, + prominence=0.05, + height=0.1, + σ_power=σ_power_spectrum, ) - peak_info, lm_result = refine_peaks( - peak_info, ringdown_params, height_cutoff=0.05, σ=σ_power_spectrum - ) + peak_info = refine_peaks(peak_info, ringdown_params, height_cutoff=0.05) - peak_info.power = average_power_spectrum - plot_spectrum_and_peak_info( - ax_spectrum, peak_info, ringdown_params, annotate=annotate - ) - if lm_result is not None: - # print(lm_result.fit_report()) + plot_spectrum_and_peak_info(ax, peak_info, annotate=annotate) + + if peak_info.lm_result is not None: fine_freq = np.linspace(freq.min(), freq.max(), 5000) - fine_fit = lm_result.eval(ω=fine_freq) - ax_spectrum.plot(fine_freq, fine_fit, color="red") - ax_spectrum.set_ylim(-0.1, max(1, fine_fit.max() * 1.1)) + fine_fit = peak_info.lm_result.eval(ω=fine_freq) + ax.plot(fine_freq, fine_fit - peak_info.noise_floor, color="C3", zorder=-100) + ax.set_ylim(-0.1, max(1, fine_fit.max() * 1.1)) - print(runtime.Ωs.real / (2 * np.pi)) + ax.set_xlabel("Frequency (MHz)") - for i, peak_freq in enumerate(runtime.Ωs): - pos = np.abs( - params.measurement_detuning - - peak_freq.real / (2 * np.pi) - + params.δ * params.Ω - + params.laser_detuning, - ) - - ax_spectrum.axvline( - pos, - color="black", - alpha=0.5, - linestyle="--", - zorder=-100, - ) - - ax_spectrum.axvspan( - pos - peak_freq.imag / (2 * np.pi), - pos + peak_freq.imag / (2 * np.pi), - color="black", - alpha=0.05, - linestyle="--", - zorder=-100, - ) + if annotate: + annotate_ringodown_mode_positions(params, ax) def generate_data( + Ω=13, + η=0.2, g_0=0.5, η_factor=5, noise_amplitude=0.3, laser_detuning=0, + laser_on_mode=0, N=10, eom_ranges=(0.5, 2.5), eom_steps=20, - small_loop_detuning=0, excitation_lifetimes=2, measurement_lifetimes=4, num_freq=3, + extra_title="", ): - η = 0.2 - Ω = 13 + """Simulate and plot the ringdown spectroscopy protocol. + + The idea is to have the laser on ``laser_on_mode`` and to sweep + the EOM frequency over a range of values given in ``eom_ranges`` + in ``eom_steps`` steps. For each step, the laser and EOM are + inputting into the system for a time given by + ``excitation_lifetimes``. Then, the ringdown signal is collected + for a time given by ``measurement_lifetimes``. (Lifetime units + are given by ``η``.) The resulting power spectra are averaged and + then fitted. + + :param Ω: The FSR of the system. + :param η: The decay rate of the system. + :param g_0: The coupling strength of the system in units of + :any:`Ω`. Note that the effective coupling strength between + the ``A`` site and the bath modes is reduced by a factor of + :math:`\sqrt{2}`. + + :param η_factor: The factor by which the decay rate of the A site + is greater. + :param noise_amplitude: The amplitude of the noise to add to the + signal. See :any:`process_shots` for details. + :param laser_detuning: The detuning of the laser from the the mode + it is exciting. + :param laser_on_mode: The mode that the laser is exciting. + :param N: The number of bath modes. + :param eom_ranges: The range of EOM frequencies in units of + :any:`Ω`. + :param eom_steps: The number of steps in the EOM frequency range. + :param excitation_lifetimes: The time the EOM is driving the + system. + :param measurement_lifetimes: The time the system is left to ring + down. + + Note that the laser is not turned off during the ringdown. + + :param num_freq: The number of frequencies to drive. See + :any:`make_shots` for details. + :param extra_title: A string to add to the title of the plot. + + :returns: The figure containing the plot. + """ + + final_laser_detuning = laser_detuning + ( + 0 if laser_on_mode == 0 else (laser_on_mode - 1 / 4) * Ω + ) params = Params( η=η, @@ -233,8 +263,8 @@ def generate_data( Ω=Ω, δ=1 / 4, ω_c=0.1, - g_0=g_0, - laser_detuning=laser_detuning, + g_0=g_0 * num_freq, # as it would be normalized otherwise + laser_detuning=final_laser_detuning, N=N, N_couplings=N, measurement_detuning=0, @@ -242,8 +272,8 @@ def generate_data( rwa=False, flat_energies=False, correct_lamb_shift=0, - laser_off_time=0, - small_loop_detuning=small_loop_detuning, + laser_off_time=None, + small_loop_detuning=0, drive_override=(np.array([]), np.array([])), ) @@ -258,84 +288,42 @@ def generate_data( num_freq, ) - (sol_on_res) = make_shots( - params, - excitation_lifetimes + measurement_lifetimes, - ((1 + params.δ), (1 + params.δ)), - 1, - num_freq, - ) - - (sol_on_res_bath) = make_shots( - params, - excitation_lifetimes + measurement_lifetimes, - ((1 + params.δ * 1.1), (1 + params.δ * 1.1)), - 1, - num_freq, - ) - freq, average_power_spectrum, σ_power_spectrum = process_shots( - solutions, noise_amplitude, params - ) - _, spectrum_on_resonance, σ_power_spectrum_on_resonance = process_shots( - sol_on_res, noise_amplitude, params - ) - _, spectrum_on_resonance_bath, σ_power_spectrum_on_resonance_bath = process_shots( - sol_on_res_bath, noise_amplitude, params + params, + solutions, + noise_amplitude, + num_freq, ) - fig = make_figure() + fig = make_figure(extra_title, figsize=(10, 6)) fig.clear() + ax = fig.subplots() - fig.suptitle(f""" - Spectroscopy Protocol V2 - - Ω/2π = {params.Ω}MHz, η/2π = {params.η}MHz, g_0 = {params.g_0}Ω, N = {params.N} - noise amplitude = {noise_amplitude} * 2/η, η_A = {η_factor} x η, EOM stepped from {eom_ranges[0]:.2f}Ω to {eom_ranges[1]:.2f}Ω in {eom_steps} steps - total time = {(excitation_lifetimes + measurement_lifetimes) * eom_steps / (params.η * 1e6)}s - """) - ax_multi, ax_single, ax_single_bath = fig.subplot_mosaic("AA\nBC").values() - - plot_power_spectrum( - ax_multi, freq, average_power_spectrum, σ_power_spectrum, params + process_and_plot_results(params, ax, freq, average_power_spectrum, σ_power_spectrum) + ax.text( + 0.01, + 0.95, + f"""$Ω/2π = {params.Ω}$MHz +$η/2π = {params.η}MHz$ +$g_0 = {params.g_0}Ω$ +$N = {params.N}$ +noise = ${noise_amplitude * 2}$ +$η_A = {η_factor}η$ +EOM range = {eom_ranges[0]:.2f}Ω to {eom_ranges[1]:.2f}Ω +EOM steps = {eom_steps} +excitation time = {excitation_lifetimes} lifetimes +measurement time = {measurement_lifetimes} lifetimes +on mode = {laser_on_mode} +laser detuning = {laser_detuning} +num freq = {num_freq} +total time = {(excitation_lifetimes + measurement_lifetimes) * eom_steps / (params.η * 1e6)}s""", + transform=ax.transAxes, + ha="left", + va="top", + size=10, + bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.5), ) - plot_power_spectrum( - ax_single, - freq, - spectrum_on_resonance, - σ_power_spectrum_on_resonance, - params, - annotate=False, - ) - plot_power_spectrum( - ax_single_bath, - freq, - spectrum_on_resonance_bath, - σ_power_spectrum_on_resonance_bath, - params, - annotate=False, - ) - - runtime = RuntimeParams(params) - for ax in [ax_multi, ax_single, ax_single_bath]: - ax.set_xlabel("Frequency (MHz)") - ax.sharex(ax_multi) - ax.sharey(ax_multi) - - ax_ticks = ax.twiny() - ax_ticks.sharey(ax) - ax_ticks.set_xticks(runtime.ringdown_frequencies) - ax_ticks.set_xticklabels( - [mode_name(i, params.N) for i in range(2 * params.N + 2)] - ) - ax_ticks.plot(freq, np.zeros_like(freq), alpha=0) - ax_ticks.set_xlim(ax.get_xlim()) - - ax_multi.set_title("Averaged Power Spectrum") - ax_single.set_title("Single-shot, No detuning") - ax_single_bath.set_title("Single-shot, EOM 10% detuned") - - # ax_spectrum.set_yscale(yscale) + ax.set_title(extra_title) fig.tight_layout() return fig @@ -344,15 +332,31 @@ def generate_data( # %% save if __name__ == "__main__": fig = generate_data( - g_0=0.5, + g_0=0.2, η_factor=5, - noise_amplitude=5e-3, + noise_amplitude=0.3, N=5, - eom_ranges=(0.7, 0.9), # (1.9, 2.1), + eom_ranges=(0.7, 0.9), eom_steps=100, - small_loop_detuning=0, laser_detuning=0, - excitation_lifetimes=1, - measurement_lifetimes=3, + laser_on_mode=0, + excitation_lifetimes=2, + measurement_lifetimes=4, num_freq=4, + extra_title="Laser on A site", + ) + + fig = generate_data( + g_0=0.2, + η_factor=5, + noise_amplitude=0.3, + N=5, + eom_ranges=(1.2, 1.3), + eom_steps=100, + laser_detuning=0, + laser_on_mode=-1, + excitation_lifetimes=2, + measurement_lifetimes=4, + num_freq=1, + extra_title="Laser on Bath Mode", ) diff --git a/src/rabifun/analysis.py b/src/rabifun/analysis.py index 711cfc8..c9ba113 100644 --- a/src/rabifun/analysis.py +++ b/src/rabifun/analysis.py @@ -121,7 +121,7 @@ class RingdownPeakData: """The fft frequency array.""" power: np.ndarray - """The power spectrum of the fft.""" + """The normalized power spectrum of the fft.""" peaks: np.ndarray """The indices of the peaks.""" @@ -132,6 +132,9 @@ class RingdownPeakData: peak_info: dict """The information from :any:`scipy.signal.find_peaks`.""" + σ_power: np.ndarray | None = None + """The standard deviation of the power spectrum.""" + peak_widths: np.ndarray | None = None """ The widths of the peaks. @@ -146,10 +149,23 @@ class RingdownPeakData: lorentz_params: list | None = None """The lorentzian fit params to be fed into :any:`lorentzian`.""" + lm_result: lmfit.model.ModelResult | None = None + """The fit result from :any:`lmfit`.""" + + noise_floor: float = 0 + """The noise floor of the spectrum.""" + @property def is_refined(self) -> bool: """Whether the peaks have been refined with :any:`refine_peaks`.""" - return self.peak_widths is not None + return self.lm_result is not None + + def __post_init__(self): + norm = np.max(self.power) + + self.power /= norm + if self.σ_power is not None: + self.σ_power /= norm def find_peaks( @@ -158,6 +174,7 @@ def find_peaks( params: RingdownParams, prominence: float = 0.005, height: float = 0.1, + σ_power: np.ndarray | None = None, ) -> RingdownPeakData: """Determine the peaks of the power spectrum of the ringdown data. @@ -168,7 +185,7 @@ def find_peaks( :param prominence: The prominence (vertical distance of peak from surrounding valleys) of the peaks. :param height: The minimum height of the peaks. - + :param σ_power: The standard deviation of the power spectrum. """ freq_step = freq[1] - freq[0] @@ -195,6 +212,7 @@ def find_peaks( peak_freqs=peak_freqs, peak_info=peak_info, power=power_spectrum, + σ_power=σ_power, ) @@ -210,6 +228,11 @@ def filter_peaks( to_be_deleted: list = [], ): deleted_peaks = [] + if not peaks.is_refined: + return peaks + + peaks = dataclasses.replace(peaks) + for i in reversed(range(len(peaks.peak_freqs))): A, ω0, γ = peaks.lorentz_params[i] Δω0, Δγ = peaks.Δpeak_freqs[i], peaks.Δpeak_widths[i] @@ -221,11 +244,11 @@ def filter_peaks( or A > 5 or Δγ > uncertainty_threshold * params.fΩ_guess ): - np.delete(peaks.peaks, i) - np.delete(peaks.peak_freqs, i) - np.delete(peaks.Δpeak_freqs, i) - np.delete(peaks.peak_widths, i) - np.delete(peaks.Δpeak_widths, i) + peaks.peaks = np.delete(peaks.peaks, i) + peaks.peak_freqs = np.delete(peaks.peak_freqs, i) + peaks.Δpeak_freqs = np.delete(peaks.Δpeak_freqs, i) + peaks.peak_widths = np.delete(peaks.peak_widths, i) + peaks.Δpeak_widths = np.delete(peaks.Δpeak_widths, i) del peaks.lorentz_params[i] deleted_peaks.append(i) @@ -243,10 +266,14 @@ def refine_peaks( params: RingdownParams, uncertainty_threshold: float = 0.2, height_cutoff: float = 0.1, - σ: np.ndarray | None = None, -): +) -> RingdownPeakData: """ - Refine the peak positions and frequencies by fitting Lorentzians. + Refine the peak positions and frequencies by fitting a sum of + Lorentzians. The peaks are filtered according to the + ``height_cutoff``, ``uncertainty_threshold`` and other criteria + and the fit repeated until nothing changes. The results are + stored in a copy of ``peaks``, among them the last successful + :any:`lmfit` fit result. :param peaks: The peak data. :param params: The ringdown parameters. @@ -256,7 +283,7 @@ def refine_peaks( """ if len(peaks.peaks) == 0: - return peaks, None + return peaks peaks = dataclasses.replace(peaks) freqs = peaks.freq @@ -268,8 +295,7 @@ def refine_peaks( scaled_power = power - if σ is None: - σ = np.zeros_like(power) + σ = np.zeros_like(power) if peaks.σ_power is None else peaks.σ_power for i, (A, ω0) in enumerate(zip(peaks.peak_info["peak_heights"], peak_freqs)): model = lmfit.Model(lorentzian, prefix=f"peak_{i}_") @@ -319,7 +345,7 @@ def refine_peaks( peaks.lorentz_params = [None] * len(peaks.peak_freqs) - for i in reversed(range(len(peaks.peak_freqs))): + for i in range(len(peaks.peak_freqs)): peak_prefix = f"peak_{i}_" A, ω0, γ = ( @@ -342,13 +368,16 @@ def refine_peaks( peaks.lorentz_params[i] = A, ω0, γ + peaks.lm_result = lm_result + peaks.noise_floor = lm_result.best_values["offset"] + before_filter = len(peaks.peaks) peaks = filter_peaks(peaks, params, uncertainty_threshold, height_cutoff) if len(peaks.peaks) < before_filter: return refine_peaks(peaks, params, uncertainty_threshold, height_cutoff) - return peaks, lm_result + return peaks class StepType(Enum): diff --git a/src/rabifun/plots.py b/src/rabifun/plots.py index 6a62209..2fc696f 100644 --- a/src/rabifun/plots.py +++ b/src/rabifun/plots.py @@ -213,53 +213,43 @@ def plot_rwa_vs_real_amplitudes(ax, solution_no_rwa, solution_rwa, params, **kwa return no_rwa_lines, rwa_lines -def plot_spectrum_and_peak_info( - ax, peaks: RingdownPeakData, params: RingdownParams, annotate=False -): +def plot_spectrum_and_peak_info(ax, peaks: RingdownPeakData, annotate=False): """Plot the fft spectrum with peaks. :param ax: The axis to plot on. :param peaks: The peak data. - :param params: The ringdown parameters. :param annotate: Whether to annotate the peaks. """ ax.clear() ax.plot( peaks.freq, - peaks.power, + peaks.power - peaks.noise_floor, label="FFT Power", color="C0", - linewidth=0.5, + linewidth=0.1, marker="o", - markersize=2, + markersize=3, ) + fine_freq = np.linspace(peaks.freq[0], peaks.freq[-1], 1000) if annotate: for i, (freq, Δfreq, lorentz) in enumerate( zip(peaks.peak_freqs, peaks.Δpeak_freqs, peaks.lorentz_params) ): - # ax.plot( - # freq, - # max(lorentz[0], 1), - # "x", - # label="Peaks", - # color="C2", - # ) - roundfreq, rounderr = scientific_round(freq, Δfreq) t = ax.text( freq, - ax.get_ylim()[1] * 0.9, + lorentz[0] * 1.1, rf"{i} (${roundfreq}\pm {rounderr}$)", fontsize=20, ) - t.set_bbox(dict(facecolor="white", alpha=1, edgecolor="white")) + t.set_bbox(dict(facecolor="white", alpha=0.8, edgecolor="white")) ax.plot( - peaks.freq, - lorentzian(peaks.freq, *lorentz), + fine_freq, + lorentzian(fine_freq, *lorentz), "--", color="C2", alpha=0.5, @@ -268,12 +258,43 @@ def plot_spectrum_and_peak_info( ax.set_title("FFT Spectrum") ax.set_xlabel("ω [linear]") ax.set_ylabel("Power") - ax.axvline( - params.fω_shift, - color="gray", - linestyle="--", - zorder=-10, - label="Frequency Shift", - ) ax.set_xlim(peaks.freq[0], peaks.freq[-1]) ax.legend() + + +def annotate_ringodown_mode_positions(params: Params, ax): + """ + Add y ticks and vertical lines indicating the theoretical + frequencies and widths of the modes in that can appear in the + ringdown spectrum. + + :param params: The system parameters. + :param ax: The pyplot axis to annotate. + """ + + runtime = RuntimeParams(params) + ax_ticks = ax.twiny() + ax_ticks.sharey(ax) + ax_ticks.set_xticks(runtime.ringdown_frequencies) + ax_ticks.set_xticklabels([mode_name(i, params.N) for i in range(2 * params.N + 2)]) + ax_ticks.set_xlim(ax.get_xlim()) + + for pos, peak_freq in zip(runtime.ringdown_frequencies, runtime.Ωs): + ax.axvline( + pos, + color="black", + alpha=0.5, + linestyle="--", + zorder=-100, + ) + + ax.axvspan( + pos - peak_freq.imag / (2 * np.pi), + pos + peak_freq.imag / (2 * np.pi), + color="black", + alpha=0.05, + linestyle="--", + zorder=-100, + ) + + return ax diff --git a/src/rabifun/system.py b/src/rabifun/system.py index 63f072a..1ad34bb 100644 --- a/src/rabifun/system.py +++ b/src/rabifun/system.py @@ -125,7 +125,7 @@ class Params: """ return n / self.Ω - def lifetimes(self, n: float): + def lifetimes(self, n: float) -> float: """ Returns the number of lifetimes of the system that correspond to `n` cycles. diff --git a/src/ringfit/utils.py b/src/ringfit/utils.py index 5e3fa38..038dd0d 100644 --- a/src/ringfit/utils.py +++ b/src/ringfit/utils.py @@ -46,3 +46,59 @@ def smoothe_signal( window = int(window_size / time_step) return uniform_filter1d(signal, window) + + +class WelfordAggregator: + """A class to aggregate values using the Welford algorithm. + + The Welford algorithm is an online algorithm to calculate the mean + and variance of a series of values. + + The aggregator keeps track of the number of samples the mean and + the variance. Aggregation of identical values is prevented by + checking the sample index. Tracking can be disabled by setting + the initial index to ``None``. + + See also the `Wikipedia article + `_. + + :param first_value: The first value to aggregate. + """ + + __slots__ = ["n", "mean", "_m_2"] + + def __init__(self, first_value: np.ndarray): + self.n: int = 1 + self.mean: np.ndarray = first_value + self._m_2 = np.zeros_like(first_value) + + def update(self, new_value: np.ndarray): + """Updates the aggregator with a new value.""" + + self.n += 1 + delta = new_value - self.mean + self.mean += delta / self.n + delta2 = new_value - self.mean + self._m_2 += np.abs(delta) * np.abs(delta2) + + @property + def sample_variance(self) -> np.ndarray: + """ + The empirical sample variance. (:math:`\sqrt{N-1}` + normalization.) + """ + + if self.n == 1: + return np.zeros_like(self.mean) + + return self._m_2 / (self.n - 1) + + @property + def ensemble_variance(self) -> np.ndarray: + """The ensemble variance.""" + return self.sample_variance / self.n + + @property + def ensemble_std(self) -> np.ndarray: + """The ensemble standard deviation.""" + return np.sqrt(self.ensemble_variance)