first shot at the spectral matching...

flake.lock

@@ -59,11 +59,11 @@
     },
     "nixpkgs": {
       "locked": {
-        "lastModified": 1716792620,
-        "narHash": "sha256-wQmXzee/veETYJv93TkRYsAQkEdt2QYCJeJil5SrJfg=",
+        "lastModified": 1718083092,
+        "narHash": "sha256-EQsPXycAbmby4meQUNLYfFaGOiqz2J9AlwFRV4UiHnY=",
         "owner": "NixOS",
         "repo": "nixpkgs",
-        "rev": "7d7cf1590c05d799745bf456f2b95b798f48d3bb",
+        "rev": "aa1ebdaf49a606e21c06e0f6ed7aece9a41831c3",
         "type": "github"
       },
       "original": {
@@ -84,11 +84,11 @@
         "treefmt-nix": "treefmt-nix"
       },
       "locked": {
-        "lastModified": 1716813403,
-        "narHash": "sha256-9+G8tEOh3QkjSUV2UMC+TpvzKOR8IUFlkJJTMpVQMkc=",
+        "lastModified": 1717965289,
+        "narHash": "sha256-62VsS1MvwcsyYWnxxOD9rX5yqDUvMXH7qE3Kj8HECYE=",
         "owner": "nix-community",
         "repo": "poetry2nix",
-        "rev": "12599ecaa9ec641c29dc8fd07f8267b23874bf3a",
+        "rev": "304f8235fb0729fd48567af34fcd1b58d18f9b95",
         "type": "github"
       },
       "original": {
@@ -156,11 +156,11 @@
         ]
       },
       "locked": {
-        "lastModified": 1715940852,
-        "narHash": "sha256-wJqHMg/K6X3JGAE9YLM0LsuKrKb4XiBeVaoeMNlReZg=",
+        "lastModified": 1717278143,
+        "narHash": "sha256-u10aDdYrpiGOLoxzY/mJ9llST9yO8Q7K/UlROoNxzDw=",
         "owner": "numtide",
         "repo": "treefmt-nix",
-        "rev": "2fba33a182602b9d49f0b2440513e5ee091d838b",
+        "rev": "3eb96ca1ae9edf792a8e0963cc92fddfa5a87706",
         "type": "github"
       },
       "original": {

scripts/experiments/.#007_generate_fake_analysis_data.py

@@ -0,0 +1 @@
+hiro@Phillips.106575:1718131872

scripts/experiments/006_test_new_analysis.py

@@ -16,7 +16,7 @@ path = "../../data/22_05_24/ringdown_try_2"
 scan = ScanData.from_dir(path)
 
 
-# %% Set Window
+# %% Extract Rough Peaks
 window = tuple(
     np.array([0.016244684251065847 + 0.000002, 0.016248626903395593 + 49e-5])
     + 8e-3
@@ -30,5 +30,5 @@ 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 FSR and hybridization amplitude
 extract_Ω_δ(peak_info, ringdown_params)

scripts/experiments/007_generate_fake_analysis_data.py

@@ -0,0 +1,104 @@
+from rabifun.system import *
+from rabifun.plots import *
+from rabifun.utilities import *
+from rabifun.analysis import *
+from ringfit.data import ScanData
+import functools
+
+# %% interactive
+
+
+@functools.lru_cache()
+def make_params_and_solve(
+    total_lifetimes, eom_off_lifetime, ω_c=0.1 / 2, N=10, g_0=1 / 4
+):
+    """
+    Make a set of parameters for the system with the current
+    best-known settings.
+    """
+
+    Ω = 13
+
+    params = Params(
+        η=0.5,
+        Ω=Ω,
+        δ=1 / 4,
+        ω_c=ω_c,
+        g_0=g_0,
+        laser_detuning=0,
+        N=2 * N + 2,
+        N_couplings=N,
+        measurement_detuning=Ω * (3),
+        α=0,
+        rwa=False,
+        flat_energies=False,
+        correct_lamb_shift=0,
+        laser_off_time=0,
+    )
+
+    params.laser_off_time = params.lifetimes(eom_off_lifetime)
+    params.drive_off_time = params.lifetimes(eom_off_lifetime)
+
+    params.drive_override = (
+        np.array([params.Ω, params.Ω * (1 - params.δ), params.δ * 2 * params.Ω]),
+        np.repeat(params.Ω, 3),
+    )
+
+    t = time_axis(params, lifetimes=total_lifetimes, resolution=0.01)
+    solution = solve(t, params)
+    return params, t, solution
+
+
+def generate_phase_one_data():
+    """
+    This generates some intensity data for phase one of the
+    calibration protocol, where the aim is to extract the resonator
+    spectrum.
+    """
+
+    total_lifetimes = 50
+    eom_off_lifetime = total_lifetimes * 2 / 3
+    fluct_size = 0.05
+
+    params, t, solution = make_params_and_solve(
+        total_lifetimes, eom_off_lifetime, N=10, g_0=10
+    )
+    signal = output_signal(t, solution.y, params)
+
+    rng = np.random.default_rng(seed=0)
+    signal += np.mean(abs(signal)) * rng.standard_normal(len(signal)) * fluct_size * 100
+
+    fig = make_figure()
+    ax_realtime, ax_spectrum = fig.subplots(2, 1)
+
+    ax_realtime.plot(t[::500], signal[::500])
+    ax_realtime.axvline(params.laser_off_time, color="black", linestyle="--")
+
+    # now we plot the power spectrum
+    window = (float(params.laser_off_time or 0), float(t[-1]))
+
+    ringdown_params = RingdownParams(
+        fω_shift=params.measurement_detuning,
+        mode_window=(params.N + 2, params.N + 2),
+        fΩ_guess=params.Ω * (1 + rng.standard_normal() * fluct_size),
+        fδ_guess=params.Ω * params.δ * (1 + rng.standard_normal() * fluct_size),
+        η_guess=params.η * (1 + rng.standard_normal() * fluct_size),
+    )
+
+    scan = ScanData(np.ones_like(signal), signal, t)
+    peak_info = find_peaks(scan, ringdown_params, window, prominence=0.1)
+    peak_info = refine_peaks(peak_info, ringdown_params)
+    plot_spectrum_and_peak_info(ax_spectrum, peak_info, ringdown_params)
+
+    Ω, ΔΩ, ladder = extract_Ω_δ(peak_info, ringdown_params, threshold=0.2)
+
+    for index, type, _ in ladder:
+        freq_index = peak_info.peaks[index]
+        ax_spectrum.plot(
+            peak_info.freq[freq_index],
+            peak_info.normalized_power[freq_index],
+            "o" if type == 0 else "*",
+            color="C4" if type == 0 else "C5",
+            label=type,
+        )
+    return Ω, ΔΩ, ladder

src/rabifun/analysis.py

@@ -30,8 +30,6 @@ def fourier_transform(
         t = t[mask]
         signal = signal[mask]  # * scipy.signal.windows.hamming(len(t))
 
-        #
-
     freq = scipy.fft.rfftfreq(len(t), t[2] - t[1])
     fft = scipy.fft.rfft(signal, norm="forward", workers=os.cpu_count())
 
@@ -94,7 +92,7 @@ class RingdownParams:
     @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
+        return max(self.fω_shift - self.mode_window[0] * self.fΩ_guess, 0)
 
     @property
     def high_cutoff(self) -> float:
@@ -204,7 +202,8 @@ def refine_peaks(peaks: RingdownPeakData, params: RingdownParams):
     Δwidths = []
 
     window = params.η_guess * 3
-    for peak_freq in peak_freqs:
+    deleted_peaks = []
+    for i, peak_freq in enumerate(peak_freqs):
         mask = (freqs > peak_freq - window) & (freqs < peak_freq + window)
         windowed_freqs = freqs[mask]
         windowed_power = power[mask]
@@ -215,6 +214,7 @@ def refine_peaks(peaks: RingdownPeakData, params: RingdownParams):
             [np.inf, windowed_freqs[-1], np.inf, 1],
         )
 
+        try:
             popt, pcov = scipy.optimize.curve_fit(
                 lorentzian,
                 windowed_freqs,
@@ -229,6 +229,13 @@ def refine_peaks(peaks: RingdownPeakData, params: RingdownParams):
 
             new_widths.append(popt[2])
             Δwidths.append(perr[2])
+        except:
+            deleted_peaks.append(i)
+
+    peaks.peaks = np.delete(peaks.peaks, deleted_peaks)
+    for key, value in peaks.peak_info.items():
+        if isinstance(value, np.ndarray):
+            peaks.peak_info[key] = np.delete(value, deleted_peaks)
 
     peaks.peak_freqs = np.array(new_freqs)
     peaks.Δpeak_freqs = np.array(Δfreqs)
@@ -238,11 +245,15 @@ def refine_peaks(peaks: RingdownPeakData, params: RingdownParams):
     return peaks
 
 
+import matplotlib.pyplot as plt
+
+
 def extract_Ω_δ(
     peaks: RingdownPeakData, params: RingdownParams, threshold: float = 0.1
 ):
     """
-    Extract the FSR and mode splitting from the peaks.
+    Extract the FSR and mode splitting from the peaks.  The threshold
+    regulates the maximum allowed deviation from the expected FSR.
 
     :param peaks: The peak data.
     :param params: The ringdown parameters.
@@ -251,12 +262,101 @@ def extract_Ω_δ(
     if not peaks.is_refined:
         raise ValueError("Peaks must be refined.")
 
+    peak_indices = peaks.peaks
+    peak_freqs = peaks.peak_freqs
+    Δpeak_freqs = (
+        np.zeros_like(peak_freqs) if peaks.Δpeak_freqs is None else peaks.Δpeak_freqs
+    )
+
     Ω_guess = params.fΩ_guess
     δ_guess = params.fδ_guess
 
-    all_diff = np.abs(peaks.peak_freqs[:, None] - peaks.peak_freqs[None, :])
+    # first step: we extract the most common frequency spacing
+    all_diff = np.abs(peak_freqs[:, None] - peak_freqs[None, :])
     all_diff = np.triu(all_diff)
+    all_ΔΩ = np.abs(Δpeak_freqs[:, None] ** 2 + Δpeak_freqs[None, :] ** 2)
 
-    bath_mask = (all_diff - Ω_guess) / Ω_guess < threshold
+    bath_mask = (np.abs((all_diff - Ω_guess)) / Ω_guess < threshold) & (all_diff > 0)
+    candidates = all_diff[bath_mask]
+    Δcandidates = all_ΔΩ[bath_mask]
 
-    return np.mean((all_diff[bath_mask]))
+    Ω = np.mean(candidates)
+    ΔΩ = max(np.sqrt(np.sum(Δcandidates**2)) / len(candidates), np.std(candidates))
+
+    if np.isnan(Ω):
+        raise ValueError("No FSR found")
+
+    # second step: we walk through the peaks and label them as for the
+    total_peaks = len(peak_indices)
+    peak_pool = list(range(total_peaks))
+
+    ladders = []
+
+    current_peak = 0
+    current_ladder = []
+
+    possible_diffs = np.array([Ω, Ω - δ_guess, 2 * δ_guess])
+
+    while len(peak_pool) > 1:
+        if current_peak == len(peak_pool):
+            if current_ladder:
+                ladders.append(current_ladder)
+            current_ladder = []
+            current_peak = peak_pool[0]
+
+        filtered_freqs = peak_freqs[peak_pool]
+
+        diffs = filtered_freqs - filtered_freqs[current_peak]
+        diffs[diffs <= 0] = np.inf
+
+        diffs = (
+            np.abs(diffs[:, None] - possible_diffs[None, :]) / possible_diffs[None, :]
+        )
+
+        min_coords = np.unravel_index(np.argmin(diffs), diffs.shape)
+        min_diff = diffs[min_coords]
+
+        mode_index, step_type = min_coords
+
+        if min_diff > threshold:
+            if current_ladder:
+                ladders.append(current_ladder)
+            current_ladder = []
+            del peak_pool[current_peak]
+            continue
+
+        current_ladder.append((peak_pool[current_peak], step_type, min_diff))
+        del peak_pool[current_peak]
+        current_peak = mode_index - 1  # we have deleted one peak
+
+    if current_ladder:
+        ladders.append(current_ladder)
+
+    # we want at least one bath mode before the A site
+    ladders = list(
+        filter(
+            lambda ladder: sum([1 if x[1] == 1 else 0 for x in ladder]) > 0,
+            ladders,
+        )
+    )
+
+    invalid = []
+    for lad_index, ladder in enumerate(ladders):
+        length = len(ladder)
+        for i, elem in enumerate(ladder):
+            if elem[1] == 1:
+                if (i + 2) >= length or not (
+                    ladder[i + 1][1] == 2 and ladder[i + 2][1] == 1
+                ):
+                    invalid.append(lad_index)
+                break
+
+    ladders = [ladder for i, ladder in enumerate(ladders) if i not in invalid]
+    costs = [sum([x[2] for x in ladder]) / len(ladder) for ladder in ladders]
+
+    if len(costs) == 0:
+        raise ValueError("No matching modes found.")
+
+    best = ladders[np.argmin(costs)]
+
+    return Ω, ΔΩ, best

src/rabifun/system.py

@@ -67,6 +67,17 @@ class Params:
     correct_lamb_shift: float = 1
     """Whether to correct for the Lamb shift by tweaking the detuning."""
 
+    drive_override: tuple[np.ndarray, np.ndarray] | None = None
+    """
+    Override the drive frequencies (first array, linear frequency) and
+    amplitudes (second array, linear frequency units).
+
+    In this case the detunings of the rotating frame will be zeroed
+    and the parameters :any:`α` and :any:`ω_c` will be ignored.
+
+    The drive strength is normalized to :any:`g_0`.
+    """
+
     def __post_init__(self):
         if self.N_couplings > self.N:
             raise ValueError("N_couplings must be less than or equal to N.")
@@ -74,6 +85,18 @@ class Params:
         if self.initial_state and len(self.initial_state) != 2 * self.N + 2:
             raise ValueError("Initial state must have length 2N + 2.")
 
+        if self.drive_override is not None:
+            if len(self.drive_override) != 2:
+                raise ValueError("Drive override must be a tuple of two arrays.")
+
+            if len(self.drive_override[0]) != len(self.drive_override[1]):
+                raise ValueError(
+                    "Drive frequencies and amplitudes must have the same length."
+                )
+
+            if self.rwa:
+                raise ValueError("Drive override is not compatible with the RWA.")
+
     def periods(self, n: float):
         """
         Returns the number of periods of the system that correspond to
@@ -116,6 +139,13 @@ class RuntimeParams:
             drive_frequencies_and_amplitudes(params)
         )  # linear frequencies!
 
+        if params.drive_override is not None:
+            self.drive_frequencies = params.drive_override[0]
+            self.g = 2 * np.pi * params.drive_override[1]
+            a_shift = 0
+            self.detunings *= 0
+            self.g /= params.g_0 * np.sqrt(np.sum(self.g**2))
+
         self.g *= 2 * np.pi
         self.Ωs = Ωs
         self.ε = (
@@ -136,9 +166,12 @@ class RuntimeParams:
         self.bath_ε = 2 * np.pi * self.detunings - 1j * params.η / 2
         self.a_shift = 2 * np.pi * a_shift
         self.detuned_Ωs = freqs - self.ε.real
+
         self.RWA_H = np.zeros((2 * params.N + 2, 2 * params.N + 2), np.complex128)
+        if not params.drive_override:
             self.RWA_H[1, 2 : 2 + params.N_couplings] = self.g
             self.RWA_H[2 : 2 + params.N_couplings, 1] = np.conj(self.g)
+
         self.detuning_matrix = self.detuned_Ωs[:, None] - self.detuned_Ωs[None, :]
 
     def __repr__(self):
@@ -245,9 +278,9 @@ def make_righthand_side(runtime_params: RuntimeParams, params: Params):
             differential[0:2] += laser[:2] / np.sqrt(2)
             differential[2:] += laser[2:]
 
-        # if params.rwa:
-        #     differential[0] = 0
-        #     differential[2 + params.N_couplings :] = 0
+        if params.rwa:
+            differential[0] = 0
+            differential[2 + params.N_couplings :] = 0
 
         return -1j * differential
 
@@ -314,7 +347,7 @@ def output_signal(t: np.ndarray, amplitudes: np.ndarray, params: Params):
     """
 
     runtime = RuntimeParams(params)
-    rotating = amplitudes * np.exp(-1j * runtime.detuned_Ωs * t)
+    rotating = amplitudes * np.exp(-1j * runtime.detuned_Ωs[:, None] * t)
 
     return (
         np.sum(rotating, axis=0)