diff --git a/scripts/experiments/004_full_nm_walk_tests.py b/scripts/experiments/004_full_nm_walk_tests.py
index 60efbd4..6dfe48b 100644
--- a/scripts/experiments/004_full_nm_walk_tests.py
+++ b/scripts/experiments/004_full_nm_walk_tests.py
@@ -7,24 +7,29 @@ from rabifun.utilities import *
 
 def test_collective_mode_rabi():
     """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."""
+    ω_c = 0.1 / 2
     params = Params(
-        η=0.01,
-        Ω=1,
+        η=0.5,
+        Ω=13,
         δ=1 / 4,
-        g_0=0.01,
+        g_0=ω_c / 5,
         laser_detuning=0,
-        ω_c=0.00,
+        ω_c=ω_c,
         N=2,
         N_couplings=2,
         measurement_detuning=0,
+        α=0.1,
         rwa=False,
         flat_energies=True,
     )
 
-    # params.laser_off_time = params.lifetimes(5)
-    t = time_axis(params, 15, 0.01)
-    solution = solve(t, params)
-    print(RuntimeParams(params))
+    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)
+    # solution = solve(t, params)
+    # print(RuntimeParams(params))
     # signal = output_signal(t, solution.y, params)
 
     # f, (_, ax) = plot_simulation_result(
@@ -36,11 +41,9 @@ def test_collective_mode_rabi():
     fig = make_figure()
     ax = fig.subplots()
 
-    rotsol = in_rotating_frame(t, solution.y, params)
+    sol_nonrwa, sol_nonrwa, *_ = plot_rwa_vs_real_amplitudes(ax, t, params)
+    ax.axvline(params.laser_off_time, color="black", linestyle="--")
 
-    ax.plot(t, np.abs(rotsol)[1])
-    ax.plot(t, np.abs(rotsol)[2])
-    ax.plot(t, np.abs(rotsol)[3])
-    # ax.plot(t, np.abs(rotsol)[4])
-    # ax.plot(t, np.abs(in_rotating_frame(t, solution.y, params))[2])
-    # ax.plot(t, np.abs(in_rotating_frame(t, solution.y, params))[3])
+    print(sol_nonrwa.y[2][-1] / sol_nonrwa.y[3][-1])
+    runtime = RuntimeParams(params)
+    print(runtime.drive_amplitudes[0] / runtime.drive_amplitudes[1])
diff --git a/src/rabifun/plots.py b/src/rabifun/plots.py
index 2346190..b1adcfa 100644
--- a/src/rabifun/plots.py
+++ b/src/rabifun/plots.py
@@ -1,5 +1,5 @@
 from plot_utils import *
-from .system import Params, RuntimeParams
+from .system import Params, RuntimeParams, solve, coupled_mode_indices, mode_name
 from .analysis import fourier_transform
 import matplotlib.pyplot as plt
 import numpy as np
@@ -97,3 +97,89 @@ def plot_rabi_sidebands(ax, params: Params):
         )
 
     ax.legend()
+
+
+def clone_color_or_default(lines: dict, mode: int):
+    """
+    Get the color of a mode or a default color if it doesn't exist.
+
+    :param lines: A dictionary of lines indexed by mode index.
+    :param mode: The mode to get the color of.
+    """
+
+    line = lines.get(mode, None)
+    if line is None:
+        return f"C{mode}"
+
+    return line.get_color()
+
+
+def plot_rotating_modes(ax, solution, params, plot_uncoupled=False, clone_colors=None):
+    """Plot the amplitude of the modes in the rotating frame.
+
+    :param ax: The axis to plot on.
+    :param solution: The solution to plot.
+    :param params: The system parameters.
+    :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.
+    """
+
+    lines = dict()
+    if clone_colors is None:
+        clone_colors = dict()
+
+    for mode in coupled_mode_indices(params):
+        lines[mode] = ax.plot(
+            solution.t,
+            np.abs(solution.y[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:
+        for mode in uncoupled_mode_indices(params):
+            lines[mode] = ax.plot(
+                solution.t,
+                np.abs(solution.y[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.set_xlabel("Time (1/Ω)")
+    ax.set_ylabel("Amplitude")
+    ax.set_title("Mode Amplitudes in the Rotating Frame")
+
+    return lines
+
+
+def plot_rwa_vs_real_amplitudes(ax, t, 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 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
+              the dictionaries of the lines plotted indexed by mode
+              index.
+    """
+    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
diff --git a/src/rabifun/system.py b/src/rabifun/system.py
index 246dabb..c338095 100644
--- a/src/rabifun/system.py
+++ b/src/rabifun/system.py
@@ -25,8 +25,7 @@ class Params:
     """Mode splitting in units of :any:`Ω`."""
 
     η: float = 0.1
-    """Decay rate :math:`\eta/2` of the system in frequency units (no
-    :math:`2 \pi`)."""
+    """Decay rate :math:`\eta/2` of the system in angular frequency units."""
 
     g_0: float = 0.01
     """Drive amplitude in units of :any:`Ω`."""
@@ -61,10 +60,15 @@ class Params:
     flat_energies: bool = False
     """Whether to use a flat distribution of bath energies."""
 
+    initial_state: np.ndarray | None = None
+
     def __post_init__(self):
         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.")
+
     def periods(self, n: float):
         """
         Returns the number of periods of the system that correspond to
@@ -91,14 +95,30 @@ class Params:
 class RuntimeParams:
     """Secondary Parameters that are required to run the simulation."""
 
-    __slots__ = ["Ωs", "drive_frequencies", "drive_amplitudes"]
-
     def __init__(self, params: Params):
-        Ωs = 2 * np.pi * params.Ω * np.concatenate(
-            [[-1 * params.δ, params.δ], np.arange(1, params.N + 1)]
-        ) - 1j * np.repeat(params.η / 2, params.N + 2)
+        freqs = (
+            2
+            * np.pi
+            * 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.Ωs = Ωs
+        self.diag_energies = (
+            2
+            * np.pi
+            * np.concatenate(
+                [
+                    [0, 0],
+                    drive_detunings(params),
+                    np.zeros(params.N - params.N_couplings),
+                ]
+            )
+            + decay_rates
+        )
+        self.detuned_Ωs = freqs - self.diag_energies.real
 
         self.drive_frequencies, self.drive_amplitudes = (
             drive_frequencies_and_amplitudes(params)
@@ -125,7 +145,7 @@ def time_axis(params: Params, lifetimes: float, resolution: float = 1):
     )
 
 
-def eom_drive(t, x, ds, ωs, rwa):
+def eom_drive(t, x, ds, ωs, rwa, detuned_Ωs):
     """The electrooptical modulation drive.
 
     :param t: time
@@ -134,15 +154,15 @@ def eom_drive(t, x, ds, ωs, rwa):
     :param ωs: linear drive frequencies
     """
 
-    stacked = np.repeat([x], len(x), 0)
-
     if rwa:
-        stacked[1, 2 : 2 + len(ωs)] *= ds * np.exp(-1j * 2 * np.pi * ωs * t)
-        stacked[2 : 2 + len(ωs), 1] *= ds * np.exp(1j * 2 * np.pi * ωs * t)
-        driven_x = np.sum(stacked, axis=1)
+        det_matrix = np.zeros((len(x), len(x)))
+        det_matrix[1, 2:] = ds / 2
+        det_matrix[2:, 1] = ds / 2
+        driven_x = det_matrix @ x
     else:
-        stacked = np.sum(stacked, axis=1)
-        driven_x = np.sum(ds * np.sin(2 * np.pi * ωs * t)) * stacked
+        freqs = np.exp(-1j * detuned_Ωs * t)
+        det_matrix = np.outer(np.conj(freqs), freqs)
+        driven_x = np.sum(ds * np.sin(2 * np.pi * ωs * t)) * (det_matrix @ x)
 
     return driven_x
 
@@ -165,7 +185,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.Ωs * x
+        differential = runtime_params.diag_energies * x
 
         if params.rwa:
             x[0] = 0
@@ -178,24 +198,41 @@ def make_righthand_side(runtime_params: RuntimeParams, params: Params):
                 runtime_params.drive_amplitudes,
                 runtime_params.drive_frequencies,
                 params.rwa,
+                runtime_params.detuned_Ωs,
             )
 
         if (params.laser_off_time is None) or (t < params.laser_off_time):
-            laser = np.exp(-1j * laser_frequency(params, t) * t)
-            differential[0:2] += laser / np.sqrt(2)
-            differential[2:] += laser
+            freqs = laser_frequency(params, t) - runtime_params.detuned_Ωs.real
 
-        if params.rwa:
-            differential[0] = 0
-            differential[2 + params.N_couplings :] = 0
+            laser = np.exp(
+                -1j * (laser_frequency(params, t) - runtime_params.detuned_Ωs.real) * t
+            )
+
+            if params.rwa:
+                index = np.argmin(abs(freqs))
+                laser[0:index] = 0
+                laser[index + 1 :] = 0
+
+            differential[0:2] += laser[:2] / np.sqrt(2)
+            differential[2:] += laser[2:]
+
+        # if params.rwa:
+        #     differential[0] = 0
+        #     differential[2 + params.N_couplings :] = 0
 
         return -1j * differential
 
     return rhs
 
 
-def solve(t: np.ndarray, 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)
+
+
+def solve(t: np.ndarray, params: Params, **kwargs):
     """Integrate the equation of motion.
+    The keyword arguments are passed to :any:`scipy.integrate.solve_ivp`.
 
     :param t: time array
     :param params: system parameters
@@ -204,15 +241,23 @@ def solve(t: np.ndarray, params: Params):
     runtime = RuntimeParams(params)
     rhs = make_righthand_side(runtime, params)
 
-    initial = np.zeros(params.N + 2, np.complex128)
+    initial = (
+        make_zero_intial_state(params)
+        if params.initial_state is None
+        else params.initial_state
+    )
 
     return scipy.integrate.solve_ivp(
         rhs,
         (np.min(t), np.max(t)),
         initial,
         vectorized=False,
-        max_step=0.01 * np.pi / (params.Ω * params.N),
+        # max_step=2 * np.pi / (params.Ω * params.N_couplings),
         t_eval=t,
+        method="DOP853",
+        atol=1e-7,
+        rtol=1e-4,
+        **kwargs,
     )
 
 
@@ -248,24 +293,16 @@ def output_signal(t: np.ndarray, amplitudes: np.ndarray, params: Params):
     ).imag
 
 
-def bath_energies(params: Params) -> np.ndarray:
+def bath_energies(N_couplings: int, ω_c: float) -> np.ndarray:
     """Return the energies (drive detunings) of the bath modes."""
 
-    return np.arange(1, params.N_couplings + 1) * params.ω_c / params.N_couplings
+    return np.arange(1, N_couplings + 1) * ω_c / N_couplings
 
 
-def ohmic_spectral_density(
-    ω: np.ndarray, g_0: float, α: float, ω_c: float
-) -> np.ndarray:
-    """The spectral density of an Ohmic bath."""
+def ohmic_spectral_density(ω: np.ndarray, α: float) -> np.ndarray:
+    """The unnormalized spectral density of an Ohmic bath."""
 
-    mask = (ω > ω_c) | (ω < 0)
-
-    ret = np.empty_like(ω)
-    ret[mask] = 0
-    ret[~mask] = g_0**2 * (α + 1) / (ω_c ** (α + 1)) * (ω**α)
-
-    return ret
+    return ω**α
 
 
 def drive_detunings(params: Params) -> np.ndarray:
@@ -274,7 +311,7 @@ def drive_detunings(params: Params) -> np.ndarray:
     if params.flat_energies:
         Δs = np.repeat(params.ω_c, params.N_couplings)
     else:
-        Δs = bath_energies(params)
+        Δs = bath_energies(params.N_couplings, params.ω_c)
 
     return Δs * params.Ω
 
@@ -286,20 +323,58 @@ def drive_frequencies_and_amplitudes(params: Params) -> tuple[np.ndarray, np.nda
     """
 
     Δs = drive_detunings(params)
-    if params.flat_energies:
-        amplitudes = np.ones_like(Δs)
-    else:
-        amplitudes = (
-            ohmic_spectral_density(
-                Δs, params.g_0 * params.Ω, params.α, params.ω_c * params.Ω
-            )
-            * params.ω_c
-            * params.Ω
-            / Params.N_couplings
-        )
+
+    amplitudes = ohmic_spectral_density(
+        bath_energies(params.N_couplings, 1),
+        params.α,
+    )
 
     amplitudes /= np.sum(amplitudes)
     amplitudes = 2 * np.pi * params.Ω * params.g_0 * np.sqrt(amplitudes)
 
     ωs = ((np.arange(1, params.N_couplings + 1) - params.δ) * params.Ω) - Δs
     return ωs, amplitudes
+
+
+def mode_name(mode: int):
+    """Return the name of the mode."""
+    if mode == 0:
+        return "anti-A"
+    if mode == 1:
+        return "A"
+
+    return f"B{mode - 1}"
+
+
+def a_site_indices(params: Params):
+    """Return the indices of the A sites."""
+    return [0, 1]
+
+
+def coupled_bath_mode_indices(params: Params):
+    """Return the indices of the bath modes that are coupled to the A site."""
+    return np.arange(2, 2 + params.N_couplings)
+
+
+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)
+
+
+def uncoupled_mode_indices(params: Params):
+    """Return the indices of the modes that not coupled to the A site."""
+    return np.concatenate(
+        [[a_site_indices(params)[0]], uncoupled_bath_mode_indices(params)]
+    )
+
+
+def coupled_mode_indices(params: Params):
+    """Return the indices of the modes that are coupled to the A site."""
+    return np.concatenate(
+        [[a_site_indices(params)[1]], coupled_bath_mode_indices(params)]
+    )
+
+
+def dimension(params: Params):
+    """Return the dimension of the system."""
+    return params.N + 2