def get_Hamiltonian(options_dict):
"""
Function that gets the hamiltonian from a file specified in the options
dictionary and returns it as a function of electric and magnetic fields
"""
H_fname = options_dict["H_fname"]
run_dir = options_dict["run_dir"]
H = make_hamiltonian(run_dir + H_fname)
return H
def simulate_SPB(
R_t=None, #Position of molecule as function of time in XYZ coordinate system
T=None, #Total time for which molecule is simulated
E_R=None,
B_R=None, #Spatial profiles of the static electric and magnetic fields
microwave_fields=None, #List of microwave fields for the simulation
initial_states=None, #Initial states for the system (can simulate multiple initial states with little additional cost)
final_state=None, #Desired final state after the state preparation
N_steps=int(1e4), #Number of timesteps for time-evolution
H_path=None, #Path to Hamiltonian if not the default
verbose=False, #Whether or not need to print a bunch of output
plots=False #Whether or not to save plots
):
"""
Function that simulates state preparation B in CeNTREX.
The structure of the cod is as follows:
0. Define the position of the molecule as a function of time. Position is defined in
the capitalized coordinate system(XYZ where Z is along the beamline). Helps with
defining time-dependence of e.g. microwave intensities and electric and magnetic fields.
1. Define electric and magnetic fields as a function of time. This is done based on
spatial profiles of the EM-fields. The coordinate system for the direction of the
the fields is different from the coordinate system used for defining the spatial
profile:
- For defining spatial profile use XYZ (Z is along beamline)
- For defining direction of fields use xyz where z is along the electric field
in the Main Interaction region
For electric fields should use V/cm and for B-fields G.
2. Define slowly time-varying Hamiltonian. This is the Hamiltonian due to the internal
dynamics of the molecule and the external DC electric and magnetic fields
(maybe split into internal Hamiltonian and EM-field Hamiltonian?)
3. Define the microwave Hamiltonian. Simply the couplings due to the microwaves.
Rotating wave approximation is applied and the rotating frame is used.
4. Time-evolve the system by using matrix exponentiation for solving Schroedinger eqn.
5. Make plots etc. if needed
"""
### 0. Position as function of time
#Check if position as function of time was provided, otherwise use a default trajectory
if R_t is None:
def molecule_position(t, r0, v):
"""
Functions that returns position of molecule at a given time for given initial position and velocity.
inputs:
t = time in seconds
r0 = position of molecule at t = 0 in meters
v = velocity of molecule in meters per second
returns:
r = position of molecule in metres
"""
r = r0 + v * t
return r
#Define the position of the molecule as a function of time
r0 = np.array((0, 0, -100e-3))
v = np.array((0, 0, 200))
R_t = lambda t: molecule_position(t, r0, v)
#Define the total time for which the molecule is simulated
z0 = r0[2]
vz = v[2]
T = np.abs(2 * z0 / vz)
### 1. Electric and magnetic fields as function of time
#Here we determine the E- and B-fields as function of time based on provided spatial
#profiles of the fields
#Electric field:
if E_R is not None:
E_t = lambda t: E_R(R_t(t))
#If spatial profile is not provided, assume a uniform 50 V/cm along z
else:
E_t = lambda t: np.array((0, 0, 50.))
#Magnetic field:
if B_R is not None:
B_t = lambda t: B_R(R_t(t))
#If spatial profile is not provided assume uniform 0.001 G along z
else:
B_t = lambda t: np.array((0, 0, 0.001))
### 2. Slowly time-varying Hamiltonian
#Make list of quantum numbers that defines the basis for the matrices
QN = make_QN(0, 3, 1 / 2, 1 / 2)
dim = len(QN)
#Get H_0 from file (H_0 should be in rad/s)
if H_path is None:
H_0_EB = make_hamiltonian(
"./hamiltonians/TlF_X_state_hamiltonian0to3_2020_03_03.py")
else:
H_0_EB = make_hamiltonian(H_path)
#Calculate the field free Hamiltonian
H_0 = H_0_EB(np.array((0, 0, 0)), np.array((0, 0, 0)))
#Calculate part of the Hamiltonian that is due to electromagnetic fields
H_EB_t = lambda t: H_0_EB(E_t(t), B_t(t)) - H_0
#Check that H_0 is hermitian
H_test = H_0 + H_EB_t(T / 2)
is_herm = np.allclose(H_test, np.conj(H_test.T))
if not is_herm:
print("H_0 is non-hermitian!")
### 3. Microwave Hamiltonians
#Generate the part of the Hamiltonian that describes the effect of the microwave
#field(s)
if microwave_fields is not None:
counter = 0
#Containers for microwave coupling Hamiltonians and the microwave frequencies
microwave_couplings = []
omegas = []
D_mu = np.zeros(H_0.shape)
for microwave_field in microwave_fields:
if verbose:
counter += 1
print("\nMicrowave {:}:".format(counter))
#Find out where the peak intensity of the field is located
Z_peak = microwave_field.z0
R_peak = np.array((0, 0, Z_peak))
#Determine the Hamiltonian at z0
H_peak = H_0_EB(E_R(R_peak), B_R(R_peak))
#Find the exact ground and excited states for the field at z_peak
microwave_field.find_closest_eigenstates(H_peak, QN)
#Calculate angular part of matrix element for main transition
ME_main = microwave_field.calculate_ME_main(
pol_vec=microwave_field.p_t(
0)) #Angular part of ME for main transition
#Find the coupling matrices due to the microwave
H_list = microwave_field.generate_couplings(QN)
H_list = [H / ME_main for H in H_list]
Hu_list = [np.triu(H) for H in H_list]
Hl_list = [np.tril(H) for H in H_list]
#Find some necessary parameters and then define the coupling matrix as a function of time
Omega_t = microwave_field.find_Omega_t(
R_t) #Time dependence of Rabi rate
p_t = microwave_field.p_t #Time dependence of polarization of field
omega = np.abs(microwave_field.calculate_frequency(
H_peak, QN)) #Calculate frequency of transition
omegas.append(omega)
D_mu += microwave_field.generate_D(
np.sum(omegas), H_peak,
QN) #Matrix that shifts energies for rotating frame
#Define the coupling matrix as function of time
def H_mu_t_func(Hu_list, Hl_list, p_t, Omega_t, t):
return (
1 / 2 * Omega_t(t) *
(Hu_list[0] * p_t(t)[0] + Hu_list[1] * p_t(t)[1] +
Hu_list[2] * p_t(t)[2] + Hl_list[0] * p_t(t)[0].conj() +
Hl_list[1] * p_t(t)[1].conj() +
Hl_list[2] * p_t(t)[2].conj()))
H_mu_t = partial(H_mu_t_func, Hu_list, Hl_list, p_t, Omega_t)
microwave_couplings.append(H_mu_t)
#Print output for checks
if verbose:
print("ME_main = {:.3E}".format(ME_main))
print("muW frequency: {:.9E} GHz".format(omega /
(2 * np.pi * 1e9)))
print("ground_main:")
ground_main = microwave_field.ground_main
ground_main.print_state()
print("excited_main:")
excited_main = microwave_field.excited_main
excited_main.print_state()
ME = excited_main.state_vector(QN).conj().T @ H_mu_t(
T / 2) @ ground_main.state_vector(QN)
print("Omega_T/2 = {:.3E}".format(ME / (2 * np.pi)))
#Generate function that gives couplings due to all microwaves
def H_mu_tot_t(t):
H_mu_tot = microwave_couplings[0](t)
if len(microwave_couplings) > 1:
for H_mu_t in microwave_couplings[1:]:
H_mu_tot = H_mu_tot + H_mu_t(t)
return H_mu_tot
else:
H_mu_tot_t = lambda t: np.zeros(H_0.shape)
if verbose:
time = timeit.timeit("H_mu_tot_t(T/2)", number=10,
globals=locals()) / 10
print("Time to generate H_mu_tot_t: {:.3e} s".format(time))
H_test = H_mu_tot_t(T / 2)
non_zero = H_test[np.abs(H_test) > 0].shape[0]
print("Non-zero elements at T/2: {}".format(non_zero))
print("Rotating frame energy shifts:")
print(np.diag(D_mu) / (2 * np.pi * 1e9))
print(np.max(np.abs(H_mu_tot_t(T / 2))) / (2 * np.pi))
### 4. Time-evolution
#Now time-evolve the system by repeatedly applying the time-evolution operator with small
# timestep
#Make array of times at which time-evolution is performed
t_array = np.linspace(0, T, N_steps)
#Calculate timestep
dt = T / N_steps
#Set the state vector to the initial state
H_t0 = H_0 + H_EB_t(0)
psis = np.zeros((len(initial_states), H_t0.shape[0]), dtype=complex)
ini_index = np.zeros(len(initial_states))
for i, initial_state in enumerate(initial_states):
initial_state = find_closest_state(H_t0, initial_state, QN)
psis[i, :] = initial_state.state_vector(QN)
#Also figure out the state indices that correspond to the initial states
ini_index[i] = find_state_idx_from_state(H_t0, initial_state, QN)
#Set up containers to store results
state_probabilities = np.zeros(
(len(initial_states), len(t_array), H_t0.shape[0]))
state_energies = np.zeros((len(t_array), H_t0.shape[0]))
psi_t = np.zeros((len(initial_states), len(t_array), H_t0.shape[0]),
dtype=complex)
#Initialize the reference matrix of eigenvectors
E_ref, V_ref = np.linalg.eigh(H_t0)
E_ref_0 = E_ref
V_ref_0 = V_ref
#Loop over timesteps to evolve system in time
for i, t in enumerate(tqdm(t_array)):
#Calculate the necessary Hamiltonians at this time
H_t = H_0 + H_EB_t(t)
H_mu = H_mu_tot_t(t)
#Diagonalize H_t and transform to that basis
D_t, V_t, info_t = zheevd(H_t)
if info_t != 0:
print("zheevd didn't work for H_0")
D_t, V_t = np.linalg.eigh(H_t)
#Make intermediate hamiltonian by transforming H to the basis where H_0 is diagonal
H_I = V_t.conj().T @ H_t @ V_t
#Sort the eigenvalues so they are in ascending order
index = np.argsort(D_t)
D_t = D_t[index]
V_t = V_t[:, index]
#Find the microwave coupling matrix in the new basis:
H_mu = V_t.conj().T @ H_mu @ V_t
#Make total intermediate hamiltonian
H_I = H_I + H_mu
#Transform H_I to the rotating basis
H_R = H_I + D_mu
#Diagonalize H_R
D_R, V_R, info_R = zheevd(H_R)
if info_R != 0:
print("zheevd didn't work for H_R")
D_R, V_R = np.linalg.eigh(H_R)
#Reorder the eigenstates so that a given index corresponds to the same "adiabatic" state
energies, evecs = D_t, V_t
energies, evecs = reorder_evecs(evecs, energies, V_ref)
#Propagate state vector in time
propagator = V_t @ V_R @ np.diag(np.exp(
-1j * D_R * dt)) @ V_R.conj().T @ V_t.conj().T
for j in range(0, len(initial_states)):
psis[j, :] = propagator @ psis[j, :]
#Calculate the overlap of the state vector with all of the eigenvectors of the Hamiltonian
overlaps = np.dot(np.conj(psis[j, :]).T, evecs)
probabilities = overlaps * np.conj(overlaps)
#Store the probabilities, energies and the state vector
state_probabilities[j, i, :] = np.real(probabilities)
psi_t[j, i, :] = psis[j, :].T
state_energies[i, :] = energies
V_ref = evecs
#Find the exact form of the desired final state and calculate overlap with state vector
final_state = find_closest_state(H_0 + H_EB_t(T), final_state, QN)
final_state_vec = final_state.state_vector(QN)
if len(initial_states) == 1:
overlap = final_state_vec.conj().T @ psis[0, :]
probability = np.abs(overlap)**2
if verbose:
print(
"Probability to end up in desired final state: {:.3f}".format(
probability))
print("desired final state|fin> = ")
final_state.print_state()
return t_array, state_probabilities, V_ref, QN
def simulate_RAP(r0=np.array((0, 0, -100e-3)),
v=np.array((0, 0, 200)),
ring_z1=-71.4375e-3,
ring_z2=85.725e-3,
ring_V1=4e3,
ring_V2=3.6e2,
B_earth=np.array((0, 0, 0)),
Omega1=2 * np.pi * 200e3,
Omega2=2 * np.pi * 200e3,
Delta=0,
delta=0,
z0_mu1=0.0,
z0_mu2=0.03,
N_steps=1e5):
"""
Function that runs the rapid adiabatic passage simulation for the given
parameters
inputs:
r0 = initial position of molecule [m]
v = velocity of molecule (assumed to be constant) [m/s]
ring_z1 = position of first ring electrode [m]
ring_z2 = position of second ring electrode [m]
ring_V1 = voltage on ring 1 [V]
ring_V2 = voltage on ring 2 [V]
B_earth = magnetic field of earth [G]
Omega1 = Rabi rate for microwave 1 [rad/s]
Omega2 = Rabi rate for microwave 2 [rad/s]
Delta = 1-photon detuning [rad/s]
delta = 2-photon detuning [rad/s]
z0_mu1 = z-position of microwave 1 [m]
z0_mu2 = z-position of microwave 2 [m]
returns:
probability = probability of being in the target final state
"""
#Define the position of the molecule as a fucntion of time
r0 = np.array((0, 0, -100e-3))
v = np.array((0, 0, 200))
r_t = lambda t: molecule_position(t, r0, v)
#Define the total time for which the molecule is simulated
z0 = r0[2]
vz = v[2]
T = np.abs(2 * z0 / vz)
#Define electric field as function of position
E_r = lambda r: (E_field_ring(r, z0=ring_z1, V=ring_V1) + E_field_ring(
r, z0=ring_z2, V=ring_V2))
#Next define electric field as function of time
E_t = lambda t: E_r(r_t(t))
#Define the magnetic field
B_r = lambda r: B_earth
B_t = lambda t: B_r(r_t(t))
#Make list of quantum numbers that defines the basis for the matrices
QN = make_QN(0, 3, 1 / 2, 1 / 2)
#Get H_0 from file (H_0 should be in rad/s)
H_0_EB = make_hamiltonian("../utility/TlF_X_state_hamiltonian0to3.py")
#Make H_0 as function of time
H_0_t = lambda t: H_0_EB(E_t(t), B_t(t))
dim = H_0_t(0).shape[0]
#Fetch the initial state from file
with open("../utility/initial_state0to3.pickle", 'rb') as f:
initial_state_approx = pickle.load(f)
#Fetch the intermediate state from file
with open("../utility/intermediate_state0to3.pickle", "rb") as f:
intermediate_state_approx = pickle.load(f)
#Fetch the final state from file
with open("../utility/final_state0to3.pickle", 'rb') as f:
final_state_approx = pickle.load(f)
#Find the eigenstate of the Hamiltonian that most closely corresponds to initial_state at T=0. This will be used as the
#actual initial state
initial_state = find_closest_state(H_0_t(0), initial_state_approx, QN)
initial_state_vec = initial_state.state_vector(QN)
#Find the eigenstate of the Hamiltonian that most closely corresponds to final_state_approx at t=T. This will be used to
#determine what fraction of the molecules end up in the correct final state
final_state = find_closest_state(H_0_t(T), final_state_approx, QN)
final_state_vec = final_state.state_vector(QN)
intermediate_state = find_closest_state(H_0_t(0),
intermediate_state_approx, QN)
intermediate_state_vec = intermediate_state.state_vector(QN)
#Find the energies of ini and int so that suitable microwave frequency for mu1 can be calculated
H_z0 = H_0_EB(E_r(np.array((0, 0, z0_mu1))), B_r(np.array((0, 0, z0_mu1))))
D_z0, V_z0 = np.linalg.eigh(H_z0)
ini_index = find_state_idx_from_state(H_z0, initial_state_approx, QN)
int_index = find_state_idx_from_state(H_z0, intermediate_state_approx, QN)
fin_index = find_state_idx_from_state(H_z0, final_state_approx, QN)
#Note: the energies are in 2*pi*[Hz]
E_ini = D_z0[ini_index]
E_int = D_z0[int_index]
omega_mu1 = E_int - E_ini
#Find the energies of int and fin so that suitable microwave frequency for mu2 can be calculated
H_z0 = H_0_EB(E_r(np.array((0, 0, z0_mu2))), B_r(np.array((0, 0, z0_mu2))))
D_z0, V_z0 = np.linalg.eigh(H_z0)
ini_index = find_state_idx_from_state(H_z0, initial_state_approx, QN)
int_index = find_state_idx_from_state(H_z0, intermediate_state_approx, QN)
fin_index = find_state_idx_from_state(H_z0, final_state_approx, QN)
#Note: the energies are in 2*pi*[Hz]
E_int = D_z0[int_index]
E_fin = D_z0[fin_index]
omega_mu2 = E_fin - E_int
#Define dipole moment of TlF
D_TlF = 2 * np.pi * 4.2282 * 0.393430307 * 5.291772e-9 / 4.135667e-15 # [rad/s/(V/cm)]
#Calculate the approximate power required for each set of microwaves to get a specified peak Rabi rate for the transitions
ME1 = np.abs(
calculate_microwave_ED_matrix_element_mixed_state_uncoupled(
initial_state_approx, intermediate_state_approx, reduced=False))
ME2 = np.abs(
calculate_microwave_ED_matrix_element_mixed_state_uncoupled(
intermediate_state_approx, final_state_approx, reduced=False))
P1 = calculate_power_needed(Omega1, ME1)
P2 = calculate_power_needed(Omega2, ME2)
#Define the microwave electric field as a function of time
E_mu1_t = lambda t: microwave_field(
r_t(t), power=P1, z0=z0_mu1, fwhm=1.1 * 0.0254)
E_mu2_t = lambda t: microwave_field(
r_t(t), power=P2, z0=z0_mu2, fwhm=1.1 * 0.0254)
#Define matrix for microwaves coupling J = 0 to 1
J1_mu1 = 0
J2_mu1 = 1
omega_mu1 = omega_mu1 + Delta
H1 = make_H_mu(J1_mu1, J2_mu1, omega_mu1, QN)
H_mu1 = lambda t: H1(0) * D_TlF * E_mu1_t(t)
#Define matrix for microwaves coupling J = 1 to 2
J1_mu2 = 1
J2_mu2 = 2
omega_mu2 = omega_mu2 + delta - Delta
H2 = make_H_mu(J1_mu2, J2_mu2, omega_mu2, QN)
H_mu2 = lambda t: H2(0) * D_TlF * E_mu2_t(t)
#Make the matrices used to transform to rotating frame
U1, D1 = make_transform_matrix(J1_mu1, J2_mu1, omega_mu1, QN)
U2, D2 = make_transform_matrix(J1_mu2, J2_mu2, omega_mu2 + omega_mu1, QN)
U = lambda t: U1(t) @ U2(t)
D = lambda t: D1 + D2
#Define number of timesteps and make a time-array
t_array = np.linspace(0, T, N_steps)
#Calculate timestep
dt = T / N_steps
#Set system to its initial state
psi = initial_state_vec
#Loop over timesteps to evolve system in time
for i, t in enumerate(t_array):
#Calculate the necessary Hamiltonians at this time
H_0 = H_0_t(t)
H_mu1_i = H_mu1(t)
H_mu2_i = H_mu2(t)
#Diagonalize H_0 and transform to that basis
D_0, V_0, info_0 = zheevd(H_0)
if info_0 != 0:
print("zheevd didn't work for H_0")
D_0, V_0 = np.linalg.eigh(H_0)
#Make intermediate hamiltonian by transforming H to the basis where H_0 is diagonal
H_I = V_0.conj().T @ H_0 @ V_0
#Sort the eigenvalues so they are in ascending order
index = np.argsort(D_0)
D_0 = D_0[index]
V_0 = V_0[:, index]
#Find the microwave coupling matrix:
H_mu1_i = V_0.conj().T @ H_mu1_i @ V_0 * np.exp(1j * omega_mu1 * t)
H_mu1_i = block_diag(H_mu1_i[0:16, 0:16], np.zeros(
(dim - 16, dim - 16)))
H_mu1_i = np.triu(H_mu1_i) + np.triu(
H_mu1_i).conj().T #+ np.diag(np.diag(H1))
H_mu2_i = V_0.conj().T @ H_mu2_i @ V_0 * np.exp(1j * omega_mu2 * t)
H_mu2_i = block_diag(np.zeros((4, 4)), H_mu2_i[4:36, 4:36],
np.zeros((dim - 36, dim - 36)))
H_mu2_i = np.triu(H_mu2_i) + np.triu(
H_mu2_i).conj().T #+ np.diag(np.diag(H1))
#Make total hamiltonian
H_I = H_I + H_mu1_i + H_mu2_i
#Find transformation matrices for rotating basis
U_t = U(t)
D_t = D(t)
#Transform H_I to the rotating basis
H_R = U_t.conj().T @ H_I @ U_t + D_t
#Diagonalize H_R
D_R, V_R, info_R = zheevd(H_R)
if info_R != 0:
print("zheevd didn't work for H_R")
D_R, V_R = np.linalg.eigh(H_R)
#Propagate state vector in time
psi = V_0 @ V_R @ np.diag(np.exp(
-1j * D_R * dt)) @ V_R.conj().T @ V_0.conj().T @ psi
psi_fin = vector_to_state(psi, QN)
#Calculate overlap between final target state and psi
overlap = final_state_vec.conj().T @ U(T) @ psi
probability = np.abs(overlap)**2
return probability
def OBE_integrator(
r0=np.array((0., 0., 0.)),
r1=np.array((0, 0, 3e-2)),
v=np.array(
(0, 0, 200)
), #Parameters for defining position of molecule as function of time
X_states=None,
B_states=None, #States that are needed for simulation
microwave_fields=None,
laser_fields=None, #Lists of MicrowaveFields and LaserFields
Gamma=2 * np.pi * 1.6e6, #Natural linewidth of the excited state
states_pop=None,
pops=None, #States that are populated and their populations
Nsteps=int(5e3),
dt=None,
dt_max=1e-5, #Number of timesteps, or alternatively size of timesteps
method='exp', #Method to use for time-integration
verbose=True, #Whether or not to print outputs for debugging
):
"""
Function that integrates optical Bloch equations for TlF in Centrex.
The structure of the code is as follows:
0. Define position of molecule as function of time. Makes it easier to
convert spatial dependence of e.g. laser/microwave intensity to a
time dependence.
1. Define the internal Hamiltonian of the molecule. The Hamiltonian
for the X- and B-states of TlF is fetched from file and diagonalized.
The eigenstates of the Hamiltonians are used as the basis in the rest
of the calculation. The size of the basis is reduced by keeping only
states that couple to one of the time-dependent fields, or which can
be reached via spontaneous decay from the excited state.
1.1 Define X-state Hamiltonian and find the diagonal basis
1.2 Define B-state Hamiltonian and find the diagonal basis
2. Microwave couplings. Construct Hamiltonian for microwave couplings between different
rotational states within X based on the list of microwave fields that is provided.
3. Optical couplings. Construct the Hamiltonian for laser couplings between X and B. Can
have multiple laser fields provided in a list of LaserField objects.
4. Generate the total Hamiltonian that contains the couplings due to the fields, and the
internal molecular Hamiltonian
5. Generate collapse operators that describe spontaneous decay
6. Generate initial density matrix
7. Transform to Liouville space if using exponentiation method
8. Time-evolution
inputs:
r0 = initial position of molecule [m]
r1 = final positon of molecule [m]
v = velocity of molecule [m/s]
X_states = list of states in X-state of TlF used in the simulation (list of State objects)
B_states = list of states in B-state of TlF used in the simulation (list of State objects)
laser_fields = list of OpticalField objects (see OBE_classes.py)
microwave_fields = list of MicrowaveField objects (see OBE_classes.py)
states_pops = states that are initially populated
pops = initial populations of states_pops. If none, Boltzmann distribution assumed
outputs:
t_array = array of times at which we have datapoints [s]
pop_results = populations in each state at the times stored in t_array
"""
### 0. Define molecular position as function of time ###
def molecule_position(t, r0, v):
"""
Function that returns position of molecule at a given time for given initial position and
velocity.
inputs:
t = time in seconds
r0 = position of molecule at t = 0 in meters
v = velocity of molecule in meters per second
returns:
r = position of molecule in metres
"""
r = r0 + v * t
return r
#Define a lambda function that gives position as function of time
r_t = lambda t: molecule_position(t, r0, v)
#Calculate total time for simulation [s]
z0 = r0[2]
z1 = r1[2]
vz = v[2]
T = np.abs((z1 - z0) / vz)
# If want printed output, print it
if verbose:
print("Total simulation time is {:.3E} s".format(T))
### 1. Internal Hamiltonian
## 1.1 X-state of TlF
#Load Hamiltonian from file
H_X_uc = make_hamiltonian(
"./utilities/TlF_X_state_hamiltonian_J0to4.pickle")
#Hamiltonian on file is in uncoupled angular momentum basis. Transform it to coupled.
#Load transform matrix
with open("./utilities/UC_to_C_j0to4.pickle", "rb") as f:
S_trans = pickle.load(f)
#Transform matrix
E = np.array((0, 0, 0))
B = np.array(
(0, 0, 0.001)
) #Very small magnetic field is used to ensure mF is a good quantum number
H_X = S_trans.conj().T @ H_X_uc(E, B) @ S_trans
#Define basis to which the Hamiltonian was transformed
Jmin = 0
Jmax = 4
I_F = 1 / 2
I_Tl = 1 / 2
QN_X = [
CoupledBasisState(F,
mF,
F1,
J,
I_F,
I_Tl,
electronic_state='X',
P=(-1)**J,
Omega=0) for J in ni_range(Jmin, Jmax + 1)
for F1 in ni_range(np.abs(J - I_F), J + I_F + 1)
for F in ni_range(np.abs(F1 - I_Tl), F1 + I_Tl + 1)
for mF in ni_range(-F, F + 1)
]
#Diagonalize the Hamiltonian (also making sure V is as close as identity matrix as possible
# in terms of ordering of states)
D, V = np.linalg.eigh(H_X)
V_ref_X = np.eye(V.shape[0])
D, V = reorder_evecs(V, D, V_ref_X)
H_X_diag = V.conj().T @ H_X @ V
#Define new basis based on eigenstates of H_X:
QN_X_diag = matrix_to_states(V, QN_X)
#Sometimes only a subset of states is needed for the simulation. Determine the X-states
#that are needed here.
if X_states is None:
ground_states = QN_X_diag
else:
ground_states = find_exact_states(X_states,
H_X_diag,
QN_X_diag,
V_ref=V_ref_X)
#Find the Hamiltonian in the reduced basis
H_X_red = reduced_basis_hamiltonian(QN_X_diag, H_X_diag, ground_states)
#Set small off diagonal terms to zero
H_X_red[np.abs(H_X_red) < 0.1] = 0
## 1.2 B-state of TlF
#Load Hamiltonian from file
H_B = make_hamiltonian_B(
"./utilities/B_hamiltonians_symbolic_coupled_P_1to3.pickle")
#Define the basis that the Hamiltonian is in
Jmin = 1
Jmax = 3
I_F = 1 / 2
I_Tl = 1 / 2
Ps = [-1, 1]
QN_B = [
CoupledBasisState(F,
mF,
F1,
J,
I_F,
I_Tl,
P=P,
Omega=1,
electronic_state='B')
for J in ni_range(Jmin, Jmax + 1)
for F1 in ni_range(np.abs(J - I_F), J + I_F + 1)
for F in ni_range(np.abs(F1 - I_Tl), F1 + I_Tl + 1)
for mF in ni_range(-F, F + 1) for P in Ps
]
#Diagonalize the Hamiltonian
D, V = np.linalg.eigh(H_B)
V_ref_B = np.eye(H_B.shape[0])
D, V = reorder_evecs(V, D, V_ref_B)
H_B_diag = V.conj().T @ H_B @ V
#Define new basis based on eigenstates of H_B
QN_B_diag = matrix_to_states(V, QN_B)
#Sometimes only a subset of states is needed for the simulation. Determine the X-states
#that are needed here.
if B_states is None:
excited_states = QN_B_diag
else:
excited_states = find_exact_states(B_states,
H_B_diag,
QN_B_diag,
V_ref=V_ref_B)
#Find the Hamiltonian in the reduced basis
H_B_red = reduced_basis_hamiltonian(QN_B_diag, H_B_diag, excited_states)
#Set small off diagonal terms to zero
H_B_red[np.abs(H_B_red) < 0.1] = 0
## 1.3 Define total internal Hamiltonian
H_int = scipy.linalg.block_diag(H_X_red, H_B_red)
V_ref_int = np.eye(H_int.shape[0])
#Define Hamiltonian in the rotating frame (transformation not applied yet)
H_rot = H_int
#Define QN for the total Hamiltonian that includes both X and B
QN = ground_states + excited_states
if verbose:
print("Diagonal of H_int:")
print(np.diag(H_int) / (2 * np.pi))
### 2. Couplings due to microwaves
#If there are microwave fields, loop over them. Otherwise skip this section.
if microwave_fields is not None:
microwave_couplings = []
omegas = []
for microwave_field in microwave_fields:
#Find the exact ground and excited states for the field
microwave_field.find_closest_eigenstates(H_rot, QN, V_ref_int)
#Calculate angular part of matrix element for main transition
ME_main = microwave_field.calculate_ME_main(
) #Angular part of ME for main transition
#Find the coupling matrices due to the laser
H_list = microwave_field.generate_couplings(QN)
H_list = [H / ME_main for H in H_list]
Hu_list = [np.triu(H) for H in H_list]
Hl_list = [np.tril(H) for H in H_list]
#Find some necessary parameters and then define the coupling matrix as a function of time
Omega_t = microwave_field.find_Omega_t(
r_t) #Time dependence of Rabi rate
p_t = microwave_field.p_t #Time dependence of polarization of field
omega = microwave_field.calculate_frequency(
H_rot, QN) #Calculate frequency of transition
D_mu = microwave_field.generate_D(
omega, H_rot, QN,
V_ref_int) #Matrix that shifts energies for rotating frame
H_rot += D_mu
#Define the coupling matrix as function of time
def H_mu_t_func(Hu_list, Hl_list, p_t, Omega_t, t):
return (
Omega_t(t) *
(Hu_list[0] * p_t(t)[0] + Hu_list[1] * p_t(t)[1] +
Hu_list[2] * p_t(t)[2] + Hl_list[0] * p_t(t)[0].conj() +
Hl_list[1] * p_t(t)[1].conj() +
Hl_list[2] * p_t(t)[2].conj()))
H_mu_t = partial(H_mu_t_func, Hu_list, Hl_list, p_t, Omega_t)
microwave_couplings.append(H_mu_t)
#Print output for checks
if verbose:
print("ME_main = {:.3E}".format(ME_main))
i_e = QN.index(microwave_field.excited_main)
i_g = QN.index(microwave_field.ground_main)
print(H_mu_t(T / 2)[i_e, i_g] / (2 * np.pi * 1e6))
print(Omega_t(T / 2))
#Generate function that gives couplings due to all microwaves
def H_mu_tot_t(t):
H_mu_tot = microwave_couplings[0](t)
if len(microwave_couplings) > 1:
for H_mu_t in microwave_couplings[1:]:
H_mu_tot = H_mu_tot + H_mu_t(t)
return H_mu_tot
if verbose:
with open("H_mu_tot.pickle", 'wb+') as f:
pickle.dump(H_mu_tot_t(T / 2), f)
#Shift energies in H_int in accordance with the rotating frame
#H_rot = H_rot + D_mu
else:
H_mu_tot_t = lambda t: np.zeros(H_rot.shape)
if verbose:
time = timeit.timeit("H_mu_tot_t(T/2)", number=10,
globals=locals()) / 10
print("Time to generate H_mu_tot_t: {:.3e} s".format(time))
H_test = H_mu_tot_t(T / 2)
non_zero = H_test[np.abs(H_test) > 0].shape[0]
print("Non-zero elements at T/2: {}".format(non_zero))
### 3. Optical couplings due to laser
#If there are laser fields, loop over them. Otherwise skip this section
if laser_fields is not None:
optical_couplings = []
optical_Ls = []
D_laser = np.zeros(H_rot.shape)
for laser_field in laser_fields:
#Find the exact ground and excited states for the field
laser_field.find_closest_eigenstates(H_rot, QN, V_ref_int)
#Calculate angular part of matrix element for main transition
ME_main = laser_field.calculate_ME_main(
) #Angular part of ME for main transition
#Find the coupling matrices due to the laser
H_list = laser_field.generate_couplings(QN)
H_list = [H / ME_main for H in H_list]
Hu_list = [np.triu(H) for H in H_list]
Hl_list = [np.tril(H) for H in H_list]
#Find some necessary parameters and then define the coupling matrix as a function of time
p_t = laser_field.p_t #Time dependence of polarization of field (includes phase modulation)
Omega_t = laser_field.find_Omega_t(
r_t) #Time dependence of Rabi rate
D_laser += laser_field.generate_D(
H_rot, QN) #Matrix that shifts energies for rotating frame
H_rot = H_rot + laser_field.generate_D(H_rot, QN)
#Define the optical coupling Hamiltonian as function of time
def H_oc_t_func(Hu_list, Hl_list, p_t, Omega_t, t):
return (
Omega_t(t) *
(Hu_list[0] * p_t(t)[0] + Hu_list[1] * p_t(t)[1] +
Hu_list[2] * p_t(t)[2] + Hl_list[0] * p_t(t)[0].conj() +
Hl_list[1] * p_t(t)[1].conj() +
Hl_list[2] * p_t(t)[2].conj()))
H_oc_t = partial(H_oc_t_func, Hu_list, Hl_list, p_t, Omega_t)
optical_couplings.append(H_oc_t)
#Print output for checks
if verbose:
print("ME_main = {:.3E}".format(ME_main))
i_e = QN.index(laser_field.excited_main)
i_g = QN.index(laser_field.ground_main)
print(H_oc_t(T / 2)[i_e, i_g] / (2 * np.pi * 1e6))
print(Omega_t(T / 2) / (1e6 * 2 * np.pi))
print("excited main:")
laser_field.ground_main.print_state()
#Functions for total Hamiltonian and Lindbladian due to optical couplings
def H_oc_tot_t(t):
H_oc_tot = optical_couplings[0](t)
if len(optical_couplings) > 1:
for H_oc_t in optical_couplings[1:]:
H_oc_tot = H_oc_tot + H_oc_t(t)
return H_oc_tot
#Shift energies in H_rot in accordance with the rotating frame
#Also shift the energies so that ground_main is at zero energy
i_g = QN.index(laser_field.ground_main)
H_rot = H_rot - np.eye(H_rot.shape[0]) * H_rot[i_g, i_g]
#If no laser fields are defined, set coupling matrix to zeros
else:
H_oc_tot_t = lambda t: np.zeros(H_rot.shape)
if verbose:
time = timeit.timeit("H_oc_tot_t(T/2)", number=10,
globals=locals()) / 10
print("Time to generate H_oc_tot_t: {:.3e} s".format(time))
print("Diagonal of H_rot in rotating frame of laser:")
print(np.diag(H_rot) / (2 * np.pi))
print("D_laser:")
print(np.diag(D_laser))
# with open("H_oc_tot.pickle",'wb+') as f:
# pickle.dump(H_oc_tot_t(T/2.3156165),f)
### 4. Total Hamiltonian
#Define the total Hamiltonian (including the oscillating fields) in the rotating frame
# as a function of time
H_tot_t = lambda t: H_rot + H_oc_tot_t(t) + H_mu_tot_t(t)
if verbose:
time = timeit.timeit("H_tot_t(T/2)", number=10, globals=locals()) / 10
print("Time to generate H_tot_t: {:.3e} s".format(time))
print("Diagonal of H_tot_t(T/2) in rotating frame of laser:")
print(np.diag(H_tot_t(T / 2)) / (2 * np.pi))
# print("D_laser:")
# print(np.diag(D_laser))
### 5. Collapse operators
#Here we generate the matrices that describe spontaneous decay from the excited
#states to the ground states
C_list = collapse_matrices(QN, ground_states, excited_states, gamma=Gamma)
#Generate the superoperator that contains spontaneous decay
#This is constant in time so only generated once
L_collapse = np.zeros((len(QN)**2, len(QN)**2), dtype=complex)
for C in tqdm(C_list):
L_collapse += (generate_superoperator(C,
C.conj().T) - 1 / 2 *
(generate_flat_superoperator(C.conj().T @ C) +
generate_sharp_superoperator(C.conj().T @ C)))
#Make the collapse operator into a sparse matrix
L_collapse = csr_matrix(L_collapse)
### 6. Initial populations
#Find the exact forms of the states that are initially populated
states_pop = find_exact_states(states_pop, H_rot, QN, V_ref=V_ref_int)
#If populations in states are not specified, assume Boltzmann distribution
if pops is None:
pops = find_boltzmann_pops(states_pop)
pops = pops / np.sum(pops)
#Generate initial density matrix
rho_ini = generate_density_matrix(QN, states_pop, pops)
if verbose:
print("Initial population in")
states_pop[0].print_state()
print("is {:.5f}".format(pops[0]))
if method == 'exp':
### 7. Transfer to Liouville space
#We transfer to Liouville space where the density matrix is a vector
#and time-evolution is found by matrix exponentiation
#Generate the Lindbladian
#Define a function that gives the time dependent part of the Liouvillian
#at time t
L_t = lambda t: (generate_commutator_superoperator(
coo_matrix(-1j * H_tot_t(t))) + L_collapse)
if verbose:
time = timeit.timeit("L_t(T/2)", number=10, globals=locals()) / 10
print("Time to generate L_t: {:.3e} s".format(time))
L_test = L_t(T / 2)
non_zero = L_test[np.abs(L_test) > 0]
### 8. Time-evolution
#Here we perform the time-evolution of the system
#Set rho vector to its initial value
rho_vec = generate_rho_vector(rho_ini)
#Calculate timestep
if dt is None:
dt = T / Nsteps
else:
Nsteps = int(T / dt) + 1
#Generate array of times
t_array = np.linspace(0, T, Nsteps)
#Pre-calculate some parameters for the matrix exponentiation
#Calculate onenorm estimate
norm = onenormest(L_t(T / 2))
#Calculate wsp and iwsp
m = 20 #maximum size of Krylov subspace
n = rho_vec.shape[0]
wsp = np.zeros(7 + n * (m + 2) + 5 * (m + 2) * (m + 2),
dtype=np.complex128)
iwsp = np.zeros(m + 2, dtype=np.int32)
#Array for storing results
pop_results = np.zeros((len(QN), len(t_array)), dtype=float)
pop_results[:, 0] = np.real(np.diag(rho_ini))
#Loop over timesteps
for i, t_i in enumerate(tqdm(t_array[1:])):
#Calculate the Lindbladian at this time
L_sparse = L_t(t_i)
#Time evolve the density vector
rho_vec = py_zgexpv(rho_vec,
L_sparse,
t=dt,
anorm=norm,
wsp=wsp,
iwsp=iwsp,
m=m)
#Convert back to density matrix
rho = rho_vec.reshape(len(QN), len(QN))
#Find populations in each state
pop_results[:, i + 1] = np.real(np.diag(rho))
# if verbose:
# time = timeit.timeit("py_zgexpv(rho_vec, L_sparse, t = dt, anorm = norm, wsp = wsp, iwsp = iwsp,m = m)",
# number = 10, globals = locals())/10
# print("Time for exponentiating: {:.3E}".format(time))
elif method == 'RK45':
### 7. Setting up ODE solver
#Define a function that generates the RHS of the OBEs in vector format
#Still useful to use Liouville space for the collapse operators since they can be
#fully precalculated (see rhs_C below)
def Lindblad_rhs(t, rho_vec):
dim = int(np.sqrt(len(rho_vec)))
rho = rho_vec.reshape((dim, dim))
rhs_H = (-1j * (H_tot_t(t) @ rho - rho @ H_tot_t(t))).flatten()
rhs_C = L_collapse @ rho_vec
rhs = rhs_H + rhs_C
return rhs
if verbose:
time = timeit.timeit(
"rhs_test = Lindblad_rhs(T/2,rho_ini.flatten())",
number=10,
globals=locals()) / 10
print("Time to generate RHS: {:.3e} s".format(time))
### 8. Time-evolution
#Perform time evolution using IVP solver from scipy
t_span = (0, T)
sol = solve_ivp(Lindblad_rhs,
t_span,
rho_ini.flatten(),
dense_output=True,
max_step=dt_max,
method='RK45')
#Get t_array and populations from solution object
t_array = sol.t
pop_results = np.real(
np.einsum(
'jji->ji',
sol.y.reshape(
(rho_ini.shape[0], rho_ini.shape[1], sol.y.shape[1]))))
else:
raise NotImplementedError("Time integation method not implemented")
return t_array, pop_results