From 1df783ab2d2e69a1dab6e091f121c41253121f79 Mon Sep 17 00:00:00 2001
From: Valentin Boettcher <hiro@protagon.space>
Date: Wed, 7 Aug 2024 17:19:28 -0400
Subject: [PATCH] begin implementing new calibration

---
 scripts/experiments/012_all_to_all.py   | 201 ++++++++++++++++++++++++
 scripts/experiments/013_step_through.py | 201 ++++++++++++++++++++++++
 src/rabifun/system.py                   |  17 +-
 3 files changed, 411 insertions(+), 8 deletions(-)
 create mode 100644 scripts/experiments/012_all_to_all.py
 create mode 100644 scripts/experiments/013_step_through.py

diff --git a/scripts/experiments/012_all_to_all.py b/scripts/experiments/012_all_to_all.py
new file mode 100644
index 0000000..17e7de5
--- /dev/null
+++ b/scripts/experiments/012_all_to_all.py
@@ -0,0 +1,201 @@
+from rabifun.system import *
+from rabifun.plots import *
+from rabifun.utilities import *
+from rabifun.analysis import *
+from ringfit.data import ScanData
+import functools
+
+from scipy.ndimage import rotate
+
+# %% interactive
+
+
+def make_params_and_solve(
+    total_lifetimes,
+    eom_off_lifetime,
+    ω_c=0.1 / 2,
+    N=10,
+    g_0=1 / 4,
+    small_loop_detuning=0,
+    laser_detuning=0,
+    η=1.0,
+    η_hybrid=1.0,
+    drive_detuning=0,
+):
+    """
+    Make a set of parameters for the system with the current
+    best-known settings.
+    """
+
+    Ω = 13
+
+    params = Params(
+        η=η,
+        η_hybrid=η_hybrid,
+        Ω=Ω,
+        δ=1 / 4,
+        ω_c=ω_c,
+        g_0=g_0,
+        laser_detuning=laser_detuning,
+        N=N,
+        N_couplings=N,
+        measurement_detuning=0,
+        α=0,
+        rwa=False,
+        flat_energies=False,
+        correct_lamb_shift=0,
+        laser_off_time=0,
+        small_loop_detuning=small_loop_detuning,
+    )
+
+    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.δ * params.Ω * 1.4,
+                # params.Ω,
+            ]
+        ),
+        np.array([1.0]),
+    )
+
+    print(params.drive_override)
+
+    t = time_axis(params, lifetimes=total_lifetimes, resolution=0.001)
+    solution = solve(t, params)
+    return params, t, solution
+
+
+def generate_phase_one_data(
+    drive_detuning=0,
+    g_0=0.3,
+    off_factor=0.5,
+    noise=False,
+    extra_title="",
+    laser_detuning=0,
+    yscale="linear",
+):
+    """
+    This generates some intensity data for phase one of the
+    calibration protocol, where the aim is to extract the resonator
+    spectrum.
+    """
+
+    total_lifetimes = 10
+    eom_off_lifetime = total_lifetimes * off_factor
+    fluct_size = 0.05
+    noise_amp = 0.3
+
+    params, t, solution = make_params_and_solve(
+        total_lifetimes,
+        eom_off_lifetime,
+        N=2,
+        g_0=g_0,
+        small_loop_detuning=0,
+        laser_detuning=laser_detuning,
+        η=0.18,
+        η_hybrid=0.18 * 5,
+        drive_detuning=drive_detuning,
+    )
+
+    rng = np.random.default_rng(seed=0)
+    raw_signal = output_signal(t, solution.y, params)
+
+    signal = raw_signal.copy()
+    if noise:
+        signal += noise_amp * rng.standard_normal(len(signal))
+
+    fig = make_figure(f"simulation_noise_{noise}", figsize=(20, 3 * 5))
+
+    ax_realtime, ax_rotating_bath, ax_rotating_system, ax_spectrum = (
+        fig.subplot_mosaic("""
+    AA
+    BC
+    DD
+    DD
+    """).values()
+    )
+
+    ax_rotating_system.sharey(ax_rotating_bath)
+    ax_rotating_system.sharex(ax_rotating_bath)
+
+    for mode in range(2):
+        ax_rotating_system.plot(
+            t[::50],
+            abs((solution.y[mode, ::50]) * (params.η / params.η_hybrid)) ** 2 / 2,
+        )
+
+    for mode in range(2, solution.y.shape[0]):
+        ax_rotating_bath.plot(
+            t[::50], abs((solution.y[mode, ::50])) ** 2, label=mode_name(mode)
+        )
+
+    for ax in [ax_rotating_bath, ax_rotating_system]:
+        ax.set_xlabel("Time")
+        ax.set_ylabel("Intensity")
+        ax.legend()
+
+    ax_rotating_bath.set_title("Bath Modes")
+    ax_rotating_system.set_title(
+        "Hybridized Modes [corrected for magnitude visible in FFT]"
+    )
+
+    ax_realtime.plot(t[::50], signal[::50])
+    ax_realtime.axvline(params.drive_off_time, color="black", linestyle="--")
+    ax_realtime.set_xlabel("Time")
+    ax_realtime.set_ylabel("Intensity")
+    ax_realtime.set_title("Photo-diode AC Intensity")
+
+    # now we plot the power spectrum
+    window = (float(params.laser_off_time) or 0, t[-1])
+    # window = (0, float(params.laser_off_time or 0))
+
+    freq, fft, t = fourier_transform(
+        t,
+        signal,
+        low_cutoff=1,
+        high_cutoff=params.Ω * 5,
+        window=window,
+        ret_time=True,
+    )
+
+    ax_spectrum.clear()
+    power = np.abs(fft) ** 2
+    ax_spectrum.plot(freq, power)
+
+    runtime = RuntimeParams(params)
+
+    for i, freq in enumerate(runtime.Ωs.real):
+        pos = np.abs(
+            params.measurement_detuning
+            - freq / (2 * np.pi)
+            + params.δ * params.Ω
+            + params.laser_detuning,
+        )
+
+        ax_spectrum.axvline(
+            pos,
+            color="red",
+            linestyle="--",
+            zorder=-100,
+        )
+        ax_spectrum.annotate(
+            mode_name(i),
+            (pos + 0.1, power.max() * (0.9 if freq < 0 else 1)),
+        )
+
+    ax_spectrum.set_yscale(yscale)
+
+    return fig
+
+
+# %% save
+if __name__ == "__main__":
+    fig = generate_phase_one_data(
+        noise=True,
+        g_0=0.8,
+        off_factor=0.5,
+        laser_detuning=13 * (1 - 1 / 4),
+    )
diff --git a/scripts/experiments/013_step_through.py b/scripts/experiments/013_step_through.py
new file mode 100644
index 0000000..2683299
--- /dev/null
+++ b/scripts/experiments/013_step_through.py
@@ -0,0 +1,201 @@
+from rabifun.system import *
+from rabifun.plots import *
+from rabifun.utilities import *
+from rabifun.analysis import *
+from ringfit.data import ScanData
+import functools
+import multiprocessing
+import copy
+from scipy.ndimage import rotate
+
+# %% interactive
+
+
+def make_shots(params, total_lifetimes, eom_range, eom_steps):
+    solutions = []
+    t = time_axis(params, lifetimes=total_lifetimes, resolution=0.01)
+
+    pool = multiprocessing.Pool()
+
+    shot_params = []
+    for step in range(eom_steps):
+        current_params = copy.deepcopy(params)
+        current_params.drive_override = (
+            np.array(
+                [
+                    params.Ω
+                    * (eom_range[0] + (eom_range[1] - eom_range[0]) * step / eom_steps)
+                ]
+            ),
+            np.array([1.0]),
+        )
+
+        shot_params.append((t, current_params))
+
+    solutions = pool.starmap(solve, shot_params)
+    return t, solutions
+
+
+def process_shots(t, solutions, noise_amplitude, params):
+    average_power_spectrum = None
+    window = (float(params.laser_off_time) or 0, t[-1])
+
+    rng = np.random.default_rng(seed=0)
+
+    # let us get a measure calibrate the noise strength
+    signal_amp = 0
+    signals = []
+    for solution in solutions:
+        signal = output_signal(t, solution.y, params)
+        signals.append(signal)
+        signal_amp += abs(signal).max()
+
+    signal_amp /= len(solutions)
+
+    for signal in signals:
+        signal += rng.normal(scale=noise_amplitude * signal_amp, size=len(signal))
+
+        freq, fft = fourier_transform(
+            t,
+            signal,
+            low_cutoff=0.0 * params.Ω,
+            high_cutoff=params.Ω * (params.N + 1),
+            window=window,
+        )
+
+        power = np.abs(fft) ** 2
+        if average_power_spectrum is None:
+            average_power_spectrum = power
+
+        else:
+            average_power_spectrum += power
+
+    return freq, (average_power_spectrum / len(solutions))
+
+
+def plot_power_spectrum(ax_spectrum, freq, average_power_spectrum, params):
+    offset = np.median(average_power_spectrum)
+    ax_spectrum.plot(freq, average_power_spectrum)
+
+    runtime = RuntimeParams(params)
+    lorentz_freqs = np.linspace(freq.min(), freq.max(), 10000)
+    fake_spectrum = np.zeros_like(lorentz_freqs)
+
+    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="red",
+            linestyle="--",
+            zorder=-100,
+        )
+
+        amplitude = average_power_spectrum[np.argmin(abs(freq - pos))] - offset
+        ax_spectrum.annotate(
+            mode_name(i),
+            (
+                pos + 0.1,
+                (amplitude + offset) * (1 if peak_freq.real < 0 else 1.1),
+            ),
+        )
+
+        fake_spectrum += lorentzian(
+            lorentz_freqs, amplitude, pos, peak_freq.imag / np.pi
+        )
+
+    ax_spectrum.plot(lorentz_freqs, fake_spectrum + offset)
+
+
+def generate_data(
+    g_0=0.3,
+    η_factor=5,
+    noise_amplitude=0.3,
+    laser_detuning=0,
+    N=10,
+    eom_ranges=(0.5, 2.5),
+    eom_steps=20,
+    small_loop_detuning=0,
+    excitation_lifetimes=2,
+    measurement_lifetimes=4,
+):
+    η = 0.2
+    Ω = 13
+
+    params = Params(
+        η=η,
+        η_hybrid=η_factor * η,
+        Ω=Ω,
+        δ=1 / 4,
+        ω_c=0.1,
+        g_0=g_0,
+        laser_detuning=13 * (-1 - 1 / 4) + laser_detuning,
+        N=N,
+        N_couplings=N,
+        measurement_detuning=0,
+        α=0,
+        rwa=False,
+        flat_energies=False,
+        correct_lamb_shift=0,
+        laser_off_time=0,
+        small_loop_detuning=small_loop_detuning,
+        drive_override=(np.array([1.0]), np.array([1.0])),
+    )
+
+    params.laser_off_time = params.lifetimes(excitation_lifetimes)
+    params.drive_off_time = params.lifetimes(excitation_lifetimes)
+
+    t, solutions = make_shots(
+        params, excitation_lifetimes + measurement_lifetimes, eom_ranges, eom_steps
+    )
+
+    _, (sol_on_res) = make_shots(
+        params,
+        excitation_lifetimes + measurement_lifetimes,
+        ((1 + params.δ), (1 + params.δ)),
+        1,
+    )
+
+    freq, average_power_spectrum = process_shots(t, solutions, noise_amplitude, params)
+    _, spectrum_on_resonance = process_shots(t, sol_on_res, noise_amplitude, params)
+
+    fig = make_figure()
+    fig.clear()
+
+    ax_multi, ax_single = fig.subplot_mosaic("AA\nBB").values()
+    ax_multi.set_title("Averaged Power Spectrum")
+    ax_single.set_title("Single-shot Power-Spectrum with EOM on resonance")
+
+    for ax in [ax_multi, ax_single]:
+        ax.set_xlabel("Frequency (MHz)")
+        ax.set_ylabel("Power")
+        ax.set_yscale("log")
+
+    plot_power_spectrum(ax_multi, freq, average_power_spectrum, params)
+    plot_power_spectrum(ax_single, freq, spectrum_on_resonance, params)
+
+    runtime = RuntimeParams(params)
+    # ax_spectrum.set_yscale(yscale)
+
+    return fig
+
+
+# %% save
+if __name__ == "__main__":
+    fig = generate_data(
+        g_0=1,
+        η_factor=5,
+        noise_amplitude=0.25,
+        N=2,
+        eom_ranges=(1.1, 1.3),
+        eom_steps=100,
+        small_loop_detuning=0,
+        laser_detuning=0,
+        excitation_lifetimes=2,
+        measurement_lifetimes=20,
+    )
diff --git a/src/rabifun/system.py b/src/rabifun/system.py
index abd05ca..c2b378c 100644
--- a/src/rabifun/system.py
+++ b/src/rabifun/system.py
@@ -260,17 +260,17 @@ def eom_drive(t, x, ds, ωs, det_matrix, a_weights):
     # print(np.min(test))
     # print(np.argmin(test, keepdims=True))
 
-    det_matrix = np.exp(-1j * det_matrix * t)
+    rot_matrix = np.exp(-1j * det_matrix * t)
     for i, weight in enumerate(a_weights):
-        det_matrix[i, 2:] *= weight
-        det_matrix[2:, i] *= weight.conjugate()
+        rot_matrix[i, 2:] *= weight
+        rot_matrix[2:, i] *= weight.conjugate()
 
-    # FIXME: that's not strictly right for the non symmetric damping
+    # # FIXME: that's not strictly right for the non symmetric damping
     prod = a_weights[0] * a_weights[1].conj()
-    det_matrix[0, 1] *= prod
-    det_matrix[1, 0] *= prod.conjugate()
+    rot_matrix[0, 1] *= prod
+    rot_matrix[1, 0] *= prod.conjugate()
 
-    driven_x = np.sum(2 * ds * np.sin(2 * np.pi * ωs * t)) * (det_matrix @ x)
+    driven_x = np.sum(2 * ds * np.sin(2 * np.pi * ωs * t)) * (rot_matrix @ x)
 
     return driven_x
 
@@ -360,7 +360,7 @@ def solve(t: np.ndarray, params: Params, **kwargs):
         (np.min(t), np.max(t)),
         initial,
         vectorized=False,
-        max_step=np.pi / (params.Ω * params.N),
+        max_step=1 / 2 * np.pi / (params.Ω * (2 * params.N + 2)),
         t_eval=t,
         method="DOP853",
         atol=1e-7,
@@ -383,6 +383,7 @@ def output_signal(t: np.ndarray, amplitudes: np.ndarray, params: Params):
     rotating = amplitudes.copy() * np.exp(
         -1j * (runtime.detuned_Ωs[:, None] * t[None, :] - laser_times[None, :])
     )
+
     rotating[0:2, :] *= runtime.a_weights[:, None].conjugate()
 
     return (np.sum(rotating, axis=0)).imag