abstract/automate some of the ringdown analysis

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)

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]))

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()