feature: transform sketch into proper python module

flake.nix

@@ -45,6 +45,10 @@
             pyright
             python3Packages.jupyter
           ];
+
+          shellHook = ''
+          export PYTHONPATH=$(pwd)/src:$PYTHONPATH
+          '';
         };
 
         # Shell for poetry.

scripts/experiments/002_rabi_detuning_scan.py

@@ -1,175 +1,84 @@
-import numpy as np
-import matplotlib.pyplot as plt
-from dataclasses import dataclass
-import scipy.integrate
-
-
-@dataclass
-class Params:
-    N: int = 2
-    """Number of bath modes."""
-
-    Ω: float = 1
-    """FSR"""
-
-    η: float = 0.1
-    """Decay rate."""
-
-    d: float = 0.01
-    """Drive amplitude."""
-
-    Δ: float = 0.0
-    """Detuning."""
-
-    laser_detuning: float = 0.0
-    """Detuning of the laser."""
-
-    laser_off_time: float | None = None
-    """Time at which the laser is turned off."""
-
-    drive_off_time: float | None = None
-    """Time at which the drive is turned off."""
-
-    def periods(self, n: float):
-        return n * 2 * np.pi / self.Ω
-
-    def lifetimes(self, n: float):
-        return n / self.η
-
-
-class RuntimeParams:
-    def __init__(self, params: Params):
-        self.Ωs = np.arange(0, params.N) * params.Ω - 1j * np.repeat(params.η, params.N)
-
-
-def drive(t, x, d, Δ, Ω):
-    stacked = np.repeat([x], len(x), 0)
-
-    np.fill_diagonal(stacked, 0)
-
-    stacked = np.sum(stacked, axis=1)
-    driven_x = d * np.sin((Ω - Δ) * t) * stacked
-
-    return driven_x
-
-
-def make_righthand_side(Ωs, params: Params):
-    def rhs(t, x):
-        differential = Ωs * x
-
-        if (params.drive_off_time is None) or (t < params.drive_off_time):
-            differential += drive(t, x, params.d, params.Δ, params.Ω)
-
-        if (params.laser_off_time is None) or (t < params.laser_off_time):
-            differential += np.exp(-1j * params.laser_detuning * t)
-
-        return -1j * differential
-
-    return rhs
-
-
-def solve(t: np.ndarray, params: Params):
-    runtime = RuntimeParams(params)
-    rhs = make_righthand_side(runtime.Ωs, params)
-
-    initial = np.zeros(params.N, np.complex128)
-
-    sol = scipy.integrate.solve_ivp(
-        rhs,
-        (0, np.max(t)),
-        initial,
-        vectorized=False,
-        max_step=0.01 * 2 * np.pi / np.max(abs(runtime.Ωs.real)),
-        method="BDF",
-        t_eval=t,
-    )
-    return sol
-
-
-def in_rotating_frame(t, amplitudes, params: Params):
-    Ωs = RuntimeParams(params).Ωs
-
-    return amplitudes * np.exp(1j * Ωs[:, None].real * t[None, :])
-
-
-def output_signal(t: np.ndarray, amplitudes: np.ndarray, laser_detuning: float):
-    return (np.sum(amplitudes, axis=0) * np.exp(1j * laser_detuning * t)).imag
-
-
-def reflect(center, value):
-    diff = center - value
-    return center - diff
-
+from rabifun.plots import plot_simulation_result
+from rabifun.system import *
+from rabifun.plots import *
+from rabifun.utilities import *
+from plot_utils import autoclose
 
 # %% interactive
-f, (ax1, ax2) = plt.subplots(2, 1)
 
 
-def transient_analysis():
+@autoclose
+def transient_rabi():
+    """A transient rabi oscillation without noise."""
+
     params = Params(η=0.001, d=0.1, laser_detuning=0, Δ=0.05)
-    params.laser_off_time = params.lifetimes(2)
-    params.drive_off_time = params.lifetimes(2)
-    t = np.arange(0, params.lifetimes(5), 0.5 / (params.Ω * params.N))
+    t = time_axis(params, 2, 1)
     solution = solve(t, params)
-
     signal = output_signal(t, solution.y, params.laser_detuning)
-    # signal += np.random.normal(0, .1 * np.max(abs(signal)), len(signal))
 
-    window = (params.lifetimes(2) > t) & (t > params.lifetimes(0))
-    # window = t > params.lifetimes(2)
-    t = t[window]
-    signal = signal[window]
+    f, (_, ax) = plot_simulation_result(t, signal, params)
+    plot_sidebands(ax, params)
+    f.suptitle("Transient Rabi oscillation")
 
-    ax1.clear()
-    # ax1.plot(
-    #     solution.t, np.real(in_rotating_frame(solution.t, solution.y, params)[1, :])
-    # )
-    ax1.plot(t, signal)
-    ax1.set_title(f"Output signal\n {params}")
-    ax1.set_xlabel("t")
-    ax1.set_ylabel("Intensity")
 
-    ax2.clear()
-    freq = np.fft.rfftfreq(len(t), t[1] - t[0]) * 2 * np.pi
-    fft = np.fft.rfft(signal)
-    ax2.set_xlim(0, params.Ω * (params.N))
+@autoclose
+def steady_rabi():
+    """A steady state rabi oscillation without noise."""
 
-    energy = 1 / 2 * np.sqrt(4 * params.d**2 + 4 * params.Δ**2)
-    for n in range(params.N):
-        sidebands = (
-            (n) * params.Ω
-            - params.laser_detuning
-            + np.array([1, -1]) * energy / 2
-            - params.Δ / 2
-        )
+    params = Params(η=0.001, d=0.1, laser_detuning=0, Δ=0.05)
+    t = time_axis(params, lifetimes=10, resolution=1)
 
-        right_sidebands = (
-            (n) * params.Ω
-            + params.laser_detuning
-            + np.array([1, -1]) * energy / 2
-            - params.Δ / 2
-        )
+    solution = solve(t, params)
+    signal = output_signal(t, solution.y, params.laser_detuning)
 
-        for sideband in sidebands:
-            ax2.axvline(
-                sideband,
-                color=f"C{n}",
-                label=f"rabi-sideband {n}",
-            )
+    f, (_, ax) = plot_simulation_result(
+        t, signal, params, window=(params.lifetimes(8), t[-1])
+    )
+    plot_sidebands(ax, params)
 
-        for sideband in right_sidebands:
-            ax2.axvline(
-                sideband,
-                color=f"C{n}",
-                linestyle="--",
-            )
-    ax2.axvline(params.Ω - params.Δ, color="black", label="")
-    ax2.axvline(2 * params.Ω - params.Δ, color="black", label="")
-    ax2.plot(freq, np.abs(fft) ** 2)
-    ax2.set_yscale("log")
-    ax2.set_title("FFT")
-    ax2.set_xlabel("ω")
-    ax2.set_ylabel("Power")
-    ax2.legend()
+    f.suptitle("Steady State Rabi oscillation. No Rabi Sidebands.")
 
-    return solution
+
+@autoclose
+def noisy_transient_rabi():
+    """A transient rabi oscillation with noise."""
+
+    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)
+
+    noise_strength = 0.1
+    signal = add_noise(signal, noise_strength)
+
+    f, (_, ax) = plot_simulation_result(t, signal, params)
+    plot_sidebands(ax, params)
+
+    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
+    laser_detuning = 0.1
+
+    params = Params(η=0.001, d=0.1, laser_detuning=laser_detuning, Δ=0.5)
+    params.laser_off_time = params.lifetimes(off_lifetime)
+    params.drive_off_time = params.lifetimes(off_lifetime)
+
+    t = time_axis(params, lifetimes=5, resolution=1)
+    solution = solve(t, params)
+    signal = output_signal(t, solution.y, params.laser_detuning)
+
+    # noise_strength = 0.1
+    # signal = add_noise(signal, noise_strength)
+
+    f, (_, fftax) = plot_simulation_result(
+        t, signal, params, window=(params.lifetimes(off_lifetime), t[-1])
+    )
+
+    fftax.axvline(params.Ω - params.laser_detuning, color="black")
+    fftax.axvline(params.laser_detuning, color="black")
+
+    f.suptitle(f"Ringdown after rabi osci EOM after {off_lifetime} lifetimes.")

scripts/transients_with_blackout.py

@@ -1,6 +1,5 @@
 import sys
 
-sys.path.append("../")
 from ringfit import data
 import matplotlib.pyplot as plt
 from ringfit.data import *

scripts/transients_without_amplification.py

@@ -1,6 +1,5 @@
 import sys
 
-sys.path.append("../")
 from ringfit import data
 import matplotlib.pyplot as plt
 from ringfit.data import *

scripts/weird_oscillations.py

@@ -1,6 +1,5 @@
 import sys
 
-sys.path.append("../")
 from ringfit import data
 import matplotlib.pyplot as plt
 from ringfit.data import *

src/plot_utils.py

@@ -0,0 +1,42 @@
+import matplotlib.pyplot as plt
+
+
+def wrap_plot(f):
+    def wrapped(*args, ax=None, setup_function=plt.subplots, **kwargs):
+        fig = None
+        if not ax:
+            fig, ax = setup_function()
+
+        ret_val = f(*args, ax=ax, **kwargs)
+        return (fig, ax, ret_val) if ret_val else (fig, ax)
+
+    return wrapped
+
+
+def autoclose(f):
+    def wrapped(*args, **kwargs):
+        plt.close()
+        return f(*args, **kwargs)
+
+    return wrapped
+
+
+def lighten_color(color, amount=0.5):
+    """
+    Lightens the given color by multiplying (1-luminosity) by the given amount.
+    Input can be matplotlib color string, hex string, or RGB tuple.
+
+    Examples:
+    >> lighten_color('g', 0.3)
+    >> lighten_color('#F034A3', 0.6)
+    >> lighten_color((.3,.55,.1), 0.5)
+    """
+    import matplotlib.colors as mc
+    import colorsys
+
+    try:
+        c = mc.cnames[color]
+    except:
+        c = color
+    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
+    return colorsys.hls_to_rgb(c[0], 1 - amount * (1 - c[1]), c[2])

src/rabifun/analysis.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+
+def fourier_transform(
+    t: np.ndarray, signal: np.ndarray, window: tuple[float, float] | None = None
+):
+    """
+    Compute the Fourier transform of a signal from the time array
+    ``t`` and the real signal ``signal``.  Optionally, a time window
+    can be specified through ``window = ([begin], [end])`` .
+
+    :returns: The (linear) frequency array and the Fourier transform.
+    """
+
+    if window:
+        mask = (window[1] > t) & (t > window[0])
+        t = t[mask]
+        signal = signal[mask]
+
+    freq = np.fft.rfftfreq(len(t), t[1] - t[0])
+    fft = np.fft.rfft(signal)
+
+    return freq, fft

src/rabifun/plots.py

@@ -0,0 +1,53 @@
+from plot_utils import *
+from .system import Params, output_signal
+from .analysis import fourier_transform
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def plot_simulation_result(
+    t: np.ndarray, signal: np.ndarray, params: Params, window=None
+):
+    f, (ax1, ax2) = plt.subplots(2, 1)
+
+    ax1.plot(t, signal)
+    ax1.set_title(f"Output signal\n {params}")
+    ax1.set_xlabel("t")
+    ax1.set_ylabel("Intensity")
+
+    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.set_title("FFT")
+    ax2.set_xlabel("ω [angular]")
+    ax2.set_ylabel("Power")
+    ax2.legend()
+
+    ax3.plot(freq, np.angle(fft), linestyle="--", color="C2", alpha=0.5, zorder=-10)
+    ax3.set_ylabel("Phase")
+
+    return f, (ax1, ax2)
+
+
+def plot_sidebands(ax, params: Params):
+    energy = params.rabi_splitting
+    sidebands = (
+        params.Ω - params.laser_detuning + np.array([1, -1]) * energy / 2 - params.Δ / 2
+    )
+
+    ax.axvline(params.Ω - params.Δ, color="black", label="steady state")
+
+    for n, sideband in enumerate(sidebands):
+        ax.axvline(
+            sideband,
+            color=f"C{n}",
+            label=f"rabi-sideband {n}",
+        )
+
+    ax.legend()

src/rabifun/system.py

@@ -0,0 +1,142 @@
+import numpy as np
+import matplotlib.pyplot as plt
+from dataclasses import dataclass
+import scipy.integrate
+
+
+@dataclass
+class Params:
+    """Parameters for the system."""
+
+    N: int = 2
+    """Number of bath modes."""
+
+    Ω: float = 1
+    """Free spectral range of the system."""
+
+    η: float = 0.1
+    """Decay rate of the system."""
+
+    d: float = 0.01
+    """Drive amplitude."""
+
+    Δ: float = 0.0
+    """Detuning of the EOM drive."""
+
+    laser_detuning: float = 0.0
+    """Detuning of the laser relative to the _A_ mode."""
+
+    laser_off_time: float | None = None
+    """Time at which the laser is turned off."""
+
+    drive_off_time: float | None = None
+    """Time at which the drive is turned off."""
+
+    def periods(self, n: float):
+        return n * 2 * np.pi / self.Ω
+
+    def lifetimes(self, n: float):
+        return n / self.η
+
+    @property
+    def rabi_splitting(self):
+        return 1 / 2 * np.sqrt(4 * self.d**2 + 4 * self.Δ**2)
+
+
+class RuntimeParams:
+    """Secondary Parameters that are required to run the simulation."""
+
+    def __init__(self, params: Params):
+        self.Ωs = np.arange(0, params.N) * params.Ω - 1j * np.repeat(params.η, params.N)
+
+
+def time_axis(params: Params, lifetimes: float, resolution: float = 1):
+    """Generate a time axis for the simulation.
+
+    :param params: system parameters
+    :param lifetimes: number of lifetimes to simulate
+    :param resolution: time resolution
+
+        Setting this to `1` will give the time axis just enough
+        resolution to capture the fastest relevant frequencies.  A
+        smaller value yields more points in the time axis.
+    """
+
+    return np.arange(
+        0, params.lifetimes(lifetimes), resolution * np.pi / (params.Ω * params.N)
+    )
+
+
+def eom_drive(t, x, d, Δ, Ω):
+    """The electrooptical modulation drive.
+
+    :param t: time
+    :param x: amplitudes
+    :param d: drive amplitude
+    :param Δ: detuning
+    :param Ω: FSR
+    """
+    stacked = np.repeat([x], len(x), 0)
+
+    np.fill_diagonal(stacked, 0)
+
+    stacked = np.sum(stacked, axis=1)
+    driven_x = d * np.sin((Ω - Δ) * t) * stacked
+
+    return driven_x
+
+
+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.drive_off_time is None) or (t < params.drive_off_time):
+            differential += eom_drive(t, x, params.d, params.Δ, params.Ω)
+
+        if (params.laser_off_time is None) or (t < params.laser_off_time):
+            differential += np.exp(-1j * params.laser_detuning * t)
+
+        return -1j * differential
+
+    return rhs
+
+
+def solve(t: np.ndarray, params: Params):
+    """Integrate the equation of motion.
+
+    :param t: time array
+    :param params: system parameters
+    """
+
+    runtime = RuntimeParams(params)
+    rhs = make_righthand_side(runtime, params)
+
+    initial = np.zeros(params.N, np.complex128)
+
+    return scipy.integrate.solve_ivp(
+        rhs,
+        (np.min(t), np.max(t)),
+        initial,
+        vectorized=False,
+        max_step=0.01 * 2 * np.pi / np.max(abs(runtime.Ωs.real)),
+        t_eval=t,
+    )
+
+
+def in_rotating_frame(
+    t: np.ndarray, amplitudes: np.ndarray, params: Params
+) -> np.ndarray:
+    """Transform the amplitudes to the rotating frame."""
+    Ωs = RuntimeParams(params).Ωs
+
+    return amplitudes * np.exp(1j * Ωs[:, None].real * t[None, :])
+
+
+def output_signal(t: np.ndarray, amplitudes: np.ndarray, laser_detuning: float):
+    """
+    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

src/rabifun/utilities.py

@@ -0,0 +1,13 @@
+import numpy as np
+
+
+def add_noise(signal: np.ndarray, noise_magnitude: float = 0.1):
+    """
+    Add Gaussian noise to ``signal``.  The standard deviation is
+    ``noise_magnitude`` relative to the mean absolute signal.
+
+    :returns: The noisy signal.
+    """
+
+    signal_magnitude = np.abs(signal).mean()
+    return signal + np.random.normal(0, signal_magnitude * noise_magnitude, len(signal))

src/ringfit/plotting.py

@@ -4,38 +4,7 @@ from scipy.ndimage import uniform_filter1d
 from scipy.interpolate import make_interp_spline
 import numpy as np
 
-
-def wrap_plot(f):
-    def wrapped(*args, ax=None, setup_function=plt.subplots, **kwargs):
-        fig = None
-        if not ax:
-            fig, ax = setup_function()
-
-        ret_val = f(*args, ax=ax, **kwargs)
-        return (fig, ax, ret_val) if ret_val else (fig, ax)
-
-    return wrapped
-
-
-def lighten_color(color, amount=0.5):
-    """
-    Lightens the given color by multiplying (1-luminosity) by the given amount.
-    Input can be matplotlib color string, hex string, or RGB tuple.
-
-    Examples:
-    >> lighten_color('g', 0.3)
-    >> lighten_color('#F034A3', 0.6)
-    >> lighten_color((.3,.55,.1), 0.5)
-    """
-    import matplotlib.colors as mc
-    import colorsys
-
-    try:
-        c = mc.cnames[color]
-    except:
-        c = color
-    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
-    return colorsys.hls_to_rgb(c[0], 1 - amount * (1 - c[1]), c[2])
+from plot_utils import *
 
 
 def fancy_error(x, y, err, ax, **kwargs):