From 580753cb1253ceda892697e925af1e742664351d Mon Sep 17 00:00:00 2001 From: Valentin Boettcher Date: Thu, 6 Jun 2024 13:36:34 -0400 Subject: [PATCH] abstract/automate some of the ringdown analysis --- scripts/experiments/006_test_new_analysis.py | 34 ++++ src/rabifun/analysis.py | 198 +++++++++++++++++++ src/rabifun/plots.py | 35 +++- 3 files changed, 266 insertions(+), 1 deletion(-) create mode 100644 scripts/experiments/006_test_new_analysis.py diff --git a/scripts/experiments/006_test_new_analysis.py b/scripts/experiments/006_test_new_analysis.py new file mode 100644 index 0000000..9aab7f2 --- /dev/null +++ b/scripts/experiments/006_test_new_analysis.py @@ -0,0 +1,34 @@ +from ringfit import data +import matplotlib.pyplot as plt +from ringfit.data import * +from ringfit.plotting import * +from ringfit.fit import * +from rabifun.analysis import * +from rabifun.plots import * + +# %% setup figure +fig = make_figure("calibration") +ax_spectrum = fig.subplots(1, 1) + + +# %% load data +path = "../../data/22_05_24/ringdown_try_2" +scan = ScanData.from_dir(path) + + +# %% Set Window +window = tuple( + np.array([0.016244684251065847 + 0.000002, 0.016248626903395593 + 49e-5]) + + 8e-3 + - 12e-7 +) + +ringdown_params = RingdownParams(fω_shift=10e4, mode_window=(0, 50)) +peak_info = find_peaks(scan, ringdown_params, window, prominence=0.008) + + +peak_info = refine_peaks(peak_info, ringdown_params) +plot_spectrum_and_peak_info(ax_spectrum, peak_info, ringdown_params) + +# extract +extract_Ω_δ(peak_info, ringdown_params) diff --git a/src/rabifun/analysis.py b/src/rabifun/analysis.py index 30d08d5..da71376 100644 --- a/src/rabifun/analysis.py +++ b/src/rabifun/analysis.py @@ -1,6 +1,11 @@ import numpy as np import scipy import os +from ringfit.data import ScanData +import dataclasses +from dataclasses import dataclass + +from scipy.optimize import Bounds def fourier_transform( @@ -62,3 +67,196 @@ def complex_lorentzian(ω, A, ω0, γ): """ return A * (γ / 2) * 1 / (-1j * (ω - ω0) - (γ / (4 * np.pi))) + + +############################################################################### +# Ringdown Calibration # +############################################################################### + + +@dataclass +class RingdownParams: + fω_shift: float = 200e6 + """The laser frequency shift in Hz (linear).""" + + fΩ_guess: float = 13e6 + """The best guess of the FSR in Hz (linear).""" + + fδ_guess: float = 3e6 + """The best guess of the mode splitting in Hz (linear).""" + + η_guess: float = 0.5e6 + """The best guess of the worst decay rate in Hz (no 2πs).""" + + mode_window: tuple[int, int] = (10, 10) + """How many FSRs of frequency to consider around :any`fω_shift`.""" + + @property + def low_cutoff(self) -> float: + """The low cutoff frequency of the ringdown spectrum fft.""" + return self.fω_shift - self.mode_window[0] * self.fΩ_guess + + @property + def high_cutoff(self) -> float: + """The high cutoff frequency of the ringdown spectrum fft.""" + return self.fω_shift + self.mode_window[1] * self.fΩ_guess + + +@dataclass +class RingdownPeakData: + freq: np.ndarray + """The fft frequency array.""" + + fft: np.ndarray + """The fft amplitudes.""" + + normalized_power: np.ndarray + """The normalized power spectrum of the fft.""" + + peaks: np.ndarray + """The indices of the peaks.""" + + peak_freqs: np.ndarray + """The frequencies of the peaks.""" + + peak_info: dict + """The information from :any:`scipy.signal.find_peaks`.""" + + peak_widths: np.ndarray | None = None + """ + The widths of the peaks. + """ + + Δpeak_freqs: np.ndarray | None = None + """The uncertainty in the peak frequencies.""" + + Δpeak_widths: np.ndarray | None = None + """The uncertainty in the peak widths.""" + + @property + def is_refined(self) -> bool: + """Whether the peaks have been refined with :any:`refine_peaks`.""" + return self.peak_widths is not None + + +def find_peaks( + data: ScanData, + params: RingdownParams, + window: tuple[float, float], + prominence: float = 0.005, +): + """Determine the peaks of the normalized power spectrum of the + ringdown data. + + :param data: The oscilloscope data. + :param params: The ringdown parameters, see :any:`RingdownParams`. + :param window: The time window to consider. (to be automated) + :param prominence: The prominence (vertical distance of peak from + surrounding valleys) of the peaks. + """ + + freq, fft, t = fourier_transform( + data.time, + data.output, + window=window, + low_cutoff=params.low_cutoff, + high_cutoff=params.high_cutoff, + ret_time=True, + ) + freq_step = freq[1] - freq[0] + + power = np.abs(fft) ** 2 + power /= power.max() + + distance = params.fδ_guess / 2 / freq_step + peaks, peak_info = scipy.signal.find_peaks( + power, distance=distance, wlen=distance // 4, prominence=prominence + ) + + peak_freqs = freq[peaks] + + return RingdownPeakData( + freq=freq, + fft=fft, + peaks=peaks, + peak_freqs=peak_freqs, + peak_info=peak_info, + normalized_power=power, + ) + + +def refine_peaks(peaks: RingdownPeakData, params: RingdownParams): + """ + Refine the peak positions and frequencies by fitting Lorentzians. + + :param peaks: The peak data. + :param params: The ringdown parameters. + """ + + peaks = dataclasses.replace(peaks) + freqs = peaks.freq + peak_freqs = peaks.peak_freqs + power = peaks.normalized_power + + new_freqs = [] + new_widths = [] + Δfreqs = [] + Δwidths = [] + + window = params.η_guess * 3 + for peak_freq in peak_freqs: + mask = (freqs > peak_freq - window) & (freqs < peak_freq + window) + windowed_freqs = freqs[mask] + windowed_power = power[mask] + + p0 = [1, peak_freq, params.η_guess, 0] + bounds = ( + [0, windowed_freqs[0], 0, 0], + [np.inf, windowed_freqs[-1], np.inf, 1], + ) + + popt, pcov = scipy.optimize.curve_fit( + lorentzian, + windowed_freqs, + windowed_power, + p0=p0, + bounds=bounds, + ) + perr = np.sqrt(np.diag(pcov)) + + new_freqs.append(popt[1]) + Δfreqs.append(perr[1]) + + new_widths.append(popt[2]) + Δwidths.append(perr[2]) + + peaks.peak_freqs = np.array(new_freqs) + peaks.Δpeak_freqs = np.array(Δfreqs) + peaks.peak_widths = np.array(new_widths) + peaks.Δpeak_widths = np.array(Δwidths) + + return peaks + + +def extract_Ω_δ( + peaks: RingdownPeakData, params: RingdownParams, threshold: float = 0.1 +): + """ + Extract the FSR and mode splitting from the peaks. + + :param peaks: The peak data. + :param params: The ringdown parameters. + """ + + if not peaks.is_refined: + raise ValueError("Peaks must be refined.") + + Ω_guess = params.fΩ_guess + δ_guess = params.fδ_guess + + all_diff = np.abs(peaks.peak_freqs[:, None] - peaks.peak_freqs[None, :]) + all_diff = np.triu(all_diff) + + bath_mask = (all_diff - Ω_guess) / Ω_guess < threshold + + return np.mean((all_diff[bath_mask])) diff --git a/src/rabifun/plots.py b/src/rabifun/plots.py index ca5ce64..205f2ee 100644 --- a/src/rabifun/plots.py +++ b/src/rabifun/plots.py @@ -1,3 +1,4 @@ +from networkx import is_aperiodic from plot_utils import * from .system import ( Params, @@ -8,7 +9,7 @@ from .system import ( uncoupled_mode_indices, correct_for_decay, ) -from .analysis import fourier_transform +from .analysis import fourier_transform, RingdownPeakData, RingdownParams import matplotlib.pyplot as plt import numpy as np @@ -210,3 +211,35 @@ def plot_rwa_vs_real_amplitudes(ax, solution_no_rwa, solution_rwa, params, **kwa params.rwa = original_rwa return no_rwa_lines, rwa_lines + + +def plot_spectrum_and_peak_info(ax, peaks: RingdownPeakData, params: RingdownParams): + """Plot the fft spectrum with peaks. + + :param ax: The axis to plot on. + :param peaks: The peak data. + :param params: The ringdown parameters. + """ + + ax.clear() + ax.plot(peaks.freq, peaks.normalized_power, label="FFT Power", color="C0") + ax.plot( + peaks.peak_freqs, + peaks.normalized_power[peaks.peaks], + "x", + label="Peaks", + color="C2", + ) + ax.set_title("FFT Spectrum") + ax.set_xlabel("ω [linear]") + ax.set_ylabel("Power") + ax.set_ylim(0, 1.1) + ax.axvline( + params.fω_shift, + color="gray", + linestyle="--", + zorder=-10, + label="Frequency Shift", + ) + + ax.legend()