diff --git a/scripts/experiments/005_mean_displacement_semi_analytic.py b/scripts/experiments/005_mean_displacement_semi_analytic.py
new file mode 100644
index 0000000..4e3434a
--- /dev/null
+++ b/scripts/experiments/005_mean_displacement_semi_analytic.py
@@ -0,0 +1,104 @@
+"""
+Just a quick hack prior to the re-calculation of the full phase
+diagram.
+"""
+
+import numpy as np
+from rabifun.system import *
+from rabifun.plots import *
+from rabifun.utilities import *
+import scipy
+import matplotlib.pyplot as plt
+
+
+def self_energy(ω, g, ε):
+    """
+    Calculate the self-energy of the system.
+
+    :param g: The coupling strengths.
+    :param ε: The bath energies.
+    """
+    return np.sum(np.abs(g[None, :] / (ω[:, None] - ε[None, :])) ** 2, axis=1)
+
+
+def coefficients(ω, g, ε):
+    coeff = 1 / (1 + self_energy(ω, g, ε))
+
+    return coeff / np.sum((coeff))
+
+
+def characteristic_poly(g: np.ndarray, ε: np.ndarray, η_A: float = 0):
+    g2 = np.abs(g) ** 2
+
+    def poly(ω):
+        s = ω[0] + ω[1] * 1j
+        res = s + 1j * η_A - np.sum(g2 / (s - ε))
+        return (res * res.conjugate()).real
+
+    return poly
+
+
+def hamiltonian(g: np.ndarray, ε: np.ndarray, η_A: float = 0):
+    H = np.diag([-1j * η_A, *ε])
+    H[0, 1:] = g
+    H[1:, 0] = np.conj(g)
+    return H
+
+
+def a_site_population(t, ω, coeff):
+    return (
+        np.abs(np.sum(coeff[None, :] * np.exp(-1j * ω[None, :] * t[:, None]), axis=1))
+        ** 2
+    )
+
+
+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 test():
+    params = make_params(N=10, gbar=1 / 3)
+    params.flat_energies = False
+    params.α = 0.0
+    params.correct_lamb_shift = True
+
+    runtime = RuntimeParams(params)
+    t = time_axis(params, recurrences=1.5)
+    g = runtime.g / 2
+    ε = runtime.bath_ε
+    H = hamiltonian(g, ε.real, 0 * params.η / 2)
+
+    ω = np.linalg.eigvals(H)
+    M = np.linalg.eig(H).eigenvectors
+    Minv = np.linalg.inv(M)
+    v0 = Minv[:, 0]
+
+    coeff = M[0, :] * v0.T
+    print(coeff)
+
+    # coeff /= np.abs(np.sum(coeff))
+    # coeff = coefficients(ω, g, ε)
+
+    plt.cla()
+    plt.plot(t, a_site_population(t, ω, coeff))
+    plt.ylim(0, 1)
+    return ω
diff --git a/scripts/simulations/001_full_system_rwa_analysis.py b/scripts/simulations/001_full_system_rwa_analysis.py
index 7702c7d..08e5c93 100644
--- a/scripts/simulations/001_full_system_rwa_analysis.py
+++ b/scripts/simulations/001_full_system_rwa_analysis.py
@@ -5,7 +5,7 @@ import itertools
 # %% interactive
 
 
-def make_params(ω_c=0.1 / 2, N=10, gbar=1 / 3):
+def make_params(ω_c=0.1 / 2, N=10, gbar=1 / 3, compensate=2):
     """
     Make a set of parameters for the system with the current
     best-known settings.
@@ -23,7 +23,7 @@ def make_params(ω_c=0.1 / 2, N=10, gbar=1 / 3):
         α=0,
         rwa=False,
         flat_energies=False,
-        correct_lamb_shift=True,
+        correct_lamb_shift=compensate,
         laser_off_time=0,
     )
 
@@ -45,7 +45,7 @@ def decay_rwa_analysis():
 
     ω_c = 0.05
     Ns = [5, 10, 20]
-    gbar = 1 / 3
+    gbar = 1 / 3 / 2
 
     fig = make_figure("decay_test", figsize=(15, len(Ns) * 3))
     ax_ns = fig.subplots(len(Ns), 2)
@@ -57,7 +57,7 @@ def decay_rwa_analysis():
 
     for i, N in enumerate(Ns):
         params = make_params(ω_c=ω_c, N=N, gbar=gbar)
-        params.laser_off_time = params.lifetimes(0)
+        params.laser_off_time = 0
         params.initial_state = make_zero_intial_state(params)
         params.initial_state[1] = 1
 
@@ -65,9 +65,9 @@ def decay_rwa_analysis():
 
         ax_real, ax_corrected = ax_ns[i]
 
-        t = time_axis(params, recurrence_time(params) * 1.1 / params.lifetimes(1), 0.1)
+        t = time_axis(params, recurrences=1.1, resolution=0.1)
 
-        for α in [0, 2]:
+        for α in np.linspace(0, 2, 5):
             params.α = α
             sol_nonrwa, sol_rwa = solve_nonrwa_rwa(t, params)
 
@@ -75,6 +75,7 @@ def decay_rwa_analysis():
             results[(N, α, False)] = sol_nonrwa
             param_dict[(N, α)] = params
 
+            color = None
             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
@@ -82,14 +83,15 @@ def decay_rwa_analysis():
                 if correct:
                     y = correct_for_decay(sol, params)
 
-                ax.plot(
+                l = 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{α}",
+                    color=color,
                 )
+                color = l[0].get_color()
 
         ax_real.set_title(f"Real, N={N}")
         ax_corrected.set_title(f"Decay Removed, N={N}")
diff --git a/scripts/simulations/figs/001_decay_test.pdf b/scripts/simulations/figs/001_decay_test.pdf
index ec56749..e2679d5 100644
Binary files a/scripts/simulations/figs/001_decay_test.pdf and b/scripts/simulations/figs/001_decay_test.pdf differ
diff --git a/scripts/simulations/figs/001_decay_test.pdf.meta.yaml b/scripts/simulations/figs/001_decay_test.pdf.meta.yaml
index f745af2..8f4fca3 100644
--- a/scripts/simulations/figs/001_decay_test.pdf.meta.yaml
+++ b/scripts/simulations/figs/001_decay_test.pdf.meta.yaml
@@ -1,6 +1,6 @@
-change_id: snsvmqorzwxnvuvwsrsxqqmukqyspwkw
-commit_id: 6ec1859dba940a2ada3aefedc5adea2cb436339f
-description: 'RESULTS: make some nice demos'
+change_id: wykwpsvmtrzvwtunotxnzokztpyrnsrx
+commit_id: 94f23ecee30914b5163f14fd5af637a8777b42cc
+description: ''
 extra_meta:
   Ns:
   - 5
@@ -9,20 +9,36 @@ extra_meta:
   params:
     ? !!python/tuple
     - 5
-    - 0
-    : &id001 !!python/object:rabifun.system.Params
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - &id001 !!python/object/apply:numpy.dtype
+        args:
+        - f8
+        - false
+        - true
+        state: !!python/tuple
+        - 3
+        - <
+        - null
+        - null
+        - null
+        - -1
+        - -1
+        - 0
+      - !!binary |
+        AAAAAAAAAAA=
+    : &id002 !!python/object:rabifun.system.Params
       N: 12
       N_couplings: 5
-      correct_lamb_shift: true
+      correct_lamb_shift: 2
       drive_off_time: null
-      dynamic_detunting: &id002 !!python/tuple
+      dynamic_detunting: &id003 !!python/tuple
       - 0
       - 0
       flat_energies: false
-      g_0: 0.016666666666666666
+      g_0: 0.008333333333333333
       initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
         args:
-        - &id003 !!python/name:numpy.ndarray ''
+        - &id004 !!python/name:numpy.ndarray ''
         - !!python/tuple
           - 0
         - !!binary |
@@ -31,7 +47,7 @@ extra_meta:
         - 1
         - !!python/tuple
           - 14
-        - &id004 !!python/object/apply:numpy.dtype
+        - &id005 !!python/object/apply:numpy.dtype
           args:
           - c16
           - false
@@ -52,32 +68,62 @@ extra_meta:
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=
       laser_detuning: 0
-      laser_off_time: 0.0
+      laser_off_time: 0
       measurement_detuning: 0
       rwa: false
       "\u03A9": 13
-      "\u03B1": 2
+      "\u03B1": !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAAAEA=
       "\u03B4": 0.25
       "\u03B7": 0.5
       "\u03C9_c": 0.05
     ? !!python/tuple
     - 5
-    - 2
-    : *id001
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA4D8=
+    : *id002
+    ? !!python/tuple
+    - 5
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA8D8=
+    : *id002
+    ? !!python/tuple
+    - 5
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA+D8=
+    : *id002
+    ? !!python/tuple
+    - 5
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAAAEA=
+    : *id002
     ? !!python/tuple
     - 10
-    - 0
-    : &id005 !!python/object:rabifun.system.Params
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAAAAA=
+    : &id006 !!python/object:rabifun.system.Params
       N: 22
       N_couplings: 10
-      correct_lamb_shift: true
+      correct_lamb_shift: 2
       drive_off_time: null
-      dynamic_detunting: *id002
+      dynamic_detunting: *id003
       flat_energies: false
-      g_0: 0.016666666666666666
+      g_0: 0.008333333333333333
       initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
         args:
-        - *id003
+        - *id004
         - !!python/tuple
           - 0
         - !!binary |
@@ -86,7 +132,7 @@ extra_meta:
         - 1
         - !!python/tuple
           - 24
-        - *id004
+        - *id005
         - false
         - !!binary |
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
@@ -97,32 +143,62 @@ extra_meta:
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
       laser_detuning: 0
-      laser_off_time: 0.0
+      laser_off_time: 0
       measurement_detuning: 0
       rwa: false
       "\u03A9": 13
-      "\u03B1": 2
+      "\u03B1": !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAAAEA=
       "\u03B4": 0.25
       "\u03B7": 0.5
       "\u03C9_c": 0.05
     ? !!python/tuple
     - 10
-    - 2
-    : *id005
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA4D8=
+    : *id006
+    ? !!python/tuple
+    - 10
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA8D8=
+    : *id006
+    ? !!python/tuple
+    - 10
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA+D8=
+    : *id006
+    ? !!python/tuple
+    - 10
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAAAEA=
+    : *id006
     ? !!python/tuple
     - 20
-    - 0
-    : &id006 !!python/object:rabifun.system.Params
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAAAAA=
+    : &id007 !!python/object:rabifun.system.Params
       N: 42
       N_couplings: 20
-      correct_lamb_shift: true
+      correct_lamb_shift: 2
       drive_off_time: null
-      dynamic_detunting: *id002
+      dynamic_detunting: *id003
       flat_energies: false
-      g_0: 0.016666666666666666
+      g_0: 0.008333333333333333
       initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
         args:
-        - *id003
+        - *id004
         - !!python/tuple
           - 0
         - !!binary |
@@ -131,7 +207,7 @@ extra_meta:
         - 1
         - !!python/tuple
           - 44
-        - *id004
+        - *id005
         - false
         - !!binary |
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
@@ -148,18 +224,45 @@ extra_meta:
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
           AAAAAAAAAAAAAAAAAAAAAAAAAAA=
       laser_detuning: 0
-      laser_off_time: 0.0
+      laser_off_time: 0
       measurement_detuning: 0
       rwa: false
       "\u03A9": 13
-      "\u03B1": 2
+      "\u03B1": !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAAAEA=
       "\u03B4": 0.25
       "\u03B7": 0.5
       "\u03C9_c": 0.05
     ? !!python/tuple
     - 20
-    - 2
-    : *id006
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA4D8=
+    : *id007
+    ? !!python/tuple
+    - 20
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA8D8=
+    : *id007
+    ? !!python/tuple
+    - 20
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAA+D8=
+    : *id007
+    ? !!python/tuple
+    - 20
+    - !!python/object/apply:numpy.core.multiarray.scalar
+      - *id001
+      - !!binary |
+        AAAAAAAAAEA=
+    : *id007
 function: decay_rwa_analysis
 name: 001_decay_test
 refers_to: ./figs/001_decay_test.pdf
diff --git a/scripts/simulations/figs/001_decay_test.png b/scripts/simulations/figs/001_decay_test.png
index 9b586bc..a417183 100644
Binary files a/scripts/simulations/figs/001_decay_test.png and b/scripts/simulations/figs/001_decay_test.png differ
diff --git a/scripts/simulations/outputs/001_decay_test.pkl.meta.yaml b/scripts/simulations/outputs/001_decay_test.pkl.meta.yaml
index dfceb90..d84ee19 100644
--- a/scripts/simulations/outputs/001_decay_test.pkl.meta.yaml
+++ b/scripts/simulations/outputs/001_decay_test.pkl.meta.yaml
@@ -1,24 +1,40 @@
-change_id: snsvmqorzwxnvuvwsrsxqqmukqyspwkw
-commit_id: 049b547f220987c6bb63165a9a020b20b1c4377b
-description: 'RESULTS: make some nice demos'
+change_id: wykwpsvmtrzvwtunotxnzokztpyrnsrx
+commit_id: e4399f86a81e8dc0fbbe24226b948db60bd83807
+description: ''
 function: decay_rwa_analysis
 param_dict:
   ? !!python/tuple
   - 5
-  - 0
-  : &id001 !!python/object:rabifun.system.Params
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - &id001 !!python/object/apply:numpy.dtype
+      args:
+      - f8
+      - false
+      - true
+      state: !!python/tuple
+      - 3
+      - <
+      - null
+      - null
+      - null
+      - -1
+      - -1
+      - 0
+    - !!binary |
+      AAAAAAAAAAA=
+  : &id002 !!python/object:rabifun.system.Params
     N: 12
     N_couplings: 5
-    correct_lamb_shift: true
+    correct_lamb_shift: 2
     drive_off_time: null
-    dynamic_detunting: &id002 !!python/tuple
+    dynamic_detunting: &id003 !!python/tuple
     - 0
     - 0
     flat_energies: false
-    g_0: 0.016666666666666666
+    g_0: 0.008333333333333333
     initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
       args:
-      - &id003 !!python/name:numpy.ndarray ''
+      - &id004 !!python/name:numpy.ndarray ''
       - !!python/tuple
         - 0
       - !!binary |
@@ -27,7 +43,7 @@ param_dict:
       - 1
       - !!python/tuple
         - 14
-      - &id004 !!python/object/apply:numpy.dtype
+      - &id005 !!python/object/apply:numpy.dtype
         args:
         - c16
         - false
@@ -48,32 +64,62 @@ param_dict:
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=
     laser_detuning: 0
-    laser_off_time: 0.0
+    laser_off_time: 0
     measurement_detuning: 0
     rwa: false
     "\u03A9": 13
-    "\u03B1": 2
+    "\u03B1": !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAAAEA=
     "\u03B4": 0.25
     "\u03B7": 0.5
     "\u03C9_c": 0.05
   ? !!python/tuple
   - 5
-  - 2
-  : *id001
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA4D8=
+  : *id002
+  ? !!python/tuple
+  - 5
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA8D8=
+  : *id002
+  ? !!python/tuple
+  - 5
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA+D8=
+  : *id002
+  ? !!python/tuple
+  - 5
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAAAEA=
+  : *id002
   ? !!python/tuple
   - 10
-  - 0
-  : &id005 !!python/object:rabifun.system.Params
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAAAAA=
+  : &id006 !!python/object:rabifun.system.Params
     N: 22
     N_couplings: 10
-    correct_lamb_shift: true
+    correct_lamb_shift: 2
     drive_off_time: null
-    dynamic_detunting: *id002
+    dynamic_detunting: *id003
     flat_energies: false
-    g_0: 0.016666666666666666
+    g_0: 0.008333333333333333
     initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
       args:
-      - *id003
+      - *id004
       - !!python/tuple
         - 0
       - !!binary |
@@ -82,7 +128,7 @@ param_dict:
       - 1
       - !!python/tuple
         - 24
-      - *id004
+      - *id005
       - false
       - !!binary |
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
@@ -93,32 +139,62 @@ param_dict:
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
     laser_detuning: 0
-    laser_off_time: 0.0
+    laser_off_time: 0
     measurement_detuning: 0
     rwa: false
     "\u03A9": 13
-    "\u03B1": 2
+    "\u03B1": !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAAAEA=
     "\u03B4": 0.25
     "\u03B7": 0.5
     "\u03C9_c": 0.05
   ? !!python/tuple
   - 10
-  - 2
-  : *id005
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA4D8=
+  : *id006
+  ? !!python/tuple
+  - 10
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA8D8=
+  : *id006
+  ? !!python/tuple
+  - 10
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA+D8=
+  : *id006
+  ? !!python/tuple
+  - 10
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAAAEA=
+  : *id006
   ? !!python/tuple
   - 20
-  - 0
-  : &id006 !!python/object:rabifun.system.Params
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAAAAA=
+  : &id007 !!python/object:rabifun.system.Params
     N: 42
     N_couplings: 20
-    correct_lamb_shift: true
+    correct_lamb_shift: 2
     drive_off_time: null
-    dynamic_detunting: *id002
+    dynamic_detunting: *id003
     flat_energies: false
-    g_0: 0.016666666666666666
+    g_0: 0.008333333333333333
     initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
       args:
-      - *id003
+      - *id004
       - !!python/tuple
         - 0
       - !!binary |
@@ -127,7 +203,7 @@ param_dict:
       - 1
       - !!python/tuple
         - 44
-      - *id004
+      - *id005
       - false
       - !!binary |
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
@@ -144,17 +220,44 @@ param_dict:
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
         AAAAAAAAAAAAAAAAAAAAAAAAAAA=
     laser_detuning: 0
-    laser_off_time: 0.0
+    laser_off_time: 0
     measurement_detuning: 0
     rwa: false
     "\u03A9": 13
-    "\u03B1": 2
+    "\u03B1": !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAAAEA=
     "\u03B4": 0.25
     "\u03B7": 0.5
     "\u03C9_c": 0.05
   ? !!python/tuple
   - 20
-  - 2
-  : *id006
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA4D8=
+  : *id007
+  ? !!python/tuple
+  - 20
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA8D8=
+  : *id007
+  ? !!python/tuple
+  - 20
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAA+D8=
+  : *id007
+  ? !!python/tuple
+  - 20
+  - !!python/object/apply:numpy.core.multiarray.scalar
+    - *id001
+    - !!binary |
+      AAAAAAAAAEA=
+  : *id007
 refers_to: outputs/001_decay_test.pkl
 source: scripts/simulations/001_full_system_rwa_analysis.py
diff --git a/src/rabifun/system.py b/src/rabifun/system.py
index 28b8c64..01baaa0 100644
--- a/src/rabifun/system.py
+++ b/src/rabifun/system.py
@@ -63,7 +63,7 @@ class Params:
     initial_state: np.ndarray | None = None
     """The initial state of the system."""
 
-    correct_lamb_shift: bool = True
+    correct_lamb_shift: float = 1
     """Whether to correct for the Lamb shift by tweaking the detuning."""
 
     def __post_init__(self):
@@ -93,7 +93,7 @@ class Params:
         if not self.flat_energies:
             raise ValueError("Rabi splitting is only defined for flat energies.")
 
-        return np.sqrt((self.Ω * self.g_0) ** 2 + (self.ω_c * self.Ω) ** 2)
+        return np.sqrt((self.Ω * self.g_0 * 2) ** 2 + (self.ω_c * self.Ω) ** 2)
 
 
 class RuntimeParams:
@@ -110,12 +110,14 @@ class RuntimeParams:
         decay_rates = -1j * np.repeat(params.η / 2, params.N + 2)
         Ωs = freqs + decay_rates
 
-        self.drive_frequencies, self.detunings, self.drive_amplitudes = (
+        self.drive_frequencies, self.detunings, self.g = (
             drive_frequencies_and_amplitudes(params)
         )  # linear frequencies!
 
+        self.g *= 2 * np.pi
         self.Ωs = Ωs
-        self.diag_energies = (
+        self.bath_ε = 2 * np.pi * self.detunings - 1j * params.η / 2
+        self.ε = (
             2
             * np.pi
             * np.concatenate(
@@ -128,13 +130,18 @@ class RuntimeParams:
             + decay_rates
         )
 
-        self.detuned_Ωs = freqs - self.diag_energies.real
+        self.detuned_Ωs = freqs - self.ε.real
 
     def __repr__(self):
-        return f"{self.__class__.__name__}(Ωs={self.Ωs}, drive_frequencies={self.drive_frequencies}, drive_amplitudes={self.drive_amplitudes})"
+        return f"{self.__class__.__name__}(Ωs={self.Ωs}, drive_frequencies={self.drive_frequencies}, drive_amplitudes={self.g})"
 
 
-def time_axis(params: Params, lifetimes: float, resolution: float = 1):
+def time_axis(
+    params: Params,
+    lifetimes: float | None = None,
+    recurrences: float | None = None,
+    resolution: float = 1,
+):
     """Generate a time axis for the simulation.
 
     :param params: system parameters
@@ -146,9 +153,15 @@ def time_axis(params: Params, lifetimes: float, resolution: float = 1):
         smaller value yields more points in the time axis.
     """
 
-    return np.arange(
-        0, params.lifetimes(lifetimes), resolution * np.pi / (params.Ω * params.N)
-    )
+    tmax = 0
+    if lifetimes is not None:
+        tmax = params.lifetimes(lifetimes)
+    elif recurrences is not None:
+        tmax = recurrence_time(params) * recurrences
+    else:
+        raise ValueError("Either lifetimes or recurrences must be set.")
+
+    return np.arange(0, tmax, resolution * np.pi / (params.Ω * params.N))
 
 
 def eom_drive(t, x, ds, ωs, rwa, detuned_Ωs):
@@ -160,12 +173,11 @@ 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:coupled_indices] = ds / 2
-        det_matrix[2:coupled_indices, 1] = ds / 2
+        det_matrix[1, 2:coupled_indices] = ds
+        det_matrix[2:coupled_indices, 1] = ds
         driven_x = det_matrix @ x
     else:
         det_matrix = detuned_Ωs[:, None] - detuned_Ωs[None, :]
@@ -176,7 +188,7 @@ def eom_drive(t, x, ds, ωs, rwa, detuned_Ωs):
 
         det_matrix = np.exp(-1j * det_matrix * t)
 
-        driven_x = np.sum(ds * np.sin(2 * np.pi * ωs * t)) * (det_matrix @ x)
+        driven_x = np.sum(2 * ds * np.sin(2 * np.pi * ωs * t)) * (det_matrix @ x)
 
     return driven_x
 
@@ -199,7 +211,7 @@ 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.diag_energies * x
+        differential = runtime_params.ε * x
 
         if params.rwa:
             x[0] = 0
@@ -209,7 +221,7 @@ def make_righthand_side(runtime_params: RuntimeParams, params: Params):
             differential += eom_drive(
                 t,
                 x,
-                runtime_params.drive_amplitudes,
+                runtime_params.g,
                 runtime_params.drive_frequencies,
                 params.rwa,
                 runtime_params.detuned_Ωs,
@@ -266,7 +278,7 @@ def solve(t: np.ndarray, params: Params, **kwargs):
         (np.min(t), np.max(t)),
         initial,
         vectorized=False,
-        # max_step=2 * np.pi / (params.Ω * params.N_couplings),
+        max_step=2 * np.pi / (params.Ω * params.N),
         t_eval=t,
         method="DOP853",
         atol=1e-7,
@@ -350,8 +362,9 @@ def drive_frequencies_and_amplitudes(
     amplitudes /= np.sum(amplitudes)
     amplitudes = params.Ω * params.g_0 * np.sqrt(amplitudes)
 
+    # FIXME: this is twice too big
     if not params.flat_energies and params.correct_lamb_shift:
-        Δs -= np.sum(amplitudes**2 / Δs)
+        Δs -= np.sum(amplitudes**2 / Δs) * params.correct_lamb_shift**2
 
     ωs = ((np.arange(1, params.N_couplings + 1) - params.δ) * params.Ω) - Δs
     return ωs, Δs, amplitudes