WIP: fix units in sim, seems to work

scripts/experiments/.#002_rabi_detuning_scan.py

@@ -0,0 +1 @@
+hiro@Lobsang.149236:1716554985

scripts/experiments/002_rabi_detuning_scan.py

@@ -9,12 +9,25 @@ from plot_utils import wrap_plot
 def transient_rabi():
     """A transient rabi oscillation without noise."""
 
-    params = Params(η=0.0001, δ=1 / 4, d=0.1, laser_detuning=0.1, Δ=0.005, N=2)
-    t = time_axis(params, 3, 0.1)
+    params = Params(
+        η=0.7,
+        Ω=13.07,
+        δ=13.07 / 4,
+        d=13.07 / 4 * 1,
+        laser_detuning=0.1,
+        Δ=13.07 / 4 * 0.01,
+        N=3,
+        measurement_detuning=1,
+        rwa=False,
+    )
+    params.laser_off_time = params.lifetimes(10)
+    t = time_axis(params, 50, 0.1)
     solution = solve(t, params)
-    signal = output_signal(t, solution.y, params.laser_detuning)
+    signal = output_signal(t, solution.y, params)
 
-    f, (_, ax) = plot_simulation_result(make_figure(), t, signal, params)
+    f, (_, ax) = plot_simulation_result(
+        make_figure(), t, signal, params, window=(params.lifetimes(10), t[-1])
+    )
     plot_sidebands(ax, params)
     # ax.set_xlim(0.73, 0.77)
     f.suptitle("Transient Rabi oscillation")
@@ -27,7 +40,7 @@ def steady_rabi():
     t = time_axis(params, lifetimes=10, resolution=1)
 
     solution = solve(t, params)
-    signal = output_signal(t, solution.y, params.laser_detuning)
+    signal = output_signal(t, solution.y, params)
 
     f, (_, ax) = plot_simulation_result(
         make_figure(), t, signal, params, window=(params.lifetimes(8), t[-1])
@@ -43,7 +56,7 @@ def noisy_transient_rabi():
     params = Params(η=0.001, d=0.1, laser_detuning=0, Δ=0.05)
     t = time_axis(params, 2, 1)
     solution = solve(t, params)
-    signal = output_signal(t, solution.y, params.laser_detuning)
+    signal = output_signal(t, solution.y, params)
 
     noise_strength = 0.1
     signal = add_noise(signal, noise_strength)
@@ -54,7 +67,6 @@ def noisy_transient_rabi():
     f.suptitle(f"Transient Rabi oscillation with noise strength {noise_strength}.")
 
 
-@autoclose
 def ringdown_after_rabi():
     """Demonstrates the nonstationary ringdown of the resonator after turning off the EOM and laser drive."""
     off_lifetime = 4
@@ -66,7 +78,7 @@ def ringdown_after_rabi():
 
     t = time_axis(params, lifetimes=5, resolution=1)
     solution = solve(t, params)
-    signal = output_signal(t, solution.y, params.laser_detuning)
+    signal = output_signal(t, solution.y, params)
 
     # noise_strength = 0.1
     # signal = add_noise(signal, noise_strength)
@@ -79,3 +91,18 @@ def ringdown_after_rabi():
     fftax.axvline(params.laser_detuning, color="black")
 
     f.suptitle(f"Ringdown after rabi osci EOM after {off_lifetime} lifetimes.")
+
+
+def sweep():
+    """A transient rabi oscillation without noise."""
+
+    params = Params(η=1, δ=1 / 4, d=0.0, laser_detuning=100, N=1)
+    t = time_axis(params, 1000, 0.001)
+    params.dynamic_detunting = -100, t[-1]
+    solution = solve(t, params)
+    signal = output_signal(t, solution.y, params)
+
+    f, (_, ax) = plot_simulation_result(make_figure(), t, signal, params)
+    plot_sidebands(ax, params)
+    # ax.set_xlim(0.73, 0.77)
+    f.suptitle("Transient Rabi oscillation")

src/rabifun/plots.py

@@ -1,5 +1,5 @@
 from plot_utils import *
-from .system import Params, output_signal
+from .system import Params
 from .analysis import fourier_transform
 import matplotlib.pyplot as plt
 import numpy as np
@@ -34,18 +34,24 @@ def plot_simulation_result(
     freq, fft = fourier_transform(t, signal, window)
     fft = fft / np.max(np.abs(fft))
 
-    freq *= 2 * np.pi
     ax2.set_xlim(0, params.Ω * (params.N))
     ax3 = ax2.twinx()
 
-    ax2.plot(freq, np.abs(fft) ** 2)
-    ax2.set_yscale("log")
+    ax2.plot(freq, np.abs(fft))
+    # ax2.set_yscale("log")
     ax2.set_title("FFT")
-    ax2.set_xlabel("ω [angular]")
+    ax2.set_xlabel("ω [linear]")
     ax2.set_ylabel("Power")
     ax2.legend()
 
-    ax3.plot(freq, np.angle(fft), linestyle="--", color="C2", alpha=0.5, zorder=-10)
+    ax3.plot(
+        freq,
+        np.unwrap(np.angle(fft) + np.pi, 2 * np.pi),
+        linestyle="--",
+        color="C2",
+        alpha=0.5,
+        zorder=-10,
+    )
     ax3.set_ylabel("Phase")
 
     return (ax1, ax2)
@@ -57,20 +63,26 @@ def plot_sidebands(ax, params: Params):
     :param ax: axis to plot on
     :param params: system parameters
     """
-    energy = params.rabi_splitting
+    energy = params.rabi_splitting / (2 * np.pi)
 
     first_sidebands = np.abs(
-        -params.laser_detuning + np.array([1, -1]) * energy / 2 + params.Δ / 2
+        -(params.laser_detuning + params.measurement_detuning)
+        + np.array([1, -1]) * energy / 2
+        + params.Δ / 2
     )
     second_sidebands = (
         params.Ω
         - params.δ
-        - params.laser_detuning
+        - (params.laser_detuning + params.measurement_detuning)
         + np.array([1, -1]) * energy / 2
         - params.Δ / 2
     )
 
-    ax.axvline(params.ω_eom, color="black", label="steady state")
+    ax.axvline(
+        params.ω_eom / (2 * np.pi) - params.measurement_detuning,
+        color="black",
+        label="steady state",
+    )
 
     for n, sideband in enumerate(first_sidebands):
         ax.axvline(

src/rabifun/system.py

@@ -29,6 +29,8 @@ class Params:
     laser_detuning: float = 0.0
     """Detuning of the laser relative to the _A_ mode."""
 
+    measurement_detuning: float = 0.0
+
     laser_off_time: float | None = None
     """Time at which the laser is turned off."""
 
@@ -38,8 +40,10 @@ class Params:
     rwa: bool = False
     """Whether to use the rotating wave approximation."""
 
+    dynamic_detunting: tuple[float, float] = 0, 0
+
     def periods(self, n: float):
-        return n * 2 * np.pi / self.Ω
+        return n / self.Ω
 
     def lifetimes(self, n: float):
         return n / self.η
@@ -50,15 +54,16 @@ class Params:
 
     @property
     def ω_eom(self):
-        return self.Ω - self.δ - self.Δ
+        return 2 * np.pi * (self.Ω - self.δ - self.Δ)
 
 
 class RuntimeParams:
     """Secondary Parameters that are required to run the simulation."""
 
     def __init__(self, params: Params):
-        Ωs = np.arange(0, params.N) * params.Ω - 1j * np.repeat(params.η, params.N)
-        Ωs[1:] -= params.δ
+        Ωs = 2 * np.pi * np.concatenate(
+            [[-1 * params.δ, params.δ], np.arange(1, params.N + 1) * params.Ω]
+        ) - 1j * np.repeat(params.η / 2, params.N + 2)
 
         self.Ωs = Ωs
 
@@ -80,14 +85,13 @@ def time_axis(params: Params, lifetimes: float, resolution: float = 1):
     )
 
 
-def eom_drive(t, x, d, ω, rwa):
+def eom_drive(t, x, d, ω):
     """The electrooptical modulation drive.
 
     :param t: time
     :param x: amplitudes
     :param d: drive amplitude
     :param ω: drive frequency
-    :param rwa: whether to use the rotating wave approximation
     """
 
     stacked = np.repeat([x], len(x), 0)
@@ -95,25 +99,40 @@ def eom_drive(t, x, d, ω, rwa):
     stacked = np.sum(stacked, axis=1)
     driven_x = d * np.sin(ω * t) * stacked
 
-    if rwa and len(x) > 2:
-        driven_x[2:] = 0
-
     return driven_x
 
 
+def laser_frequency(params: Params, t: np.ndarray):
+    base = 2 * np.pi * (params.laser_detuning + params.δ)
+    if params.dynamic_detunting[1] == 0:
+        return base
+
+    return base + 2 * np.pi * (
+        params.dynamic_detunting[0] * t / params.dynamic_detunting[1]
+    )
+
+
 def make_righthand_side(runtime_params: RuntimeParams, params: Params):
     """The right hand side of the equation of motion."""
 
     def rhs(t, x):
         differential = runtime_params.Ωs * x
 
+        if params.rwa:
+            x[0] = 0
+            x[3:] = 0
+
         if (params.drive_off_time is None) or (t < params.drive_off_time):
-            differential += eom_drive(t, x, params.d, params.ω_eom, params.rwa)
+            differential += eom_drive(t, x, params.d, params.ω_eom)
 
         if (params.laser_off_time is None) or (t < params.laser_off_time):
-            laser = np.exp(-1j * params.laser_detuning * t)
-            differential[0] += laser / np.sqrt(2)
-            differential[1:] += laser
+            laser = np.exp(-1j * laser_frequency(params, t) * t)
+            differential[0:2] += laser / np.sqrt(2)
+            differential[2:] += laser
+
+        if params.rwa:
+            differential[0] = 0
+            differential[3:] = 0
 
         return -1j * differential
 
@@ -130,14 +149,14 @@ def solve(t: np.ndarray, params: Params):
     runtime = RuntimeParams(params)
     rhs = make_righthand_side(runtime, params)
 
-    initial = np.zeros(params.N, np.complex128)
+    initial = np.zeros(params.N + 2, np.complex128)
 
     return scipy.integrate.solve_ivp(
         rhs,
         (np.min(t), np.max(t)),
         initial,
         vectorized=False,
-        # max_step=0.1 * np.pi / (params.Ω * params.N),
+        # max_step=0.01 * np.pi / (params.Ω * params.N),
         t_eval=t,
     )
 
@@ -151,9 +170,16 @@ def in_rotating_frame(
     return amplitudes * np.exp(1j * Ωs[:, None].real * t[None, :])
 
 
-def output_signal(t: np.ndarray, amplitudes: np.ndarray, laser_detuning: float):
+def output_signal(t: np.ndarray, amplitudes: np.ndarray, params: Params):
     """
     Calculate the output signal when mixing with laser light of
     frequency `laser_detuning`.
     """
-    return (np.sum(amplitudes, axis=0) * np.exp(1j * laser_detuning * t)).imag
+    return (
+        np.sum(amplitudes, axis=0)
+        * np.exp(
+            1j
+            * (laser_frequency(params, t) + 2 * np.pi * params.measurement_detuning)
+            * t
+        )
+    ).imag