RESULTS: make some nice demos

.gitignore

@@ -10,3 +10,5 @@ devenv.local.nix
 
 data
 __pycache__
+
+*.pkl

fitting_ringdown.org

@@ -1,38 +1,10 @@
 :PROPERTIES:
 :ID:       d7630955-6ca9-4de7-9770-2d50d4847bcd
 :END:
-#+title: Fitting Ringdown
+#+title: Python Code for the Fibre Loop Experiment
 
-* Fitting transients
+This is the code for the experimental side of [[id:bbd22b2d-4fcd-4ea9-9db7-8d599511a815][Simulation of Open
-- one can solve the equation of motion for when a mode is in the
+Systems with Fiber Loops]].
-  steady state and the laser is suddenly switched. one obtains a
-  decaying oscillation which can be fitted to the data.
-- this works ok, when the system /was/ in the steady state before the jump
-- see [[file:scripts/transients_without_amplification.py]]
-- we obtain \(γ\approx\SI{2.2}{\kilo\hertz}\) but this can vary a lot
-- the detuning frequency seems to be completely off
 
-* Failure of the steady state
+It contains some analysis code which is discussed in separate notes
-
+and a partial re-implementation of the [[id:c45097d2-2599-426d-82db-6ecfb5207151][Julia Code Project for the Non-Markovian Quantum Walk in Fiber Loops]].
-On the right peak on can see that there might be something fish going
-on. This is ~data/24_05_24/nice_transient2~.
-[[file:figures/non_steady.png]]
-
-* Weird Oscillations                                                 :ATTACH:
-- when turning on the amp while the laser is running we get some
-  oscillations at \(\sim \SI{40}{\kilo\hertz}\) which is a period of
-  \(\SI{25}{\micro\second}\). this is 500 times roundrip time :(). I
-  thought it might be a wave packet traveling around the loop, i.e
-  multiple modes together...
-
-[[attachment:osci_closeup.png][osci_closeup.png]]
-[[attachment:osci_far.png][osci_far.png]]
-[[attachment:cleaned.png][cleaned.png]]
-- it isn't close to the fsr though...
-- fitting a sine to that gives pretty much the same, as does an fft
-  [[attachment:cleaned_fit.png][cleaned_fit.png]]
-
-- it is pretty constant over all steps in [[file:data/09_05_24/Nicely_hybridised_2 2024,05,09, 15h57min00sec/][data]] as can be ascertained
-  from [[file:scripts/weird_oscillations.py][this script]]
-- these are also present w/o the amp cutout!
-[[attachment:osci_without.png][osci_without.png]]

scripts/experiments/004_full_nm_walk_tests.py

@@ -1,49 +1,105 @@
 from rabifun.system import *
 from rabifun.plots import *
 from rabifun.utilities import *
-
+import itertools
 # %% interactive
 
 
-def test_collective_mode_rabi():
+def make_params(ω_c=0.1, N=10):
-    """This is couples to a linear combination of bathe modes in stead of to a single one, but the result is pretty much the same."""
+    return Params(
-    ω_c = 0.1 / 2
-    params = Params(
         η=0.5,
         Ω=13,
         δ=1 / 4,
+        ω_c=ω_c,
         g_0=ω_c / 5,
         laser_detuning=0,
-        ω_c=ω_c,
+        N=2 * N + 2,
-        N=2,
+        N_couplings=N,
-        N_couplings=2,
         measurement_detuning=0,
-        α=0.1,
+        α=0,
         rwa=False,
-        flat_energies=True,
+        flat_energies=False,
+        correct_lamb_shift=True,
+        laser_off_time=0,
     )
 
-    params.laser_off_time = params.lifetimes(10)
-    # params.initial_state = make_zero_intial_state(params)
-    # params.initial_state[1] = 1
 
-    t = time_axis(params, 10, 0.01)
+def decay_rwa_analysis():
-    # solution = solve(t, params)
+    """This is couples to a linear combination of bathe modes in stead of to a single one, but the result is pretty much the same."""
-    # print(RuntimeParams(params))
-    # signal = output_signal(t, solution.y, params)
 
-    # f, (_, ax) = plot_simulation_result(
+    ω_c = 0.1
-    #     make_figure(), t, signal, params, window=(params.lifetimes(5), t[-1])
-    # )
-    # # ax.axvline(params.laser_off_time)
-    # plot_rabi_sidebands(ax, params)
 
-    fig = make_figure()
+    Ns = [1, 4]  # [5, 10, 20]
-    ax = fig.subplots()
 
-    sol_nonrwa, sol_nonrwa, *_ = plot_rwa_vs_real_amplitudes(ax, t, params)
+    fig = make_figure("decay_test")
-    ax.axvline(params.laser_off_time, color="black", linestyle="--")
+    ax_ns = fig.subplots(len(Ns), 2)
 
-    print(sol_nonrwa.y[2][-1] / sol_nonrwa.y[3][-1])
+    have_legend = False
-    runtime = RuntimeParams(params)
+
-    print(runtime.drive_amplitudes[0] / runtime.drive_amplitudes[1])
+    results = {}
+    param_dict = {}
+
+    for i, N in enumerate(Ns):
+        params = make_params(ω_c=ω_c, N=N)
+        params.laser_off_time = params.lifetimes(0)
+        params.initial_state = make_zero_intial_state(params)
+        params.initial_state[1] = 1
+
+        a_site = a_site_indices(params)[1]
+
+        ax_real, ax_corrected = ax_ns[i]
+
+        t = time_axis(params, recurrence_time(params) * 1.1 / params.lifetimes(1), 0.1)
+
+        for α in [0, 2]:
+            params.α = α
+            sol_nonrwa, sol_rwa = solve_nonrwa_rwa(t, params)
+
+            results[(N, α, True)] = sol_rwa
+            results[(N, α, False)] = sol_nonrwa
+            param_dict[(N, α)] = params
+
+            for correct, rwa in itertools.product([True, False], [True, False]):
+                sol = sol_rwa if rwa else sol_nonrwa
+                ax = ax_corrected if correct else ax_real
+                y = sol.y
+                if correct:
+                    y = correct_for_decay(sol, params)
+
+                ax.plot(
+                    sol.t,
+                    np.abs(y[a_site]) ** 2,
+                    label=f"{'rwa' if rwa else ''} α={α}",
+                    linestyle="--" if rwa else "-",
+                    alpha=0.5 if rwa else 1,
+                    color=f"C{α}",
+                )
+
+        ax_real.set_title(f"Real, N={N}")
+        ax_corrected.set_title(f"Decay Removed, N={N}")
+        for ax in [ax_real, ax_corrected]:
+            if not have_legend:
+                ax.legend()
+                have_legend = True
+
+            ax.set_xlabel("t [1/Ω]")
+            ax.set_ylabel("Population")
+            ax.set_ylim(0, 1)
+            ax.set_xlim(0, t[-1])
+            ax.axvline(recurrence_time(params), color="black", linestyle="--")
+
+    fig.tight_layout()
+    fig.suptitle(
+        f"Decay test for η={params.η}MHz, Ω={params.Ω}MHz, δ/Ω={params.δ}, ω_c/Ω={params.ω_c}"
+    )
+
+    save_figure(fig, "decay_test", extra_meta=dict(params=param_dict, Ns=Ns))
+    quick_save_pickle(
+        dict(results=results, params=param_dict, Ns=Ns),
+        "decay_test",
+        param_dict=param_dict,
+    )
+
+
+# %% make all figures
+decay_rwa_analysis()

scripts/experiments/figs/001_decay_test.pdf.meta.yaml

@@ -0,0 +1,108 @@
+change_id: snsvmqorzwxnvuvwsrsxqqmukqyspwkw
+commit_id: 8c5536ad27086daddb7a7e83aae06271b764fff5
+description: 'RESULTS: make some nice demos'
+extra_meta:
+  Ns:
+  - 1
+  - 4
+  params:
+    ? !!python/tuple
+    - 1
+    - 0
+    : &id001 !!python/object:rabifun.system.Params
+      N: 4
+      N_couplings: 1
+      correct_lamb_shift: true
+      drive_off_time: null
+      dynamic_detunting: &id002 !!python/tuple
+      - 0
+      - 0
+      flat_energies: false
+      g_0: 0.03333333333333333
+      initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
+        args:
+        - &id003 !!python/name:numpy.ndarray ''
+        - !!python/tuple
+          - 0
+        - !!binary |
+          Yg==
+        state: !!python/tuple
+        - 1
+        - !!python/tuple
+          - 6
+        - &id004 !!python/object/apply:numpy.dtype
+          args:
+          - c16
+          - false
+          - true
+          state: !!python/tuple
+          - 3
+          - <
+          - null
+          - null
+          - null
+          - -1
+          - -1
+          - 0
+        - false
+        - !!binary |
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+      laser_detuning: 0
+      laser_off_time: 0.0
+      measurement_detuning: 0
+      rwa: false
+      "\u03A9": 13
+      "\u03B1": 2
+      "\u03B4": 0.25
+      "\u03B7": 0.5
+      "\u03C9_c": 0.1
+    ? !!python/tuple
+    - 1
+    - 2
+    : *id001
+    ? !!python/tuple
+    - 4
+    - 0
+    : &id005 !!python/object:rabifun.system.Params
+      N: 10
+      N_couplings: 4
+      correct_lamb_shift: true
+      drive_off_time: null
+      dynamic_detunting: *id002
+      flat_energies: false
+      g_0: 0.03333333333333333
+      initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
+        args:
+        - *id003
+        - !!python/tuple
+          - 0
+        - !!binary |
+          Yg==
+        state: !!python/tuple
+        - 1
+        - !!python/tuple
+          - 12
+        - *id004
+        - false
+        - !!binary |
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAA
+      laser_detuning: 0
+      laser_off_time: 0.0
+      measurement_detuning: 0
+      rwa: false
+      "\u03A9": 13
+      "\u03B1": 2
+      "\u03B4": 0.25
+      "\u03B7": 0.5
+      "\u03C9_c": 0.1
+    ? !!python/tuple
+    - 4
+    - 2
+    : *id005
+name: 001_decay_test
+refers_to: ./figs/001_decay_test.pdf
+source: scripts/simulations/001_full_system_rwa_analysis.py

scripts/simulations/001_full_system_rwa_analysis.py

@@ -0,0 +1,121 @@
+from rabifun.system import *
+from rabifun.plots import *
+from rabifun.utilities import *
+import itertools
+# %% interactive
+
+
+def make_params(ω_c=0.1 / 2, N=10, gbar=1 / 3):
+    """
+    Make a set of parameters for the system with the current
+    best-known settings.
+    """
+    return Params(
+        η=0.5,
+        Ω=13,
+        δ=1 / 4,
+        ω_c=ω_c,
+        g_0=ω_c * gbar,
+        laser_detuning=0,
+        N=2 * N + 2,
+        N_couplings=N,
+        measurement_detuning=0,
+        α=0,
+        rwa=False,
+        flat_energies=False,
+        correct_lamb_shift=True,
+        laser_off_time=0,
+    )
+
+
+def decay_rwa_analysis():
+    """
+    Plot the population of the A site for different values of N and α,
+    with and without the RWA.
+
+    The parameters are chosen corresponding to the current best-known
+    values from the experiment.
+
+    The cutoff frequency is chosen to be 0.05Ω (which is the strongest
+    that still kinda works.).
+
+    Surprisingly the whole thing works for a gbar = 1/3 instead of 1/5
+    and with much fewer modes.
+    """
+
+    ω_c = 0.05
+    Ns = [5, 10, 20]
+    gbar = 1 / 3
+
+    fig = make_figure("decay_test", figsize=(15, len(Ns) * 3))
+    ax_ns = fig.subplots(len(Ns), 2)
+
+    have_legend = False
+
+    results = {}
+    param_dict = {}
+
+    for i, N in enumerate(Ns):
+        params = make_params(ω_c=ω_c, N=N, gbar=gbar)
+        params.laser_off_time = params.lifetimes(0)
+        params.initial_state = make_zero_intial_state(params)
+        params.initial_state[1] = 1
+
+        a_site = a_site_indices(params)[1]
+
+        ax_real, ax_corrected = ax_ns[i]
+
+        t = time_axis(params, recurrence_time(params) * 1.1 / params.lifetimes(1), 0.1)
+
+        for α in [0, 2]:
+            params.α = α
+            sol_nonrwa, sol_rwa = solve_nonrwa_rwa(t, params)
+
+            results[(N, α, True)] = sol_rwa
+            results[(N, α, False)] = sol_nonrwa
+            param_dict[(N, α)] = params
+
+            for correct, rwa in itertools.product([True, False], [True, False]):
+                sol = sol_rwa if rwa else sol_nonrwa
+                ax = ax_corrected if correct else ax_real
+                y = sol.y
+                if correct:
+                    y = correct_for_decay(sol, params)
+
+                ax.plot(
+                    sol.t,
+                    np.abs(y[a_site]) ** 2,
+                    label=f"{'rwa' if rwa else ''} α={α}",
+                    linestyle="--" if rwa else "-",
+                    alpha=0.5 if rwa else 1,
+                    color=f"C{α}",
+                )
+
+        ax_real.set_title(f"Real, N={N}")
+        ax_corrected.set_title(f"Decay Removed, N={N}")
+        for ax in [ax_real, ax_corrected]:
+            if not have_legend:
+                ax.legend()
+                have_legend = True
+
+            ax.set_xlabel("t [1/Ω]")
+            ax.set_ylabel("Population")
+            ax.set_ylim(0, 1)
+            ax.set_xlim(0, t[-1])
+            ax.axvline(recurrence_time(params), color="black", linestyle="--")
+
+    fig.tight_layout()
+    fig.suptitle(
+        f"Decay test for η={params.η}MHz, Ω={params.Ω}MHz, δ/Ω={params.δ}, ω_c/Ω={params.ω_c}, g_0/ω_c={params.g_0/params.ω_c:.2f}"
+    )
+
+    save_figure(fig, "001_decay_test", extra_meta=dict(params=param_dict, Ns=Ns))
+    quick_save_pickle(
+        dict(results=results, params=param_dict, Ns=Ns),
+        "001_decay_test",
+        param_dict=param_dict,
+    )
+
+
+# %% make all figures
+decay_rwa_analysis()

scripts/simulations/figs/001_decay_test.pdf.meta.yaml

@@ -0,0 +1,166 @@
+change_id: snsvmqorzwxnvuvwsrsxqqmukqyspwkw
+commit_id: 6ec1859dba940a2ada3aefedc5adea2cb436339f
+description: 'RESULTS: make some nice demos'
+extra_meta:
+  Ns:
+  - 5
+  - 10
+  - 20
+  params:
+    ? !!python/tuple
+    - 5
+    - 0
+    : &id001 !!python/object:rabifun.system.Params
+      N: 12
+      N_couplings: 5
+      correct_lamb_shift: true
+      drive_off_time: null
+      dynamic_detunting: &id002 !!python/tuple
+      - 0
+      - 0
+      flat_energies: false
+      g_0: 0.016666666666666666
+      initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
+        args:
+        - &id003 !!python/name:numpy.ndarray ''
+        - !!python/tuple
+          - 0
+        - !!binary |
+          Yg==
+        state: !!python/tuple
+        - 1
+        - !!python/tuple
+          - 14
+        - &id004 !!python/object/apply:numpy.dtype
+          args:
+          - c16
+          - false
+          - true
+          state: !!python/tuple
+          - 3
+          - <
+          - null
+          - null
+          - null
+          - -1
+          - -1
+          - 0
+        - false
+        - !!binary |
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=
+      laser_detuning: 0
+      laser_off_time: 0.0
+      measurement_detuning: 0
+      rwa: false
+      "\u03A9": 13
+      "\u03B1": 2
+      "\u03B4": 0.25
+      "\u03B7": 0.5
+      "\u03C9_c": 0.05
+    ? !!python/tuple
+    - 5
+    - 2
+    : *id001
+    ? !!python/tuple
+    - 10
+    - 0
+    : &id005 !!python/object:rabifun.system.Params
+      N: 22
+      N_couplings: 10
+      correct_lamb_shift: true
+      drive_off_time: null
+      dynamic_detunting: *id002
+      flat_energies: false
+      g_0: 0.016666666666666666
+      initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
+        args:
+        - *id003
+        - !!python/tuple
+          - 0
+        - !!binary |
+          Yg==
+        state: !!python/tuple
+        - 1
+        - !!python/tuple
+          - 24
+        - *id004
+        - false
+        - !!binary |
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+      laser_detuning: 0
+      laser_off_time: 0.0
+      measurement_detuning: 0
+      rwa: false
+      "\u03A9": 13
+      "\u03B1": 2
+      "\u03B4": 0.25
+      "\u03B7": 0.5
+      "\u03C9_c": 0.05
+    ? !!python/tuple
+    - 10
+    - 2
+    : *id005
+    ? !!python/tuple
+    - 20
+    - 0
+    : &id006 !!python/object:rabifun.system.Params
+      N: 42
+      N_couplings: 20
+      correct_lamb_shift: true
+      drive_off_time: null
+      dynamic_detunting: *id002
+      flat_energies: false
+      g_0: 0.016666666666666666
+      initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
+        args:
+        - *id003
+        - !!python/tuple
+          - 0
+        - !!binary |
+          Yg==
+        state: !!python/tuple
+        - 1
+        - !!python/tuple
+          - 44
+        - *id004
+        - false
+        - !!binary |
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAA=
+      laser_detuning: 0
+      laser_off_time: 0.0
+      measurement_detuning: 0
+      rwa: false
+      "\u03A9": 13
+      "\u03B1": 2
+      "\u03B4": 0.25
+      "\u03B7": 0.5
+      "\u03C9_c": 0.05
+    ? !!python/tuple
+    - 20
+    - 2
+    : *id006
+function: decay_rwa_analysis
+name: 001_decay_test
+refers_to: ./figs/001_decay_test.pdf
+source: scripts/simulations/001_full_system_rwa_analysis.py

scripts/simulations/outputs/001_decay_test.pkl.meta.yaml

@@ -0,0 +1,160 @@
+change_id: snsvmqorzwxnvuvwsrsxqqmukqyspwkw
+commit_id: 049b547f220987c6bb63165a9a020b20b1c4377b
+description: 'RESULTS: make some nice demos'
+function: decay_rwa_analysis
+param_dict:
+  ? !!python/tuple
+  - 5
+  - 0
+  : &id001 !!python/object:rabifun.system.Params
+    N: 12
+    N_couplings: 5
+    correct_lamb_shift: true
+    drive_off_time: null
+    dynamic_detunting: &id002 !!python/tuple
+    - 0
+    - 0
+    flat_energies: false
+    g_0: 0.016666666666666666
+    initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
+      args:
+      - &id003 !!python/name:numpy.ndarray ''
+      - !!python/tuple
+        - 0
+      - !!binary |
+        Yg==
+      state: !!python/tuple
+      - 1
+      - !!python/tuple
+        - 14
+      - &id004 !!python/object/apply:numpy.dtype
+        args:
+        - c16
+        - false
+        - true
+        state: !!python/tuple
+        - 3
+        - <
+        - null
+        - null
+        - null
+        - -1
+        - -1
+        - 0
+      - false
+      - !!binary |
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=
+    laser_detuning: 0
+    laser_off_time: 0.0
+    measurement_detuning: 0
+    rwa: false
+    "\u03A9": 13
+    "\u03B1": 2
+    "\u03B4": 0.25
+    "\u03B7": 0.5
+    "\u03C9_c": 0.05
+  ? !!python/tuple
+  - 5
+  - 2
+  : *id001
+  ? !!python/tuple
+  - 10
+  - 0
+  : &id005 !!python/object:rabifun.system.Params
+    N: 22
+    N_couplings: 10
+    correct_lamb_shift: true
+    drive_off_time: null
+    dynamic_detunting: *id002
+    flat_energies: false
+    g_0: 0.016666666666666666
+    initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
+      args:
+      - *id003
+      - !!python/tuple
+        - 0
+      - !!binary |
+        Yg==
+      state: !!python/tuple
+      - 1
+      - !!python/tuple
+        - 24
+      - *id004
+      - false
+      - !!binary |
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+    laser_detuning: 0
+    laser_off_time: 0.0
+    measurement_detuning: 0
+    rwa: false
+    "\u03A9": 13
+    "\u03B1": 2
+    "\u03B4": 0.25
+    "\u03B7": 0.5
+    "\u03C9_c": 0.05
+  ? !!python/tuple
+  - 10
+  - 2
+  : *id005
+  ? !!python/tuple
+  - 20
+  - 0
+  : &id006 !!python/object:rabifun.system.Params
+    N: 42
+    N_couplings: 20
+    correct_lamb_shift: true
+    drive_off_time: null
+    dynamic_detunting: *id002
+    flat_energies: false
+    g_0: 0.016666666666666666
+    initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
+      args:
+      - *id003
+      - !!python/tuple
+        - 0
+      - !!binary |
+        Yg==
+      state: !!python/tuple
+      - 1
+      - !!python/tuple
+        - 44
+      - *id004
+      - false
+      - !!binary |
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAA=
+    laser_detuning: 0
+    laser_off_time: 0.0
+    measurement_detuning: 0
+    rwa: false
+    "\u03A9": 13
+    "\u03B1": 2
+    "\u03B4": 0.25
+    "\u03B7": 0.5
+    "\u03C9_c": 0.05
+  ? !!python/tuple
+  - 20
+  - 2
+  : *id006
+refers_to: outputs/001_decay_test.pkl
+source: scripts/simulations/001_full_system_rwa_analysis.py

src/plot_utils.py

@@ -6,11 +6,23 @@ import yaml
 import inspect
 import subprocess
 import pathlib
+import sys
 
 P = ParamSpec("P")
 R = TypeVar("R")
 
 
+def noop_if_interactive(f):
+    @wraps(f)
+    def wrapped(*args, **kwargs):
+        if hasattr(sys, "ps1"):
+            return
+
+        return f(*args, **kwargs)
+
+    return wrapped
+
+
 def make_figure(fig_name: str = "interactive", *args, **kwargs):
     fig = plt.figure(fig_name, *args, **kwargs)
     fig.clf()
@@ -94,19 +106,21 @@ def write_meta(path, **kwargs):
         .strip()
     )
 
-    frame = inspect.stack()[1]
+    frame = inspect.stack()[3]
     module = inspect.getmodule(frame[0])
     filename = str(
         pathlib.Path(module.__file__).relative_to(project_dir)  # type: ignore
         if module
         else "<unknown>"
     )
+    function = frame.function
 
     outpath = f"{path}.meta.yaml"
     with open(outpath, "w") as f:
         yaml.dump(
             dict(
                 source=filename,
+                function=function,
                 change_id=change_id,
                 commit_id=commit_id,
                 description=description.strip(),
@@ -119,19 +133,21 @@ def write_meta(path, **kwargs):
     print(f"Metadata written to {outpath}")
 
 
-def save_figure(fig, name, *args, **kwargs):
+@noop_if_interactive
+def save_figure(fig, name, extra_meta=None, *args, **kwargs):
     dir = pathlib.Path(f"./figs/")
     dir.mkdir(exist_ok=True)
     fig.tight_layout()
 
-    write_meta(f"./figs/{name}.pdf", name=name)
+    write_meta(f"./figs/{name}.pdf", name=name, extra_meta=extra_meta)
 
     plt.savefig(f"./figs/{name}.pdf", *args, **kwargs)
     plt.savefig(f"./figs/{name}.png", *args, dpi=600, **kwargs)
 
-    print(f"Figure saved as {name}.pdf")
+    print(f"Figure saved as ./figs/{name}.pdf")
 
 
+@noop_if_interactive
 def quick_save_pickle(obj, name, **kwargs):
     """Quickly save an object to a pickle file with metadata."""
     import pickle

src/rabifun/plots.py

@@ -1,5 +1,13 @@
 from plot_utils import *
-from .system import Params, RuntimeParams, solve, coupled_mode_indices, mode_name
+from .system import (
+    Params,
+    RuntimeParams,
+    solve,
+    coupled_mode_indices,
+    mode_name,
+    uncoupled_mode_indices,
+    correct_for_decay,
+)
 from .analysis import fourier_transform
 import matplotlib.pyplot as plt
 import numpy as np
@@ -114,7 +122,15 @@ def clone_color_or_default(lines: dict, mode: int):
     return line.get_color()
 
 
-def plot_rotating_modes(ax, solution, params, plot_uncoupled=False, clone_colors=None):
+def plot_rotating_modes(
+    ax,
+    solution,
+    params,
+    plot_uncoupled=False,
+    clone_colors=None,
+    only_A_site=False,
+    correct_for_decay=False,
+):
     """Plot the amplitude of the modes in the rotating frame.
 
     :param ax: The axis to plot on.
@@ -123,63 +139,74 @@ def plot_rotating_modes(ax, solution, params, plot_uncoupled=False, clone_colors
     :param plot_uncoupled: Whether to plot the uncoupled modes.
     :param clone_colors: A dictionary of lines indexed by mode index
         from which to clone colors from.
+    :param only_A_site: Whether to plot only the A site modes.
+    :param correct_for_decay: Whether to correct for decay by
+        multiplying with an exponential.
     """
 
     lines = dict()
     if clone_colors is None:
         clone_colors = dict()
 
-    for mode in coupled_mode_indices(params):
+    h = solution.y
+    if correct_for_decay:
+        h = correct_for_decay(solution, params)
+
+    for mode in [1] if only_A_site else coupled_mode_indices(params):
         lines[mode] = ax.plot(
             solution.t,
-            np.abs(solution.y[mode]),
+            np.abs(h[mode]),
             label=mode_name(mode) + (" (rwa)" if params.rwa else ""),
             color=clone_color_or_default(clone_colors, mode),
             linestyle="dashdot" if params.rwa else "-",
         )[0]
 
-    if plot_uncoupled:
+    if plot_uncoupled and not only_A_site:
         for mode in uncoupled_mode_indices(params):
             lines[mode] = ax.plot(
                 solution.t,
-                np.abs(solution.y[mode]),
+                np.abs(h[mode]),
                 label=mode_name(mode) + (" (rwa)" if params.rwa else ""),
                 color=clone_color_or_default(clone_colors, mode),
                 linestyle="dotted" if params.rwa else "--",
             )[0]
 
-    ax.legend()
+    # ax.legend()
     ax.set_xlabel("Time (1/Ω)")
     ax.set_ylabel("Amplitude")
-    ax.set_title("Mode Amplitudes in the Rotating Frame")
+    ax.set_title(
+        "Mode Amplitudes in the Rotating Frame"
+        + (" (corrected)" if correct_for_decay else "")
+    )
 
     return lines
 
 
-def plot_rwa_vs_real_amplitudes(ax, t, params, **kwargs):
+def plot_rwa_vs_real_amplitudes(ax, solution_no_rwa, solution_rwa, params, **kwargs):
     """Plot the amplitudes of the modes of the system ``params`` in
     the rotating frame with and without the RWA onto ``ax``.
 
     The keyword arguments are passed to :any:`plot_rotating_modes`.
 
     :param ax: The axis to plot on.
-    :param t: The time axis.
+    :param non_rwa: The solution without the rwa.
+    :param rwa: The solution with the rwa.
     :param params: The system parameters.
     :param kwargs: Additional keyword arguments to pass to
         :any:`plot_rotating_modes`.
 
-    :returns: A tuple of the solutions without and with the rwa and
+    :returns: A tuple of the line dictionaries for the non-rwa and rwa solutions.
-              the dictionaries of the lines plotted indexed by mode
-              index.
     """
+
+    original_rwa = params.rwa
     params.rwa = False
-    solution_no_rwa = solve(t, params)
     no_rwa_lines = plot_rotating_modes(ax, solution_no_rwa, params, **kwargs)
 
     params.rwa = True
-    solution_rwa = solve(t, params)
     rwa_lines = plot_rotating_modes(
         ax, solution_rwa, params, **kwargs, clone_colors=no_rwa_lines
     )
 
-    return solution_no_rwa, solution_rwa, no_rwa_lines, rwa_lines
+    params.rwa = original_rwa
+
+    return no_rwa_lines, rwa_lines

src/rabifun/system.py

@@ -18,13 +18,13 @@ class Params:
     complement of modes.
     """
 
-    Ω: float = 1
+    Ω: float = 13
     """Free spectral range of the system in *frequency units*."""
 
     δ: float = 1 / 4
     """Mode splitting in units of :any:`Ω`."""
 
-    η: float = 0.1
+    η: float = 0.5
     """Decay rate :math:`\eta/2` of the system in angular frequency units."""
 
     g_0: float = 0.01
@@ -61,6 +61,10 @@ class Params:
     """Whether to use a flat distribution of bath energies."""
 
     initial_state: np.ndarray | None = None
+    """The initial state of the system."""
+
+    correct_lamb_shift: bool = True
+    """Whether to correct for the Lamb shift by tweaking the detuning."""
 
     def __post_init__(self):
         if self.N_couplings > self.N:
@@ -81,7 +85,7 @@ class Params:
         Returns the number of lifetimes of the system that correspond to
         `n` cycles.
         """
-        return n / self.η
+        return 2 * n / self.η
 
     @property
     def rabi_splitting(self):
@@ -102,9 +106,14 @@ class RuntimeParams:
             * params.Ω
             * np.concatenate([[-1 * params.δ, params.δ], np.arange(1, params.N + 1)])
         )
+
         decay_rates = -1j * np.repeat(params.η / 2, params.N + 2)
         Ωs = freqs + decay_rates
 
+        self.drive_frequencies, self.detunings, self.drive_amplitudes = (
+            drive_frequencies_and_amplitudes(params)
+        )  # linear frequencies!
+
         self.Ωs = Ωs
         self.diag_energies = (
             2
@@ -112,17 +121,14 @@ class RuntimeParams:
             * np.concatenate(
                 [
                     [0, 0],
-                    drive_detunings(params),
+                    self.detunings,
                     np.zeros(params.N - params.N_couplings),
                 ]
             )
             + decay_rates
         )
-        self.detuned_Ωs = freqs - self.diag_energies.real
 
-        self.drive_frequencies, self.drive_amplitudes = (
+        self.detuned_Ωs = freqs - self.diag_energies.real
-            drive_frequencies_and_amplitudes(params)
-        )
 
     def __repr__(self):
         return f"{self.__class__.__name__}(Ωs={self.Ωs}, drive_frequencies={self.drive_frequencies}, drive_amplitudes={self.drive_amplitudes})"
@@ -154,14 +160,22 @@ def eom_drive(t, x, ds, ωs, rwa, detuned_Ωs):
     :param ωs: linear drive frequencies
     """
 
+    ds = 2 * np.pi * ds
     if rwa:
+        coupled_indices = 2 + len(ds)
         det_matrix = np.zeros((len(x), len(x)))
-        det_matrix[1, 2:] = ds / 2
+        det_matrix[1, 2:coupled_indices] = ds / 2
-        det_matrix[2:, 1] = ds / 2
+        det_matrix[2:coupled_indices, 1] = ds / 2
         driven_x = det_matrix @ x
     else:
-        freqs = np.exp(-1j * detuned_Ωs * t)
+        det_matrix = detuned_Ωs[:, None] - detuned_Ωs[None, :]
-        det_matrix = np.outer(np.conj(freqs), freqs)
+        # test = abs(det_matrix.copy())
+        # test[test < 1e-10] = np.inf
+        # print(np.min(test))
+        # print(np.argmin(test, keepdims=True))
+
+        det_matrix = np.exp(-1j * det_matrix * t)
+
         driven_x = np.sum(ds * np.sin(2 * np.pi * ωs * t)) * (det_matrix @ x)
 
     return driven_x
@@ -283,8 +297,12 @@ def output_signal(t: np.ndarray, amplitudes: np.ndarray, params: Params):
     Calculate the output signal when mixing with laser light of
     frequency `laser_detuning`.
     """
+
+    runtime = RuntimeParams(params)
+    rotating = amplitudes * np.exp(-1j * runtime.detuned_Ωs * t)
+
     return (
-        np.sum(amplitudes, axis=0)
+        np.sum(rotating, axis=0)
         * np.exp(
             1j
             * (laser_frequency(params, t) + 2 * np.pi * params.measurement_detuning)
@@ -305,24 +323,24 @@ def ohmic_spectral_density(ω: np.ndarray, α: float) -> np.ndarray:
     return ω**α
 
 
-def drive_detunings(params: Params) -> np.ndarray:
+def lamb_shift(amplitudes, Δs):
-    """Return the drive detunings of the bath modes in frequency units."""
+    return np.sum(amplitudes**2 / Δs)
+
+
+def drive_frequencies_and_amplitudes(
+    params: Params,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Return the linear frequencies and amplitudes of the drives based
+    on the ``params``.
+    """
 
     if params.flat_energies:
         Δs = np.repeat(params.ω_c, params.N_couplings)
     else:
         Δs = bath_energies(params.N_couplings, params.ω_c)
 
-    return Δs * params.Ω
+    Δs *= params.Ω
-
-
-def drive_frequencies_and_amplitudes(params: Params) -> tuple[np.ndarray, np.ndarray]:
-    """
-    Return the linear frequencies and amplitudes of the drives based
-    on the ``params``.
-    """
-
-    Δs = drive_detunings(params)
 
     amplitudes = ohmic_spectral_density(
         bath_energies(params.N_couplings, 1),
@@ -330,10 +348,13 @@ def drive_frequencies_and_amplitudes(params: Params) -> tuple[np.ndarray, np.nda
     )
 
     amplitudes /= np.sum(amplitudes)
-    amplitudes = 2 * np.pi * params.Ω * params.g_0 * np.sqrt(amplitudes)
+    amplitudes = params.Ω * params.g_0 * np.sqrt(amplitudes)
+
+    if not params.flat_energies and params.correct_lamb_shift:
+        Δs -= np.sum(amplitudes**2 / Δs)
 
     ωs = ((np.arange(1, params.N_couplings + 1) - params.δ) * params.Ω) - Δs
-    return ωs, amplitudes
+    return ωs, Δs, amplitudes
 
 
 def mode_name(mode: int):
@@ -378,3 +399,41 @@ def coupled_mode_indices(params: Params):
 def dimension(params: Params):
     """Return the dimension of the system."""
     return params.N + 2
+
+
+def recurrence_time(params: Params):
+    """Return the recurrence time of the system."""
+    return params.N_couplings / (params.Ω * params.ω_c)
+
+
+def solve_nonrwa_rwa(t: np.ndarray, params: Params, **kwargs):
+    """
+    Solve the system in the non-RWA and RWA cases and return the results.
+    The keyword arguments are passed to :any:`solve`.
+
+    :param t: time array
+    :param params: system parameters
+
+    :returns: non-RWA and RWA solutions
+    """
+    initial_rwa = params.rwa
+
+    params.rwa = False
+    nonrwa = solve(t, params, **kwargs)
+    rwa_params = params
+    rwa_params.rwa = True
+    rwa = solve(t, rwa_params, **kwargs)
+
+    params.rwa = initial_rwa
+    return nonrwa, rwa
+
+
+def correct_for_decay(solution, params):
+    """Correct the ``solution`` for decay.
+
+    :param solution: The solution from :any:`solve_ivp` to correct.
+    :param params: The system parameters
+
+    :returns: The corrected solution amplitudes.
+    """
+    return solution.y * np.exp(params.η / 2 * solution.t[None, :])