mirror of
https://github.com/vale981/fibre_walk_project_code
synced 2025-03-04 09:21:38 -05:00
simulate dynamics in rotating frame
This commit is contained in:
parent
bf6d138826
commit
8e66d5b4ff
3 changed files with 230 additions and 66 deletions
|
|
@ -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])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue