new recursive algorithm works a little better

scripts/experiments/.#007_generate_fake_analysis_data.py

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

scripts/experiments/007_generate_fake_analysis_data.py

@@ -61,18 +61,20 @@ def generate_phase_one_data():
     fluct_size = 0.05
 
     params, t, solution = make_params_and_solve(
-        total_lifetimes, eom_off_lifetime, N=10, g_0=10
+        total_lifetimes, eom_off_lifetime, N=3, 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
+    signal += np.mean(abs(signal)) * rng.standard_normal(len(signal)) * fluct_size * 50
 
     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="--")
+    ax_realtime.set_xlabel("Time")
+    ax_realtime.set_ylabel("Intensity")
 
     # now we plot the power spectrum
     window = (float(params.laser_off_time or 0), float(t[-1]))
@@ -90,15 +92,18 @@ def generate_phase_one_data():
     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)
+    Ω, ΔΩ, δ, Δδ, ladder = extract_Ω_δ(
+        peak_info, ringdown_params, Ω_threshold=0.1, ladder_threshold=0.1
+    )
 
-    for index, type, _ in ladder:
+    for index, type in ladder:
         freq_index = peak_info.peaks[index]
+        print(type, type is StepType.BATH)
         ax_spectrum.plot(
             peak_info.freq[freq_index],
             peak_info.normalized_power[freq_index],
-            "o" if type == 0 else "*",
+            "o" if type.value == StepType.BATH.value else "*",
-            color="C4" if type == 0 else "C5",
+            color="C4" if type.value == StepType.BATH.value else "C5",
             label=type,
         )
-    return Ω, ΔΩ, ladder
+    return Ω, ΔΩ, δ, Δδ, ladder

src/rabifun/analysis.py

@@ -6,6 +6,7 @@ import dataclasses
 from dataclasses import dataclass
 
 from scipy.optimize import Bounds
+from enum import Enum
 
 
 def fourier_transform(
@@ -257,8 +258,20 @@ def refine_peaks(
 import matplotlib.pyplot as plt
 
 
+class StepType(Enum):
+    BATH = 0
+    BATH_TO_A = 1
+    A_TO_A = 2
+
+
 def extract_Ω_δ(
-    peaks: RingdownPeakData, params: RingdownParams, threshold: float = 0.1
+    peaks: RingdownPeakData,
+    params: RingdownParams,
+    Ω_threshold: float = 0.1,
+    ladder_threshold: float = 0.1,
+    bifurcations: int = 3,
+    start_peaks: int = 2,
+    min_length: int = 4,
 ):
     """
     Extract the FSR and mode splitting from the peaks.  The threshold
@@ -266,6 +279,14 @@ def extract_Ω_δ(
 
     :param peaks: The peak data.
     :param params: The ringdown parameters.
+    :param Ω_threshold: The maximum allowed relative deviation from
+        the expected FSR for the rough search.
+    :param ladder_threshold: The maximum allowed relative deviation
+        from the expected step sizes for the ladder search.
+    :param bifurcations: The number of bifurcations to consider in the
+        ladder search, i.e. how many possible new steps are accepted at each step.
+    :param start_peaks: The number of peaks to start the ladder search (from the left).
+    :param min_length: The minimum length of a ladder to be considered valid.
     """
 
     if not peaks.is_refined:
@@ -285,7 +306,7 @@ def extract_Ω_δ(
     all_diff = np.triu(all_diff)
     all_ΔΩ = np.abs(Δpeak_freqs[:, None] ** 2 + Δpeak_freqs[None, :] ** 2)
 
-    bath_mask = (np.abs((all_diff - Ω_guess)) / Ω_guess < threshold) & (all_diff > 0)
+    bath_mask = (np.abs((all_diff - Ω_guess)) / Ω_guess < Ω_threshold) & (all_diff > 0)
     candidates = all_diff[bath_mask]
     Δcandidates = all_ΔΩ[bath_mask]
 
@@ -306,66 +327,128 @@ def extract_Ω_δ(
 
     possible_diffs = np.array([Ω, Ω - δ_guess, 2 * δ_guess])
 
-    while len(peak_pool) > 1:
+    # entry in a ladder: (peak_index, step_type)
-        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]
+    def walk_ladder(
+        current_peak,
+        last_type=StepType.BATH,
+        second_last_type=StepType.BATH,
+        bifurcations=3,
+    ):
+        if current_peak == total_peaks - 1:
+            return [], []
 
-        diffs = filtered_freqs - filtered_freqs[current_peak]
+        match last_type:
+            case StepType.BATH:
+                allowed_steps = [StepType.BATH, StepType.BATH_TO_A]
+            case StepType.BATH_TO_A:
+                match second_last_type:
+                    case StepType.BATH:
+                        allowed_steps = [StepType.A_TO_A]
+
+                    case StepType.A_TO_A:
+                        allowed_steps = [StepType.BATH]
+            case StepType.A_TO_A:
+                allowed_steps = [StepType.BATH_TO_A]
+
+        allowed_step_indices = [x.value for x in allowed_steps]
+        allowed_possible_diffs = possible_diffs[allowed_step_indices]
+
+        diffs = peak_freqs - peak_freqs[current_peak]
         diffs[diffs <= 0] = np.inf
 
         diffs = (
-            np.abs(diffs[:, None] - possible_diffs[None, :]) / possible_diffs[None, :]
+            np.abs(diffs[:, None] - allowed_possible_diffs[None, :])
+            / allowed_possible_diffs[None, :]
         )
 
-        min_coords = np.unravel_index(np.argmin(diffs), diffs.shape)
+        ladders = []
+        costs = []
+        min_candidates = np.argpartition(diffs.flatten(), bifurcations)[:bifurcations]
+
+        for min_coords in min_candidates:
+            min_coords = np.unravel_index(min_coords, diffs.shape)
             min_diff = diffs[min_coords]
+            peak_index, step_type = min_coords
 
-        mode_index, step_type = min_coords
+            step_type = StepType(allowed_step_indices[step_type])
 
-        if min_diff > threshold:
+            this_rung = [(current_peak, step_type)]
-            if current_ladder:
+            if min_diff < ladder_threshold + (ΔΩ / Ω):
-                ladders.append(current_ladder)
+                new_ladders, new_costs = walk_ladder(
-            current_ladder = []
+                    peak_index,
-            del peak_pool[current_peak]
+                    step_type,
-            continue
+                    last_type,
-
+                    bifurcations,
-        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,
                 )
+
+                if len(new_ladders) == 0:
+                    ladders.append(this_rung)
+                    costs.append(min_diff)
+                else:
+                    for new_ladder, new_cost in zip(new_ladders, new_costs):
+                        if new_ladder:
+                            ladders.append(this_rung + new_ladder)
+                            costs.append(new_cost + min_diff)
+
+        return ladders, costs
+
+    ladders = []
+    costs = []
+
+    for start_index in range(min(total_peaks, start_peaks)):
+        new_ladders, new_costs = walk_ladder(
+            start_index, StepType.BATH, StepType.BATH, bifurcations
         )
+        ladders += new_ladders
+        costs += new_costs
 
     invalid = []
     for lad_index, ladder in enumerate(ladders):
-        length = len(ladder)
+        if len(ladder) < min_length:
-        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)
+            continue
+
+        is_invalid = True
+        for elem in ladder:
+            if elem[1] == StepType.A_TO_A:
+                is_invalid = False
                 break
 
+        if is_invalid:
+            invalid.append(lad_index)
+
     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]
+    costs = [cost for i, cost in enumerate(costs) if i not in invalid]
 
-    if len(costs) == 0:
+    costs = [cost / len(ladder) for cost, ladder in zip(costs, ladders)]
-        raise ValueError("No matching modes found.")
 
-    best = ladders[np.argmin(costs)]
+    best = np.argmin(costs)
+    best_ladder = ladders[best]
 
-    return Ω, ΔΩ, best
+    Ωs = []
+    δs = []
+    Ω_m_δs = []
+
+    for (i, (begin_index, begin_type)), (end_index, _) in zip(
+        enumerate(best_ladder[:-1]), best_ladder[1:]
+    ):
+        match begin_type:
+            case StepType.BATH:
+                Ωs.append(peak_freqs[end_index] - peak_freqs[begin_index])
+            case StepType.BATH_TO_A:
+                Ω_m_δs.append(peak_freqs[end_index] - peak_freqs[begin_index])
+                if i < len(best_ladder) - 2:
+                    Ωs.append(
+                        (peak_freqs[best_ladder[i + 2][0]] - peak_freqs[begin_index])
+                        / 2
+                    )
+            case StepType.A_TO_A:
+                δs.append((peak_freqs[end_index] - peak_freqs[begin_index]) / 2)
+
+    Ω = np.mean(Ωs)
+    ΔΩ = np.std(Ωs)
+    δ = np.mean(δs)
+    Δδ = np.std(δs)
+
+    return Ω, ΔΩ, δ, Δδ, best_ladder