def setup(self):
"""Sets up matrices for Hamiltonian construction."""
with timer("Generating states"):
if self.hamiltonian.endswith("relevant"):
self.states = States(self.n, basis=Basis.N_L_ML_MS_RELEVANT)
print("Loaded relevant N L ML MS states.")
else:
self.states = States(self.n, basis=Basis.N_L_ML_MS)
print("Loaded N L ML MS states.")
with timer("Loading Hamiltonian"):
mat_1, mat_1_zeeman, mat_2, mat_2_minus, mat_2_plus = load_hamiltonian(
self.hamiltonian)
mat_2_combination = mat_2_plus + mat_2_minus
with timer("Loading transformations"):
transform_1 = load_transformation(self.n, Basis.N_L_J_MJ_RELEVANT,
Basis.N_L_ML_MS_RELEVANT)
with timer("Applying transformation to nlmlms"):
mat_1 = transform_basis(mat_1, transform_1)
mat_1_zeeman = transform_basis(mat_1_zeeman, transform_1)
mat_2 = transform_basis(mat_2, transform_1)
# mat_2_plus = transform_basis(mat_2_plus, transform_1)
# mat_2_minus = transform_basis(mat_2_minus, transform_1)
mat_2_combination = transform_basis(mat_2_combination, transform_1)
self.mat_1 = mat_1
self.mat_1_zeeman = mat_1_zeeman
self.mat_2 = mat_2
# self.mat_2_plus = mat_2_plus
# self.mat_2_minus = mat_2_minus
self.mat_2_combination = mat_2_combination
def nlmlms_to_n1n2mlms(n) -> np.ndarray:
"""
Generates the transformation matrix for the transformation from the
n, l, ml, ms basis to the n1, n2, ml, ms basis.
CG coefficients generated according to:
Park, D. Relation between the parabolic and spherical eigenfunctions of hydrogen. Z. Physik 159, 155–157 (1960).
https://doi.org/10.1007/BF01338343
:param n:
:return:
"""
# source_states = States(n, Basis.N_L_ML_MS)
source_states = States(n, Basis.N_L_ML_MS_RELEVANT)
target_states = States(n, Basis.N1_N2_ML_MS)
dimension = len(source_states.states)
identity = np.identity(dimension)
transform = np.zeros_like(identity)
for ii in trange(dimension, desc="Convert to n1n2"):
n1, n2, _ml, _ms = target_states.states[ii]
coeff_sum = 0
for jj in range(dimension):
__n, __l, __ml, __ms = source_states.states[jj]
if _ms != __ms:
continue
# Satisfies: -K <= k1, k2 <= K, m = k1 + k2, n = 2K + 1 = n1 + n2 + |m| + 1
k1 = (_ml + n1 - n2) / 2
k2 = (_ml - n1 + n2) / 2
K = (n - 1) / 2
try:
# See citation in docstring for source
coeff = CG(
j1=K, j2=K, j3=__l,
m1=k1, m2=k2, m3=__ml,
)
coeff_sum += coeff ** 2
if coeff != 0:
transform[ii] += coeff * identity[jj]
except (ValueError, AttributeError):
continue
if abs(coeff_sum - 1) > 1e-3:
logger.error(f"CG coeff sum discrepancy exceed threshold: {coeff_sum} ({n1}, {n2}, {_ml}, {_ms})")
return transform
def nljmj_to_nlmlms(n) -> np.ndarray:
"""
Generates the transformation matrix for the transformation from the
n, l, j, mj basis to the n, l, ml, ms basis.
:param n:
:return:
"""
# source_states = States(n, Basis.N_L_J_MJ)
# target_states = States(n, Basis.N_L_ML_MS)
source_states = States(n, Basis.N_L_J_MJ_RELEVANT)
target_states = States(n, Basis.N_L_ML_MS_RELEVANT)
dimension = len(source_states.states)
identity = np.identity(dimension)
transform = np.zeros_like(identity)
for ii in trange(dimension, desc="Convert to nlmlms"):
_n, _l, _ml, _ms = target_states.states[ii]
coeff_sum = 0
for jj in range(dimension):
__n, __l, __j, __mj = source_states.states[jj]
if _l != __l:
continue
try:
coeff = CG(
j1=_l, j2=0.5, j3=__j,
m1=_ml, m2=_ms, m3=__mj,
)
coeff_sum += coeff ** 2
if coeff != 0:
logger.info(
f"j,mj to ml,ms: {coeff:.2f}. \tn {_n}, \tl {_l}, \tml {_ml}, \tms {_ms}, \tj {__j}, \tmj {__mj}")
transform[ii] += coeff * identity[jj]
except (ValueError, AttributeError):
continue
if abs(coeff_sum - 1) > 1e-3:
logger.error(f"CG coeff sum discrepancy exceed threshold: {coeff_sum} ({_n}, {_l}, {_ml}, {_ms})")
return transform
from scipy.constants import e as C_e, h as C_h, hbar as C_hbar, physical_constants
from plots.presentation.utils import setup_plot, save_current_fig
from system.hamiltonians.hamiltonians import load_hamiltonian
from system.hamiltonians.utils import plot_matrices, diagonalise_by_ml, diagonalise_for_n1n2
from system.states import States, Basis
from system.transformations.utils import load_transformation, transform_basis
from timer import timer
plot_fig17a = False
plot_fig17b = True
n = 51
with timer("Generating states"):
states = States(n, basis=Basis.N_L_ML_MS_RELEVANT)
# states = States(n, basis=Basis.N_L_ML_MS)
with timer("Loading Hamiltonian"):
mat_1, mat_1_zeeman, mat_2, mat_2_minus, mat_2_plus = load_hamiltonian(
f"{n}_rubidium87_relevant")
# mat_1, mat_1_zeeman, mat_2, mat_2_minus, mat_2_plus = load_hamiltonian(f"{n}_rubidium87")
with timer("Loading transformations"):
transform_1 = load_transformation(n, Basis.N_L_J_MJ_RELEVANT,
Basis.N_L_ML_MS_RELEVANT)
# transform_1 = load_transformation(n, Basis.N_L_J_MJ, Basis.N_L_ML_MS)
with timer("Applying transformation to nlmlms"):
mat_1 = transform_basis(mat_1, transform_1)
mat_2 = transform_basis(mat_2, transform_1)
with open(f"../../system/simulation/saved_simulations/{filename}", "rb") as f:
simulation: Simulation = pickle.load(f)
if not hasattr(simulation, 'rf_field'):
# Migrate old name
simulation.rf_field = simulation.rf_energy
print(simulation.dc_field)
print(simulation.rf_field)
print(simulation.rf_freq)
print(simulation.t)
systems: List[qutip.Qobj] = simulation.results.states
states = States(51, Basis.N1_N2_ML_MS).states
indices_to_keep = []
for i, (n1, n2, ml, ms) in enumerate(states):
if (n1 == 0 or n1 == 1) and ml >= 0 and ms > 0:
indices_to_keep.append(i)
indices_to_keep = sorted(indices_to_keep, key=lambda i: (states[i][0], states[i][2]))
states = np.array(states)[indices_to_keep]
state_mls = [state[2] for state in states]
max_ml = int(max(state_mls))
t_list = np.linspace(0, simulation.t, simulation.timesteps + 1)
system_mls = []
system_ml_averages = []
system_n1s = []
def plot(file_name):
with open(f"system/simulation/saved_simulations/{file_name}", "rb") as f:
simulation: Simulation = pickle.load(f)
print(simulation.dc_field)
print(simulation.rf_field)
print(simulation.rf_freq)
print(simulation.t)
systems: List[qutip.Qobj] = simulation.results.states
states = States(51, Basis.N1_N2_ML_MS).states
indices_to_keep = []
for i, (n1, n2, ml, ms) in enumerate(states):
if (n1 == 0 or n1 == 1) and ml >= 0 and ms > 0:
indices_to_keep.append(i)
indices_to_keep = sorted(indices_to_keep,
key=lambda i: (states[i][0], states[i][2]))
states = np.array(states)[indices_to_keep]
state_mls = [state[2] for state in states]
max_ml = int(max(state_mls))
t_list = np.linspace(0, simulation.t, simulation.timesteps + 1)
system_mls = []
system_ml_averages = []
system_n1s = []
for i, t in enumerate(t_list):
system = systems[i]
ml_average = 0
mls = np.zeros(max_ml + 1)
n1s = np.zeros(2)
system_populations = np.abs(system.data.toarray())**2
for j in range(simulation.states_count):
n1, n2, ml, ms = states[j]
state_population = system_populations[j]
if state_population > 0:
ml_average += state_population * ml
if n1 == 0:
mls[int(ml)] += state_population
n1s[int(n1)] += state_population
system_mls.append(mls)
system_ml_averages.append(ml_average)
system_n1s.append(n1s)
setup_plot()
setup_upmu()
fig, (ax1, ax2, ax3) = plt.subplots(3,
1,
figsize=(5, 6),
sharex='all',
gridspec_kw={
'hspace': 0.2,
'left': 0.15,
'right': 0.85,
'top': 0.96,
'bottom': 0.1,
})
rf_field = np.array(
[simulation.rf_field_calculator(t * 1000) for t in t_list])
rf_freq = np.array(
[simulation.rf_freq_calculator(t * 1000) for t in t_list])
e_rf_t, = ax1.plot(
t_list,
# np.sin(t_list / t_list[-1] * np.pi) * simulation.rf_field * 10,
np.cos(t_list * rf_freq * 1000 * 2 * np.pi) * rf_field *
10, # Factor of 10 to convert V/m to mV/cm
c="C0",
lw=3,
)
_ax1 = ax1.twinx()
dc_field = np.array(
[simulation.dc_field_calculator(t * 1000) for t in t_list])
e_dc_t, = _ax1.plot(
t_list,
dc_field / 100,
# (simulation.dc_field[0] + t_list / t_list[-1] * (simulation.dc_field[1] - simulation.dc_field[0])) / 100,
c="C1",
lw=3,
)
ax1.set_ylabel(r"$E_{\mathrm{RF}}$ [mV $\mathrm{cm}^{-1}$]")
_ax1.set_ylabel(r"$E_{\mathrm{d.c.}}$ [V $\mathrm{cm}^{-1}$]")
ax1.yaxis.label.set_color(e_rf_t.get_color())
_ax1.yaxis.label.set_color(e_dc_t.get_color())
ax1.tick_params(axis='y', colors=e_rf_t.get_color())
_ax1.tick_params(axis='y', colors=e_dc_t.get_color())
# lines = [e_rf_t, e_dc_t]
# ax1.legend(lines, [l.get_label() for l in lines])
system_mls = np.array(system_mls).T
system_mls = np.clip(system_mls, 1e-10, 1)
im = ax2.imshow(
system_mls,
aspect='auto',
cmap=plt.get_cmap('Blues'),
# cmap=COLORMAP, norm=NORM,
norm=LogNorm(vmin=1e-3, vmax=1, clip=True),
origin='lower',
extent=(0, t_list[-1], 0, max_ml))
# plt.colorbar(mappable=im, ax=ax2)
ax2.set_ylim((0, max_ml - 1))
ax2.set_ylabel("$m_l$, $n_1 = 0$")
system_n1s = np.array(system_n1s).T
# Initial state
ax3.plot(
t_list,
system_mls[3],
label=f"$c_{3}$, $n_1 = 0$",
lw=3,
)
# Circular state
ax3.plot(
t_list,
system_mls[-1],
label="$c_{n - 1}$",
lw=3,
)
print(f"c_n-1: {system_mls[-1][-1]}")
# n1 = 0
ax3.plot(
t_list,
system_n1s[0],
'--',
label="$\sum c$, $n_1 = 0$",
lw=3,
)
# n1 = 1
ax3.plot(
t_list,
system_n1s[1],
'--',
label="$\sum c$, $n_1 = 1$",
lw=3,
)
ax3.legend(fontsize='x-small')
ax3.set_ylim((0, 1))
ax3.set_ylabel("State Population")
ax3.set_xlabel(r"$t$ [$\upmu$s]")
save_current_fig(f'_simulation_{file_name}')
def generate_matrices(n: int, stark_map: arc.StarkMap, s=0.5):
"""
Generates matrices (described in detail below) used to construct a Hamiltonian in the n, l, j, mj basis.
This function is modelled after arc.calculations_atom_single.StarkMap.defineBasis
mat_1:
Atomic energies. Diagonal elements only.
See arc.alkali_atom_functions.AlkaliAtom.getEnergy()
mat_2:
E field couplings
See arc.calculations_atom_single.StarkMap._eFieldCouplingDivE()
mat_2_minus and mat_2_plus:
Couplings to an RF field
See calculate_coupling() below.
:param n:
:param stark_map:
:param s:
:return:
"""
global wignerPrecal
wignerPrecal = True
stark_map.eFieldCouplingSaved = _EFieldCoupling()
# states = States(n, Basis.N_L_J_MJ).states
states = States(n, Basis.N_L_J_MJ_RELEVANT).states
dimension = len(states)
print(f"Dimension: {dimension}", flush=True)
mat_1 = np.zeros((dimension, dimension), dtype=np.double)
mat_1_zeeman = np.zeros((dimension, dimension), dtype=np.double)
mat_2 = np.zeros((dimension, dimension), dtype=np.double)
mat_2_minus = np.zeros((dimension, dimension), dtype=np.double)
mat_2_plus = np.zeros((dimension, dimension), dtype=np.double)
pbar = tqdm(desc="Generating matrices", total=dimension**2)
for ii in range(dimension):
pbar.update(((dimension - ii) * 2 - 1))
n1, l1, j1, mj1 = states[ii]
### mat_1
atom_energy = stark_map.atom.getEnergy(
n=n1, l=l1, j=j1, s=stark_map.s) * C_e / C_h * 1e-9
mat_1[ii][ii] = atom_energy
zeeman_energy_shift = stark_map.atom.getZeemanEnergyShift(
l=l1, j=j1, mj=mj1, magneticFieldBz=1, s=stark_map.s) / C_h * 1e-9
mat_1_zeeman[ii][ii] = zeeman_energy_shift
for jj in range(ii + 1, dimension):
n2, l2, j2, mj2 = states[jj]
### mat_2
coupling_1 = stark_map._eFieldCouplingDivE(
n1=n1,
l1=l1,
j1=j1,
mj1=mj1,
n2=n2,
l2=l2,
j2=j2,
mj2=mj2,
s=stark_map.s) * 1.e-9 / C_h
# Scaling (as is also done in the arc package) so this can be multiplied by an E field (in V/m) to yield units of GHz.
mat_2[jj][ii] = coupling_1
mat_2[ii][jj] = coupling_1
pbar.close()
pbar = tqdm(desc="Generating matrices 2", total=dimension)
for ii in range(dimension):
for jj in range(dimension):
n1, l1, j1, mj1 = states[ii]
n2, l2, j2, mj2 = states[jj]
### mat_2_minus and mat_2_plus
coupling_2 = calculate_coupling(
stark_map,
n1,
l1,
j1,
mj1,
n2,
l2,
j2,
mj2,
-1,
s,
) * C_e * C_a_0 * 1.e-9 / C_hbar
mat_2_minus[ii][jj] = coupling_2
coupling_3 = calculate_coupling(
stark_map,
n1,
l1,
j1,
mj1,
n2,
l2,
j2,
mj2,
1,
s,
) * C_e * C_a_0 * 1.e-9 / C_hbar
mat_2_plus[ii][jj] = coupling_3
pbar.update(1)
pbar.close()
return states, (mat_1, mat_1_zeeman, mat_2, mat_2_minus, mat_2_plus)
from matplotlib.cm import ScalarMappable
from matplotlib.colors import ListedColormap, Normalize
from tqdm import tqdm
from scipy.constants import e as C_e, h as C_h, hbar as C_hbar, physical_constants
from system.hamiltonians.hamiltonians import load_hamiltonian
from system.hamiltonians.utils import plot_matrices
from system.states import States, Basis
from system.transformations.utils import load_transformation, transform_basis
from timer import timer
n = 56
with timer("Generating states"):
states_n_l_ml_ms = States(n, basis=Basis.N_L_ML_MS).states
states = States(n, basis=Basis.N1_N2_ML_MS).states
with timer("Loading Hamiltonian"):
# mat_1, mat_2, mat_2_minus, mat_2_plus = load_hamiltonian(f"{n}_rubidium")
mat_1, mat_2, mat_2_minus, mat_2_plus = load_hamiltonian(f"{n}_rubidium87")
# mat_1, mat_2, mat_2_minus, mat_2_plus = load_hamiltonian(f"{n}_hydrogen")
mat_2_combination = mat_2_plus + mat_2_minus # Units of a0 e
# mat_2_combination = mat_2_plus # Units of a0 e
mat_2_combination *= C_e * physical_constants["Bohr radius"][0] / C_hbar
# Conversion from atomic units for dipole matrix elements to a Rabi freq in Hz
mat_2_combination *= 1e-9 # Convert Hz to GHz
# plot_matrices([mat_1, mat_2])
with timer("Loading transformations"):
# _raw_dc_calculator = simulation.get_calculator((270, 230))
# simulation.dc_field_calculator = lambda _t: _raw_dc_calculator(_t).round(1)
# _raw_dc_calculator = simulation.get_calculator((270, 210))
# simulation.dc_field_calculator = lambda _t: _raw_dc_calculator(_t).round(1)
#
# simulation.rf_freq_calculator = simulation.get_calculator(230e6 / 1e9)
# simulation.rf_field_calculator = lambda t: 3 * np.sin(np.pi * t / 1000 / simulation.t)
print(simulation.dc_field)
print(simulation.rf_field)
print(simulation.rf_freq)
print(simulation.t)
systems: List[qutip.Qobj] = simulation.results.states
states = States(simulation.n, Basis.N1_N2_ML_MS).states
indices_to_keep = []
for i, (n1, n2, ml, ms) in enumerate(states):
if (n1 == 0 or n1 == 1) and ml >= 0 and ms > 0:
indices_to_keep.append(i)
indices_to_keep = sorted(indices_to_keep, key=lambda i: (states[i][0], states[i][2]))
states = np.array(states)[indices_to_keep]
state_mls = [state[2] for state in states]
max_ml = int(max(state_mls))
t_list = np.linspace(0, simulation.t, simulation.timesteps + 1)
system_mls = []
system_ml_averages = []
system_n1s = []