diff --git a/scripts/experiments/005_mean_displacement_semi_analytic.py b/scripts/experiments/005_mean_displacement_semi_analytic.py
index 4e3434a..f2e17e7 100644
--- a/scripts/experiments/005_mean_displacement_semi_analytic.py
+++ b/scripts/experiments/005_mean_displacement_semi_analytic.py
@@ -38,16 +38,22 @@ def characteristic_poly(g: np.ndarray, ε: np.ndarray, η_A: float = 0):
     return poly
 
 
-def hamiltonian(g: np.ndarray, ε: np.ndarray, η_A: float = 0):
-    H = np.diag([-1j * η_A, *ε])
+def hamiltonian(g: np.ndarray, ε: np.ndarray, ε_A: float = 0, η_A: float = 0):
+    H = np.diag([ε_A - 1j * η_A, *ε])
     H[0, 1:] = g
     H[1:, 0] = np.conj(g)
     return H
 
 
-def a_site_population(t, ω, coeff):
+def a_site_population(t, ω, coeff, lower_cutoff=0):
     return (
-        np.abs(np.sum(coeff[None, :] * np.exp(-1j * ω[None, :] * t[:, None]), axis=1))
+        np.abs(
+            np.sum(
+                coeff[None, lower_cutoff:]
+                * np.exp(-1j * ω[None, lower_cutoff:] * t[:, None]),
+                axis=1,
+            )
+        )
         ** 2
     )
 
@@ -64,7 +70,7 @@ def make_params(ω_c=0.1 / 2, N=10, gbar=1 / 3):
         ω_c=ω_c,
         g_0=ω_c * gbar,
         laser_detuning=0,
-        N=2 * N + 2,
+        N=N,
         N_couplings=N,
         measurement_detuning=0,
         α=0,
@@ -76,29 +82,53 @@ def make_params(ω_c=0.1 / 2, N=10, gbar=1 / 3):
 
 
 def test():
-    params = make_params(N=10, gbar=1 / 3)
+    params = make_params(N=20, gbar=1 / 4)
     params.flat_energies = False
-    params.α = 0.0
-    params.correct_lamb_shift = True
+    params.α = 0
 
+    params.correct_lamb_shift = 1
     runtime = RuntimeParams(params)
-    t = time_axis(params, recurrences=1.5)
-    g = runtime.g / 2
+    t = time_axis(params, recurrences=1.1)
+    g = runtime.g
     ε = runtime.bath_ε
-    H = hamiltonian(g, ε.real, 0 * params.η / 2)
 
-    ω = np.linalg.eigvals(H)
-    M = np.linalg.eig(H).eigenvectors
+    H = hamiltonian(g, ε.real, ε_A=runtime.a_shift, η_A=0 * params.η / 2)
+
+    print(runtime.a_shift - ε[-1])
+    eig = np.linalg.eig(H)
+
+    ω = eig.eigenvalues
+
+    idx = np.argsort(ω.real)
+    ω = ω[idx]
+
+    M = np.linalg.eig(H).eigenvectors[:, idx]
+
     Minv = np.linalg.inv(M)
-    v0 = Minv[:, 0]
 
-    coeff = M[0, :] * v0.T
+    coeff = M[0, :] * Minv[:, 0]
+
+    f = make_figure()
+
+    U_A = np.abs(1 / (1 + np.sum((g / (ω[0] - ε.real)) ** 2))) ** 2
+    U_A_coeff = np.max(np.abs(coeff**2))
+    print(np.argmax(np.abs(coeff**2)))
     print(coeff)
+    ax_t, ax_e = f.subplots(1, 2)
+    ax_t.plot(t, a_site_population(t, ω, coeff, 0))
+    # ax_t.plot(t, U_A_coeff + .1 * np.sin(ω[0].real * t))
+    ax_t.set_ylim(0, 1.01)
+    ax_t.set_xlabel("Time [1/Ω]")
+    ax_t.set_ylabel(r"$ρ_A$")
 
-    # coeff /= np.abs(np.sum(coeff))
-    # coeff = coefficients(ω, g, ε)
+    # print((np.abs(1 / (1 + np.sum((g / ε) ** 2)))))
+    ax_t.axhline(U_A, color="green", linestyle="-.")
+    print(U_A)
+    ax_t.axhline(U_A_coeff, color="red", linestyle="--")
 
-    plt.cla()
-    plt.plot(t, a_site_population(t, ω, coeff))
-    plt.ylim(0, 1)
-    return ω
+    for ω_ in ω:
+        ax_e.axvline(ω_.real, color="blue", alpha=0.5)
+    ax_e.axvline(min(ω.real), color="black", linestyle="--")
+    ax_e.axvline(max(ω.real), color="black", linestyle="--")
+    ax_e.axvline(max(ε.real), color="green", linestyle="--")
+    ax_e.set_xlabel("Energy")
diff --git a/scripts/simulations/001_full_system_rwa_analysis.py b/scripts/simulations/001_full_system_rwa_analysis.py
index 08e5c93..67087a4 100644
--- a/scripts/simulations/001_full_system_rwa_analysis.py
+++ b/scripts/simulations/001_full_system_rwa_analysis.py
@@ -43,9 +43,9 @@ def decay_rwa_analysis():
     and with much fewer modes.
     """
 
-    ω_c = 0.05
-    Ns = [5, 10, 20]
-    gbar = 1 / 3 / 2
+    ω_c = 0.1 / 2
+    Ns = [1, 5]
+    gbar = 1 / 4
 
     fig = make_figure("decay_test", figsize=(15, len(Ns) * 3))
     ax_ns = fig.subplots(len(Ns), 2)
@@ -56,7 +56,7 @@ def decay_rwa_analysis():
     param_dict = {}
 
     for i, N in enumerate(Ns):
-        params = make_params(ω_c=ω_c, N=N, gbar=gbar)
+        params = make_params(ω_c=ω_c, N=N, gbar=gbar, compensate=1)
         params.laser_off_time = 0
         params.initial_state = make_zero_intial_state(params)
         params.initial_state[1] = 1
@@ -67,7 +67,7 @@ def decay_rwa_analysis():
 
         t = time_axis(params, recurrences=1.1, resolution=0.1)
 
-        for α in np.linspace(0, 2, 5):
+        for α in np.linspace(0, 2, 3):
             params.α = α
             sol_nonrwa, sol_rwa = solve_nonrwa_rwa(t, params)
 
diff --git a/scripts/simulations/figs/001_decay_test.pdf b/scripts/simulations/figs/001_decay_test.pdf
index e2679d5..d7eb382 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 8f4fca3..56a2ce8 100644
--- a/scripts/simulations/figs/001_decay_test.pdf.meta.yaml
+++ b/scripts/simulations/figs/001_decay_test.pdf.meta.yaml
@@ -1,14 +1,13 @@
-change_id: wykwpsvmtrzvwtunotxnzokztpyrnsrx
-commit_id: 94f23ecee30914b5163f14fd5af637a8777b42cc
+change_id: zqvmrwsoxvtmtutmpnkonxuntmvxpnou
+commit_id: fa7a31938c77430505d38b5a0415ff0666a1fc6c
 description: ''
 extra_meta:
   Ns:
+  - 1
   - 5
-  - 10
-  - 20
   params:
     ? !!python/tuple
-    - 5
+    - 1
     - !!python/object/apply:numpy.core.multiarray.scalar
       - &id001 !!python/object/apply:numpy.dtype
         args:
@@ -27,15 +26,15 @@ extra_meta:
       - !!binary |
         AAAAAAAAAAA=
     : &id002 !!python/object:rabifun.system.Params
-      N: 12
-      N_couplings: 5
-      correct_lamb_shift: 2
+      N: 4
+      N_couplings: 1
+      correct_lamb_shift: 1
       drive_off_time: null
       dynamic_detunting: &id003 !!python/tuple
       - 0
       - 0
       flat_energies: false
-      g_0: 0.008333333333333333
+      g_0: 0.0125
       initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
         args:
         - &id004 !!python/name:numpy.ndarray ''
@@ -46,7 +45,7 @@ extra_meta:
         state: !!python/tuple
         - 1
         - !!python/tuple
-          - 14
+          - 10
         - &id005 !!python/object/apply:numpy.dtype
           args:
           - c16
@@ -65,8 +64,7 @@ extra_meta:
         - !!binary |
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==
       laser_detuning: 0
       laser_off_time: 0
       measurement_detuning: 0
@@ -80,47 +78,33 @@ extra_meta:
       "\u03B7": 0.5
       "\u03C9_c": 0.05
     ? !!python/tuple
-    - 5
-    - !!python/object/apply:numpy.core.multiarray.scalar
-      - *id001
-      - !!binary |
-        AAAAAAAA4D8=
-    : *id002
-    ? !!python/tuple
-    - 5
+    - 1
     - !!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
+    - 1
     - !!python/object/apply:numpy.core.multiarray.scalar
       - *id001
       - !!binary |
         AAAAAAAAAEA=
     : *id002
     ? !!python/tuple
-    - 10
+    - 5
     - !!python/object/apply:numpy.core.multiarray.scalar
       - *id001
       - !!binary |
         AAAAAAAAAAA=
     : &id006 !!python/object:rabifun.system.Params
-      N: 22
-      N_couplings: 10
-      correct_lamb_shift: 2
+      N: 12
+      N_couplings: 5
+      correct_lamb_shift: 1
       drive_off_time: null
       dynamic_detunting: *id003
       flat_energies: false
-      g_0: 0.008333333333333333
+      g_0: 0.0125
       initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
         args:
         - *id004
@@ -131,7 +115,7 @@ extra_meta:
         state: !!python/tuple
         - 1
         - !!python/tuple
-          - 24
+          - 26
         - *id005
         - false
         - !!binary |
@@ -141,7 +125,8 @@ extra_meta:
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
           AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+          AAAAAAAAAAAAAAAAAAAAAAA=
       laser_detuning: 0
       laser_off_time: 0
       measurement_detuning: 0
@@ -155,114 +140,19 @@ extra_meta:
       "\u03B7": 0.5
       "\u03C9_c": 0.05
     ? !!python/tuple
-    - 10
-    - !!python/object/apply:numpy.core.multiarray.scalar
-      - *id001
-      - !!binary |
-        AAAAAAAA4D8=
-    : *id006
-    ? !!python/tuple
-    - 10
+    - 5
     - !!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
+    - 5
     - !!python/object/apply:numpy.core.multiarray.scalar
       - *id001
       - !!binary |
         AAAAAAAAAEA=
     : *id006
-    ? !!python/tuple
-    - 20
-    - !!python/object/apply:numpy.core.multiarray.scalar
-      - *id001
-      - !!binary |
-        AAAAAAAAAAA=
-    : &id007 !!python/object:rabifun.system.Params
-      N: 42
-      N_couplings: 20
-      correct_lamb_shift: 2
-      drive_off_time: null
-      dynamic_detunting: *id003
-      flat_energies: false
-      g_0: 0.008333333333333333
-      initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
-        args:
-        - *id004
-        - !!python/tuple
-          - 0
-        - !!binary |
-          Yg==
-        state: !!python/tuple
-        - 1
-        - !!python/tuple
-          - 44
-        - *id005
-        - false
-        - !!binary |
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-          AAAAAAAAAAAAAAAAAAAAAAAAAAA=
-      laser_detuning: 0
-      laser_off_time: 0
-      measurement_detuning: 0
-      rwa: false
-      "\u03A9": 13
-      "\u03B1": !!python/object/apply:numpy.core.multiarray.scalar
-      - *id001
-      - !!binary |
-        AAAAAAAAAEA=
-      "\u03B4": 0.25
-      "\u03B7": 0.5
-      "\u03C9_c": 0.05
-    ? !!python/tuple
-    - 20
-    - !!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 a417183..431632d 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 d84ee19..537697d 100644
--- a/scripts/simulations/outputs/001_decay_test.pkl.meta.yaml
+++ b/scripts/simulations/outputs/001_decay_test.pkl.meta.yaml
@@ -1,10 +1,10 @@
-change_id: wykwpsvmtrzvwtunotxnzokztpyrnsrx
-commit_id: e4399f86a81e8dc0fbbe24226b948db60bd83807
+change_id: zqvmrwsoxvtmtutmpnkonxuntmvxpnou
+commit_id: 7372777c2aeb7f97f97d5e754b4bfde3a6244232
 description: ''
 function: decay_rwa_analysis
 param_dict:
   ? !!python/tuple
-  - 5
+  - 1
   - !!python/object/apply:numpy.core.multiarray.scalar
     - &id001 !!python/object/apply:numpy.dtype
       args:
@@ -23,15 +23,15 @@ param_dict:
     - !!binary |
       AAAAAAAAAAA=
   : &id002 !!python/object:rabifun.system.Params
-    N: 12
-    N_couplings: 5
-    correct_lamb_shift: 2
+    N: 4
+    N_couplings: 1
+    correct_lamb_shift: 1
     drive_off_time: null
     dynamic_detunting: &id003 !!python/tuple
     - 0
     - 0
     flat_energies: false
-    g_0: 0.008333333333333333
+    g_0: 0.0125
     initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
       args:
       - &id004 !!python/name:numpy.ndarray ''
@@ -42,7 +42,7 @@ param_dict:
       state: !!python/tuple
       - 1
       - !!python/tuple
-        - 14
+        - 10
       - &id005 !!python/object/apply:numpy.dtype
         args:
         - c16
@@ -61,8 +61,7 @@ param_dict:
       - !!binary |
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA==
     laser_detuning: 0
     laser_off_time: 0
     measurement_detuning: 0
@@ -76,47 +75,33 @@ param_dict:
     "\u03B7": 0.5
     "\u03C9_c": 0.05
   ? !!python/tuple
-  - 5
-  - !!python/object/apply:numpy.core.multiarray.scalar
-    - *id001
-    - !!binary |
-      AAAAAAAA4D8=
-  : *id002
-  ? !!python/tuple
-  - 5
+  - 1
   - !!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
+  - 1
   - !!python/object/apply:numpy.core.multiarray.scalar
     - *id001
     - !!binary |
       AAAAAAAAAEA=
   : *id002
   ? !!python/tuple
-  - 10
+  - 5
   - !!python/object/apply:numpy.core.multiarray.scalar
     - *id001
     - !!binary |
       AAAAAAAAAAA=
   : &id006 !!python/object:rabifun.system.Params
-    N: 22
-    N_couplings: 10
-    correct_lamb_shift: 2
+    N: 12
+    N_couplings: 5
+    correct_lamb_shift: 1
     drive_off_time: null
     dynamic_detunting: *id003
     flat_energies: false
-    g_0: 0.008333333333333333
+    g_0: 0.0125
     initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
       args:
       - *id004
@@ -127,7 +112,7 @@ param_dict:
       state: !!python/tuple
       - 1
       - !!python/tuple
-        - 24
+        - 26
       - *id005
       - false
       - !!binary |
@@ -137,7 +122,8 @@ param_dict:
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
         AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
+        AAAAAAAAAAAAAAAAAAAAAAA=
     laser_detuning: 0
     laser_off_time: 0
     measurement_detuning: 0
@@ -151,113 +137,18 @@ param_dict:
     "\u03B7": 0.5
     "\u03C9_c": 0.05
   ? !!python/tuple
-  - 10
-  - !!python/object/apply:numpy.core.multiarray.scalar
-    - *id001
-    - !!binary |
-      AAAAAAAA4D8=
-  : *id006
-  ? !!python/tuple
-  - 10
+  - 5
   - !!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
+  - 5
   - !!python/object/apply:numpy.core.multiarray.scalar
     - *id001
     - !!binary |
       AAAAAAAAAEA=
   : *id006
-  ? !!python/tuple
-  - 20
-  - !!python/object/apply:numpy.core.multiarray.scalar
-    - *id001
-    - !!binary |
-      AAAAAAAAAAA=
-  : &id007 !!python/object:rabifun.system.Params
-    N: 42
-    N_couplings: 20
-    correct_lamb_shift: 2
-    drive_off_time: null
-    dynamic_detunting: *id003
-    flat_energies: false
-    g_0: 0.008333333333333333
-    initial_state: !!python/object/apply:numpy.core.multiarray._reconstruct
-      args:
-      - *id004
-      - !!python/tuple
-        - 0
-      - !!binary |
-        Yg==
-      state: !!python/tuple
-      - 1
-      - !!python/tuple
-        - 44
-      - *id005
-      - false
-      - !!binary |
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAPA/AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA
-        AAAAAAAAAAAAAAAAAAAAAAAAAAA=
-    laser_detuning: 0
-    laser_off_time: 0
-    measurement_detuning: 0
-    rwa: false
-    "\u03A9": 13
-    "\u03B1": !!python/object/apply:numpy.core.multiarray.scalar
-    - *id001
-    - !!binary |
-      AAAAAAAAAEA=
-    "\u03B4": 0.25
-    "\u03B7": 0.5
-    "\u03C9_c": 0.05
-  ? !!python/tuple
-  - 20
-  - !!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 01baaa0..b6f711f 100644
--- a/src/rabifun/system.py
+++ b/src/rabifun/system.py
@@ -1,3 +1,4 @@
+from pdb import run
 import numpy as np
 import matplotlib.pyplot as plt
 from dataclasses import dataclass
@@ -70,8 +71,8 @@ class Params:
         if self.N_couplings > self.N:
             raise ValueError("N_couplings must be less than or equal to N.")
 
-        if self.initial_state and len(self.initial_state) != self.N + 2:
-            raise ValueError("Initial state must have length N + 2.")
+        if self.initial_state and len(self.initial_state) != 2 * self.N + 2:
+            raise ValueError("Initial state must have length 2N + 2.")
 
     def periods(self, n: float):
         """
@@ -100,37 +101,45 @@ class RuntimeParams:
     """Secondary Parameters that are required to run the simulation."""
 
     def __init__(self, params: Params):
+        bath = np.arange(1, params.N + 1)
         freqs = (
             2
             * np.pi
             * params.Ω
-            * np.concatenate([[-1 * params.δ, params.δ], np.arange(1, params.N + 1)])
+            * np.concatenate([[-1 * params.δ, params.δ], bath, -bath])
         )
 
-        decay_rates = -1j * np.repeat(params.η / 2, params.N + 2)
+        decay_rates = -1j * np.repeat(params.η / 2, 2 * params.N + 2)
         Ωs = freqs + decay_rates
 
-        self.drive_frequencies, self.detunings, self.g = (
+        self.drive_frequencies, self.detunings, self.g, a_shift = (
             drive_frequencies_and_amplitudes(params)
         )  # linear frequencies!
 
         self.g *= 2 * np.pi
         self.Ωs = Ωs
-        self.bath_ε = 2 * np.pi * self.detunings - 1j * params.η / 2
         self.ε = (
             2
             * np.pi
             * np.concatenate(
                 [
-                    [0, 0],
+                    [-a_shift, a_shift],
                     self.detunings,
                     np.zeros(params.N - params.N_couplings),
+                    -self.detunings,
+                    np.zeros(params.N - params.N_couplings),
                 ]
             )
             + decay_rates
         )
 
+        self.bath_ε = 2 * np.pi * self.detunings - 1j * params.η / 2
+        self.a_shift = 2 * np.pi * a_shift
         self.detuned_Ωs = freqs - self.ε.real
+        self.RWA_H = np.zeros((2 * params.N + 2, 2 * params.N + 2), np.complex128)
+        self.RWA_H[1, 2 : 2 + params.N_couplings] = self.g
+        self.RWA_H[2 : 2 + params.N_couplings, 1] = np.conj(self.g)
+        self.detuning_matrix = self.detuned_Ωs[:, None] - self.detuned_Ωs[None, :]
 
     def __repr__(self):
         return f"{self.__class__.__name__}(Ωs={self.Ωs}, drive_frequencies={self.drive_frequencies}, drive_amplitudes={self.g})"
@@ -164,7 +173,7 @@ def time_axis(
     return np.arange(0, tmax, resolution * np.pi / (params.Ω * params.N))
 
 
-def eom_drive(t, x, ds, ωs, rwa, detuned_Ωs):
+def eom_drive(t, x, ds, ωs, det_matrix):
     """The electrooptical modulation drive.
 
     :param t: time
@@ -173,22 +182,14 @@ def eom_drive(t, x, ds, ωs, rwa, detuned_Ωs):
     :param ωs: linear drive frequencies
     """
 
-    if rwa:
-        coupled_indices = 2 + len(ds)
-        det_matrix = np.zeros((len(x), len(x)))
-        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, :]
-        # test = abs(det_matrix.copy())
-        # test[test < 1e-10] = np.inf
-        # print(np.min(test))
-        # print(np.argmin(test, keepdims=True))
+    # 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)
+    det_matrix = np.exp(-1j * det_matrix * t)
 
-        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)) * (det_matrix @ x)
 
     return driven_x
 
@@ -218,14 +219,16 @@ def make_righthand_side(runtime_params: RuntimeParams, params: Params):
             x[2 + params.N_couplings :] = 0
 
         if (params.drive_off_time is None) or (t < params.drive_off_time):
-            differential += eom_drive(
-                t,
-                x,
-                runtime_params.g,
-                runtime_params.drive_frequencies,
-                params.rwa,
-                runtime_params.detuned_Ωs,
-            )
+            if params.rwa:
+                differential += runtime_params.RWA_H @ x
+            else:
+                differential += eom_drive(
+                    t,
+                    x,
+                    runtime_params.g,
+                    runtime_params.drive_frequencies,
+                    runtime_params.detuning_matrix,
+                )
 
         if (params.laser_off_time is None) or (t < params.laser_off_time):
             freqs = laser_frequency(params, t) - runtime_params.detuned_Ωs.real
@@ -253,7 +256,7 @@ def make_righthand_side(runtime_params: RuntimeParams, params: Params):
 
 def make_zero_intial_state(params: Params) -> np.ndarray:
     """Make initial state with all zeros."""
-    return np.zeros(params.N + 2, np.complex128)
+    return np.zeros(2 * params.N + 2, np.complex128)
 
 
 def solve(t: np.ndarray, params: Params, **kwargs):
@@ -278,7 +281,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),
+        max_step=np.pi / (params.Ω * params.N),
         t_eval=t,
         method="DOP853",
         atol=1e-7,
@@ -326,13 +329,14 @@ def output_signal(t: np.ndarray, amplitudes: np.ndarray, params: Params):
 def bath_energies(N_couplings: int, ω_c: float) -> np.ndarray:
     """Return the energies (drive detunings) of the bath modes."""
 
-    return np.arange(1, N_couplings + 1) * ω_c / N_couplings
+    return (np.arange(1, N_couplings + 1) - 1 / 2) * ω_c / N_couplings
 
 
 def ohmic_spectral_density(ω: np.ndarray, α: float) -> np.ndarray:
     """The unnormalized spectral density of an Ohmic bath."""
 
-    return ω**α
+    ω = np.concatenate([[0], ω])
+    return np.diff(ω ** (α + 1))
 
 
 def lamb_shift(amplitudes, Δs):
@@ -341,7 +345,7 @@ def lamb_shift(amplitudes, Δs):
 
 def drive_frequencies_and_amplitudes(
     params: Params,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]:
     """
     Return the linear frequencies and amplitudes of the drives based
     on the ``params``.
@@ -362,12 +366,12 @@ def drive_frequencies_and_amplitudes(
     amplitudes /= np.sum(amplitudes)
     amplitudes = params.Ω * params.g_0 * np.sqrt(amplitudes)
 
-    # FIXME: this is twice too big
+    a_shift = 0
     if not params.flat_energies and params.correct_lamb_shift:
-        Δs -= np.sum(amplitudes**2 / Δs) * params.correct_lamb_shift**2
+        a_shift = 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
+    ωs = ((np.arange(1, params.N_couplings + 1) - params.δ) * params.Ω) - (Δs - a_shift)
+    return ωs, Δs, amplitudes, a_shift
 
 
 def mode_name(mode: int):
@@ -392,7 +396,7 @@ def coupled_bath_mode_indices(params: Params):
 
 def uncoupled_bath_mode_indices(params: Params):
     """Return the indices of the bath modes that not coupled to the A site."""
-    return np.arange(2 + params.N_couplings, 2 + params.N)
+    return np.arange(2 + params.N_couplings, 2 + 2 * params.N)
 
 
 def uncoupled_mode_indices(params: Params):
@@ -411,7 +415,7 @@ def coupled_mode_indices(params: Params):
 
 def dimension(params: Params):
     """Return the dimension of the system."""
-    return params.N + 2
+    return 2 * params.N + 2
 
 
 def recurrence_time(params: Params):