# -*- coding: utf-8 -*-
"""
qspec._simulate_cpp
===================
Created on 09.01.2022
@author: Patrick Mueller
Classes and methods for the 'simulate' module using the Python/C++ interface.
"""
import numpy as np
from matplotlib import cm
import matplotlib.pyplot as plt
from qspec._types import *
from qspec._cpp import *
from qspec import tools
from qspec import get_f, get_m
import qspec.algebra as al
__all__ = ['Polarization', 'Laser', 'Environment', 'construct_electronic_state', 'construct_hyperfine_state', 'State',
'DecayMap', 'Atom', 'Interaction']
def sr_generate_y(denominator: np.ndarray, f_theta: np.ndarray, f_phi: np.ndarray,
counts: np.ndarray, shape: np.ndarray):
"""
:param denominator: The denominator of the scattering rate.
:param f_theta: The numerator with the 'theta-polarization'.
:param f_phi: The numerator with the 'phi-polarization'.
:param counts: The number of summands.
:param shape: The shape of y.
:returns: The scattering rate.
"""
y = np.zeros((shape[0] * shape[1], ), dtype=float) # Allocate memory.
denominator_p = denominator.ctypes.data_as(c_complex_p) # Get all pointers to the first elements of the arrays.
f_theta_p = f_theta.ctypes.data_as(c_complex_p)
f_phi_p = f_phi.ctypes.data_as(c_complex_p)
counts_p = counts.ctypes.data_as(c_int_p)
shape_p = shape.ctypes.data_as(c_int_p)
y_p = y.ctypes.data_as(c_double_p)
dll.sr_generate_y(denominator_p, f_theta_p, f_phi_p, counts_p, shape_p, y_p) # Modify y "in-place" with C++.
return y
def _process_q_axis(q_axis: array_like) -> ndarray:
"""
Preprocess the quantization axis.
:param q_axis: The quantization axis. Must be an integer in {0, 1, 2} or a 3d-vector.
:returns: The quantization axis as a 3d-vector.
"""
q_axis = np.asarray(q_axis, dtype=float)
if not q_axis.shape:
if q_axis % 1 != 0:
raise ValueError('q_axis must be an integer or a 3d-vector.')
elif int(q_axis) not in {0, 1, 2}:
raise ValueError('q_axis must be in {0, 1, 2}.')
else:
q_axis = tools.unit_vector(int(q_axis), 3)
elif q_axis.shape != (3,):
raise ValueError('q_axis must be an integer or a 3d-vector.')
return q_axis
[docs]class Polarization:
"""
Class representing a polarization state of light. The property 'x' holds the polarization in cartesian coordinates.
The property 'q' holds the polarization in spherical coordinates ( sigma-, pi, sigma+ )
with respect to the chosen quantization axis.
"""
def __init__(self, vec: array_iter = None, q_axis: array_like = 2, vec_as_q: bool = True, instance=None):
"""
:param vec: The polarization vector. I.e. the amplitude of the electromagnetic wave. So to specify, e.g.,
1/3 of pi and 2/3 sigma+ light for a given 'q_axis', vec must be ( 0, sqrt(1/3), sqrt(2/3) ).
The default is linear polarization in z-direction, such that x = (0, 0, 1) and q = (0, 1, 0).
:param q_axis: The quantization axis. Must be an integer in {0, 1, 2} or a 3d-vector. The default is 2 (z-axis).
:param vec_as_q: Whether 'vec' is given as ( sigma-, pi, sigma+ ) (True) or in cartesian coordinates (False).
:param instance: A pointer to an existing Polarization instance.
If this is specified, the other parameters are omitted.
"""
self.instance = instance
if self.instance is None:
self.instance = dll.polarization_construct()
if vec is None:
vec = tools.unit_vector(1, 3, dtype=complex)
vec = np.asarray(vec, dtype=complex)
if vec.shape != (3,):
raise ValueError('vec must be a 3d-vector.')
q_axis = _process_q_axis(q_axis)
dll.polarization_init(self.instance, vec, q_axis, vec_as_q)
def __del__(self):
dll.polarization_destruct(self.instance)
[docs] def def_q_axis(self, q_axis: array_like = 2, q_fixed: bool = False):
"""
Defines the quantization axis. This changes either 'x' or 'q', depending on 'q_fixed'.
:param q_axis: The quantization axis. Must be an integer in {0, 1, 2} or a 3d-vector. The default is 2 (z-axis).
:param q_fixed: Whether 'q' should stay the same with the new quantization axis (True) or 'x' (False).
"""
q_axis = _process_q_axis(q_axis)
dll.polarization_def_q_axis(self.instance, q_axis, q_fixed)
@property
def q_axis(self) -> ndarray:
"""
:returns: The complex polarization in spherical coordinates ( sigma-, pi, sigma+ ).
"""
return dll.polarization_get_q_axis(self.instance)
@property
def x(self) -> ndarray:
"""
:returns: The complex polarization in cartesian coordinates ( x, y, z ).
"""
return dll.polarization_get_x(self.instance)
@property
def q(self) -> ndarray:
"""
:returns: The complex polarization in spherical coordinates ( sigma-, pi, sigma+ ).
"""
return dll.polarization_get_q(self.instance)
[docs]class Laser:
"""
Class representing a laser.
"""
def __init__(self, freq: scalar, intensity: scalar = 1., polarization: Polarization = None,
k: array_like = None, instance=None):
"""
:param freq: The frequency of the laser (MHz).
:param intensity: The intensity of the laser (uW / mm**2 = W / m**2).
:param polarization: The polarization of the laser.
:param k: The direction of the laser.
:param instance: A pointer to an existing Laser instance.
If this is specified, the other parameters are omitted.
"""
self.instance = instance
if self.instance is None:
self.instance = dll.laser_construct()
self._polarization = polarization
if self.polarization is None:
self._polarization = Polarization()
dll.laser_init(self.instance, c_double(freq), c_double(intensity), polarization.instance)
if k is None:
k = tools.unit_vector(0, 3)
self.k = k
else:
self._polarization = Polarization(instance=dll.laser_get_polarization(self.instance))
def __del__(self):
dll.laser_destruct(self.instance)
@property
def freq(self):
"""
:returns: The frequency of the laser.
"""
return dll.laser_get_freq(self.instance)
@freq.setter
def freq(self, value: scalar):
"""
:param value: The new frequency of the laser.
:returns: None.
"""
dll.laser_set_freq(self.instance, c_double(value))
@property
def intensity(self):
"""
:returns: The intensity of the laser.
"""
return dll.laser_get_intensity(self.instance)
@intensity.setter
def intensity(self, value: scalar):
"""
:param value: The new intensity of the laser.
:returns: None.
"""
dll.laser_set_intensity(self.instance, c_double(value))
@property
def polarization(self):
"""
:returns: The polarization of the laser.
"""
return self._polarization
@polarization.setter
def polarization(self, value: Polarization):
"""
:param value: The new polarization of the laser.
:return: None.
"""
self._polarization = value
dll.laser_set_polarization(self.instance, value.instance)
@property
def k(self):
"""
:returns: The direction of the laser. The default direction is ( 1, 0, 0 ).
"""
return dll.laser_get_k(self.instance)
@k.setter
def k(self, value: array_like):
"""
:param value: The new direction of the laser.
:return: None.
"""
_value = np.asarray(value, dtype=float)
if _value.size != 3:
raise ValueError('Interaction.k must be a 3d-vector, but has shape {}.'.format(_value.shape))
dll.laser_set_k(self.instance, _value.flatten())
def _process_hyper_const(hyper_const: array_like) -> ndarray:
"""
Preprocess the hyperfine-structure constants.
:param hyper_const: The hyperfine-structure constants. Currently, constants up to the electric quadrupole order are
supported (A, B). If 'hyper_const' is a scalar,
it is assumed to be the constant A and the other orders are 0 (MHz).
:returns: The hyperfine-structure constants as a 3d-vector.
"""
if hyper_const is None or not hyper_const:
hyper_const = [0., 0., 0.]
elif not np.asarray(hyper_const, dtype=float).shape:
hyper_const = [float(hyper_const), 0., 0.]
hyper_const = list(hyper_const)
while len(hyper_const) < 3:
hyper_const.append(0.)
return np.asarray(hyper_const, dtype=float)[:3]
# noinspection PyPep8Naming
[docs]class Environment:
"""
Class representing an electromagnetic environment.
"""
def __init__(self, E: array_like = None, B: array_like = None, instance=None):
self.instance = instance
if self.instance is None:
self.instance = dll.environment_construct()
self.E = E
self.B = B
def __del__(self):
dll.environment_destruct(self.instance)
@property
def E(self):
return dll.environment_get_E(self.instance) * dll.environment_get_e_E(self.instance)
@E.setter
def E(self, value: array_like):
if value is None:
dll.environment_set_E(self.instance, np.asarray([1, 0, 0], dtype=float))
dll.environment_set_E_double(self.instance, c_double(0.))
else:
value = np.asarray(value, dtype=float)
if not value.shape:
dll.environment_set_E_double(self.instance, value)
elif value.shape == (3, ):
dll.environment_set_E(self.instance, value)
else:
raise ValueError('E must be a scalar, 3d-vector or None, but has shape {}'.format(value.shape))
@property
def B(self):
return dll.environment_get_B(self.instance) * dll.environment_get_e_B(self.instance)
@B.setter
def B(self, value: array_like):
if value is None:
dll.environment_set_B(self.instance, np.array([0, 0, 1], dtype=float))
dll.environment_set_B_double(self.instance, c_double(0.))
else:
value = np.asarray(value, dtype=float)
if not value.shape:
dll.environment_set_B_double(self.instance, value)
elif value.shape == (3, ):
dll.environment_set_B(self.instance, value)
else:
raise ValueError('B must be a scalar, 3d-vector or None, but has shape {}'.format(value.shape))
[docs]class State:
"""
Class representing an atomic quantum state :math:`|(\\mathrm{label})SLJIFm\\rangle`.
"""
def __init__(self, freq_j: scalar, s: scalar, l: scalar, j: scalar, i: scalar, f: scalar, m: scalar,
hyper_const: array_like = None, g: scalar = 0, label: str = None, instance=None):
"""
:param freq_j: The energetic position of the state without the hyperfine structure or the environment (MHz).
:param s: The electron spin quantum number S.
:param l: The electronic angular momentum quantum number L.
:param j: The electronic total angular momentum quantum number J.
:param i: The nuclear spin quantum number I.
:param f: The total angular momentum quantum number F.
:param m: The B-field-axis component quantum number m of the total angular momentum.
:param hyper_const: The hyperfine-structure constants. Currently, constants up to the electric quadrupole order
are supported (A, B). If 'hyper_const' is a scalar, it is assumed to be the constant A
and the other orders are 0 (MHz).
:param g: The nuclear g-factor.
:param label: The label of the state. The label is used to link states via decay maps.
:param instance: A pointer to an existing State instance.
If this is specified, the other parameters are omitted.
"""
self.instance = instance
if self.instance is None:
tools.check_half_integer(s, l, j, i, f, m)
hyper_const = _process_hyper_const(hyper_const)
if label is None:
label = '{}({}, {}, {})'.format(int(np.around(freq_j, decimals=0)), s, l, j)
self.instance = dll.state_construct()
dll.state_init(self.instance, c_double(freq_j), c_double(s), c_double(l), c_double(j), c_double(i),
c_double(f), c_double(m), hyper_const, c_double(g), c_char_p(bytes(label, 'utf-8')))
def __del__(self):
dll.state_destruct(self.instance)
def __repr__(self):
return '{}({})'.format(self.label, ('{}, ' * 6)[:-2]) \
.format(*[tools.half_integer_to_str(qn, '/') for qn in [self.s, self.l, self.j, self.i, self.f, self.m]])
[docs] def update(self, environment: Environment = None):
"""
Update the shifted frequency of the state.
:returns: None.
"""
if environment is None:
dll.state_update(self.instance)
else:
dll.state_update_env(self.instance, environment.instance)
[docs] def get_shift(self):
"""
:returns: The difference between the shifted frequency of the hyperfine-structure
and the frequency of the fine-structure state.
"""
return dll.state_get_shift(self.instance)
@property
def freq_j(self):
"""
:returns: The frequency of the fine-structure state.
"""
return dll.state_get_freq_j(self.instance)
@freq_j.setter
def freq_j(self, value: scalar):
dll.state_set_freq_j(self.instance, c_double(value))
@property
def freq(self):
"""
:returns: The shifted frequency of the hyperfine-structure state.
"""
return dll.state_get_freq(self.instance)
@property
def s(self):
"""
:returns: The electron spin quantum number S.
"""
return dll.state_get_s(self.instance)
@property
def l(self):
"""
:returns: The electronic angular momentum quantum number L.
"""
return dll.state_get_l(self.instance)
@property
def j(self):
"""
:returns: The electronic total angular momentum quantum number J.
"""
return dll.state_get_j(self.instance)
@property
def i(self):
"""
:returns: The nuclear spin quantum number I.
"""
return dll.state_get_i(self.instance)
@property
def f(self):
"""
:returns: The total angular momentum quantum number F.
"""
return dll.state_get_f(self.instance)
@property
def m(self):
"""
:returns: The B-field-axis component quantum number m of the total angular momentum.
"""
return dll.state_get_m(self.instance)
@property
def hyper_const(self):
"""
:returns: The hyperfine-structure constants as a 3d-vector.
"""
return dll.state_get_hyper_const(self.instance)
@hyper_const.setter
def hyper_const(self, value: array_like):
"""
:param value: The new hyperfine-structure constants. Currently, constants up to the electric quadrupole order
are supported (A, B). If 'hyper_const' is a scalar, it is assumed to be the constant A
and the other orders are 0 (MHz).
:returns: None.
"""
value = _process_hyper_const(value)
dll.state_get_hyper_const(self.instance, value)
@property
def g(self):
"""
:returns: The nuclear g-factor.
"""
return dll.state_get_g(self.instance)
@g.setter
def g(self, value: scalar):
"""
:param value: The new nuclear g-factor.
:returns: None.
"""
dll.state_set_g(self.instance, c_double(value))
@property
def label(self):
"""
:returns: The label of the state. The label is used to link states via decay maps.
"""
return dll.state_get_label(self.instance).decode('utf-8')
@label.setter
def label(self, value: str):
"""
:param value: The label of the state. The label is used to link states via decay maps.
:returns: None.
"""
dll.state_set_label(self.instance, c_char_p(bytes(value, 'utf-8')))
[docs]def construct_electronic_state(freq_0: scalar, s: scalar, l: scalar, j: scalar, i: scalar = 0,
hyper_const: Iterable[scalar] = None, g: scalar = 0, label: str = None):
"""
Creates all substates of a fine-structure state using a common label.
:param freq_0: The energetic position of the state without the hyperfine structure or the magnetic field (MHz).
:param s: The electron spin quantum number S.
:param l: The electronic angular momentum quantum number L.
:param j: The electronic total angular momentum quantum number J.
:param i: The nuclear spin quantum number I.
:param hyper_const: The hyperfine-structure constants. Currently, constants up to the electric quadrupole order are
supported (A, B). If 'hyper_const' is a scalar,
it is assumed to be the constant A and the other orders are 0 (MHz).
:param g: The nuclear g-factor.
:param label: The label of the states. The labels are used to link states via decay maps.
:returns: A list of the created states.
"""
f = get_f(i, j)
m = [get_m(_f) for _f in f]
fm = [(_f, _m) for _f, m_f in zip(f, m) for _m in m_f]
return [State(freq_0, s, l, j, i, _f, _m, hyper_const=hyper_const, g=g, label=label) for (_f, _m) in fm]
[docs]def construct_hyperfine_state(freq_0: scalar, s: scalar, l: scalar, j: scalar, i: scalar, f: scalar,
hyper_const: Iterable[scalar] = None, g: scalar = 0, label: str = None):
"""
Creates all substates of a fine-structure state using a common label.
:param freq_0: The energetic position of the state without the hyperfine structure or the magnetic field (MHz).
:param s: The electron spin quantum number S.
:param l: The electronic angular momentum quantum number L.
:param j: The electronic total angular momentum quantum number J.
:param i: The nuclear spin quantum number I.
:param f: The hyperfine structure total angular momentum quantum number F.
:param hyper_const: The hyperfine-structure constants. Currently, constants up to the electric quadrupole order are
supported (A, B). If 'hyper_const' is a scalar,
it is assumed to be the constant A and the other orders are 0 (MHz).
:param g: The nuclear g-factor.
:param label: The label of the states. The labels are used to link states via decay maps.
:returns: A list of the created states.
"""
return [State(freq_0, s, l, j, i, f, _m, hyper_const=hyper_const, g=g, label=label) for _m in get_m(f)]
[docs]class DecayMap:
"""
Class linking sets of atomic states via Einstein-A coefficients.
"""
def __init__(self, labels: Iterable[tuple] = None, a: Iterable[scalar] = None, instance=None):
"""
:param labels: An iterable of label pairs, corresponding to atomic states which get connected.
:param a: An Iterable of Einstein-A coefficients (MHz).
:param instance: A pointer to an existing DecayMap instance.
If this is specified, the other parameters are omitted.
"""
self.instance = instance
if self.instance is None:
self.instance = dll.decaymap_construct()
if labels is None:
labels = []
if a is None:
a = []
self._labels = list(labels)
for (s0, s1), _a in zip(self._labels, a):
dll.decaymap_add_decay(self.instance, c_char_p(bytes(s0, 'utf-8')), c_char_p(bytes(s1, 'utf-8')),
c_double(float(_a)))
else:
self._labels = self._get_labels()
def __del__(self):
dll.decaymap_destruct(self.instance)
def _get_labels(self):
"""
:returns: The labels used in the C++ class.
"""
return [(dll.decaymap_get_label(self.instance, 0, i).decode('utf-8'),
dll.decaymap_get_label(self.instance, 1, i).decode('utf-8')) for i in range(self.size)]
@property
def labels(self):
"""
:returns: The list of label pairs, corresponding to atomic states which get connected.
"""
return self._labels
@property
def a(self):
"""
:returns: The list of Einstein-A coefficients.
"""
vector_d_p = np.ctypeslib.ndpointer(dtype=float, shape=(self.size, ))
set_restype(dll.decaymap_get_a, vector_d_p)
return dll.decaymap_get_a(self.instance).tolist()
@property
def size(self):
"""
:returns: The number of linked sets of atomic states.
"""
return dll.decaymap_get_size(self.instance)
[docs] def get(self, label_0: str, label_1: str):
"""
:returns: The number of linked sets of atomic states.
"""
return dll.decaymap_get_item(
self.instance, c_char_p(bytes(label_0, 'utf-8')), c_char_p(bytes(label_1, 'utf-8')))
def _gen_label_map(atom):
"""
:param atom: The atom.
:returns: A dictionary with state labels as keys
and an array of the indices of the states with the labels as values.
"""
if isinstance(atom, int):
return {'States 0 - {}'.format(atom): np.arange(atom, dtype=int)}
all_labels = [s.label for s in atom]
labels = []
for s in all_labels:
if s not in labels:
labels.append(s)
label_map = {s: np.array([i for i, _s in enumerate(all_labels) if _s == s]) for s in labels}
return label_map
[docs]class Atom:
"""
Class representing an Atom and its inner structure.
"""
def __init__(self, states: Iterable[State] = None, decay_map: DecayMap = None, mass: scalar = 0, instance=None):
"""
:param states: The states of the atom.
:param decay_map: The decay map which connects the atomic states.
:param mass: The mass of the atom (u).
:param instance: A pointer to an existing Atom instance.
If this is specified, the other parameters are omitted.
"""
self.instance = instance
if self.instance is None:
self.instance = dll.atom_construct()
if states is None:
states = []
if decay_map is None:
decay_map = DecayMap()
self.states = list(states)
self.decay_map = decay_map
self.mass = mass
self.label_map = None
self.update()
else:
self._states = None
def __del__(self):
dll.atom_destruct(self.instance)
def __iter__(self):
for state in self.states:
yield state
def __getitem__(self, key: int):
return self.states[key]
[docs] def update(self):
"""
Update the atom.
:returns: None.
"""
dll.atom_update(self.instance)
self.label_map = _gen_label_map(self)
@property
def states(self):
"""
:returns: The states of the atom.
"""
return self._states
@states.setter
def states(self, value: Iterable[State]):
"""
:param value: The new states of the atom.
:returns: None.
"""
dll.atom_clear_states(self.instance)
self._states = list(value)
for s in self._states:
dll.atom_add_state(self.instance, s.instance)
@property
def decay_map(self):
"""
:returns: The decay map which connects the atomic states.
"""
return self._decay_map
@decay_map.setter
def decay_map(self, value: DecayMap):
"""
:param value: The new decay map which connects the atomic states.
:returns: None.
"""
self._decay_map = value
if self._decay_map is None:
self._decay_map = DecayMap()
dll.atom_set_decay_map(self.instance, value.instance)
@property
def mass(self):
"""
:returns: The mass of the atom (u).
"""
return dll.atom_get_mass(self.instance)
@mass.setter
def mass(self, value: scalar):
"""
:param value: The new mass of the atom (u).
:returns: None.
"""
dll.atom_set_mass(self.instance, c_double(value))
@property
def size(self):
"""
:returns: The number of states of the atom.
"""
return dll.atom_get_size(self.instance)
@property
def gs(self) -> ndarray:
"""
:returns: The indices of the ground states.
"""
vector_i_p = np.ctypeslib.ndpointer(dtype=c_size_t, shape=(dll.atom_get_gs_size(self.instance), ))
set_restype(dll.atom_get_gs, vector_i_p)
return dll.atom_get_gs(self.instance)
@property
def dipoles(self):
"""
:returns: The dipole strengths between the atomic states in the 3 basis-components
of the spherical vector basis ( sigma-, pi, sigma+ ). This can be used to calculate Rabi-frequencies by
multiplying it with the square-root of a laser intensity in the corresponding polarization.
The resulting array has shape (3, size, size).
"""
return np.array([np.ctypeslib.as_array(dll.atom_get_m_dipole(self.instance, c_size_t(i)),
(self.size, self.size)).T for i in range(3)])
@property
def l0(self):
a = dll.atom_get_L0(self.instance)
return np.ctypeslib.as_array(a, (self.size, self.size)).T
@property
def l1(self):
a = dll.atom_get_L1(self.instance)
return np.ctypeslib.as_array(a, (self.size, self.size)).T
[docs] def get_y0(self, ground_state_labels: Union[Iterable[str], str] = None) -> np.ndarray:
"""
:param ground_state_labels: An Iterable of labels belonging to ground states.
:returns: The initial population of the atom.
"""
y0 = np.zeros(self.size)
if ground_state_labels is None:
ground_state_labels = [self.states[0].label]
indices = np.array([i for i, state in enumerate(self) if state.label in ground_state_labels])
y0[indices] = 1 / indices.size
return y0
[docs] def get_state_indexes(self, labels: Union[Iterable[str], str] = None,
f: Union[Iterable[scalar], scalar] = None) -> np.ndarray:
"""
:param labels: The labels of the states whose indexes are to be returned.
:param f: The F quantum numbers whose indexes are to be returned.
:returns: The indexes corresponding to the specified labels and F quantum numbers.
"""
if labels is None:
labels = set(s.label for s in self.states)
if f is None:
f = set(s.f for s in self.states)
try:
f = set(f)
except TypeError:
f = {f}
return np.array([i for i, s in enumerate(self.states) if s.label in labels and s.f in f], dtype=int)
[docs] def scattering_rate(self, rho: array_like, theta: array_like = None, phi: array_like = None,
as_density_matrix: bool = True, i: array_like = None, j: array_like = None, axis: int = 1):
"""
:param rho: The state vector (density matrix) of the atom. Must have the same size as the atom
along the specified 'axis' (and 'axis' + 1).
:param theta: The elevation angle of detection.
:param phi: The azimuthal angle of detection.
:param as_density_matrix: Whether 'rho' is a state vector or a density matrix.
:param i: The initially excited state indexes to consider for spontaneous decay.
If None, all states are considered.
:param j: The final decayed state indexes to consider for spontaneous decay. If None, all states are considered.
:param axis: The axis along which the population is aligned in 'rho'.
:returns: The scattering rate of the atom given the population 'rho' (MHz or Events / s).
:raises ValueError: 'rho' must have the same size as the atom along the specified 'axis'.
"""
rho = np.asarray(rho)
if i is None:
i = np.arange(self.size, dtype=int)
else:
i = np.array(i).flatten()
if j is None:
j = np.arange(self.size, dtype=int)
else:
j = np.array(j).flatten()
if axis < 0:
axis += len(rho.shape)
l0 = np.array([[1. if _i in i and _j in j else 0. for _i in range(self.size)] for _j in range(self.size)])
l0 *= self.l0
if theta is None and phi is None:
if as_density_matrix:
rho = np.diagonal(rho, axis1=axis, axis2=axis + 1)
if len(rho.shape) > axis + 1:
axes = list(range(len(rho.shape)))
axes[axis + 1:] = axes[axis:-1]
axes[axis] = len(rho.shape) - 1
rho = np.transpose(rho, axes=axes)
axes = [ax for ax in range(axis)]
if axes:
l0 = np.expand_dims(self.l0, axis=axes)
axes = [axis + ax + 2 for ax in range(len(rho.shape) - axis - 1)]
if axes:
l0 = np.expand_dims(l0, axis=axes)
sr = tools.transform(l0, rho, axis=axis)
return np.sum(sr, axis=axis)
elif theta is None or phi is None:
raise ValueError('\'theta\' and \'phi\' must either both be specified or both be None.')
if not as_density_matrix:
rho = tools.vector_to_diag_matrix(rho, axis=axis)
a_cart = [[al.a_dipole_cart(self.states[_j].i, self.states[_j].j, self.states[_j].f, self.states[_j].m,
self.states[_i].j, self.states[_i].f, self.states[_i].m)
* np.sqrt(2 * self.states[_j].i + 1) * np.sqrt(2 * self.states[_j].j + 1)
* np.sqrt(self.decay_map.get(self.states[_j].label, self.states[_i].label))
if l0[_j, _i] else np.zeros(3, dtype=complex)
if _i in i and _j in j else 0. for _i in range(self.size)] for _j in range(self.size)]
e_theta = tools.e_theta(theta, phi)
e_phi = tools.e_phi(theta, phi)
c_theta = np.array([[np.sum(e_theta * _a_cart) for _a_cart in a_cart_list]
for a_cart_list in a_cart])
c_phi = np.array([[np.sum(e_phi * _a_cart) for _a_cart in a_cart_list]
for a_cart_list in a_cart])
ct_theta = np.array([[np.sum(e_theta * np.conj(_a_cart))
for _a_cart in a_cart_list] for a_cart_list in a_cart])
ct_phi = np.array([[np.sum(e_phi * np.conj(_a_cart))
for _a_cart in a_cart_list] for a_cart_list in a_cart])
axes = [ax for ax in range(axis)]
if axes:
c_theta = np.expand_dims(c_theta, axis=axes)
c_phi = np.expand_dims(c_phi, axis=axes)
ct_theta = np.expand_dims(ct_theta, axis=axes)
ct_phi = np.expand_dims(ct_phi, axis=axes)
axes = [axis + ax + 2 for ax in range(len(rho.shape) - axis - 2)]
if axes:
c_theta = np.expand_dims(c_theta, axis=axes)
c_phi = np.expand_dims(c_phi, axis=axes)
ct_theta = np.expand_dims(ct_theta, axis=axes)
ct_phi = np.expand_dims(ct_phi, axis=axes)
sr = (np.sum([np.expand_dims(tools.get_subarray(c_theta, k, axis), axis=axis + 1)
* np.expand_dims(tools.get_subarray(rho, k, axis + 1), axis=axis)
for k in range(self.size)], axis=0)
* np.sum([np.expand_dims(tools.get_subarray(ct_theta, k, axis + 1), axis=axis)
* np.expand_dims(tools.get_subarray(rho, k, axis), axis=axis + 1)
for k in range(self.size)], axis=0))
sr += (np.sum([np.expand_dims(tools.get_subarray(c_phi, k, axis), axis=axis + 1)
* np.expand_dims(tools.get_subarray(rho, k, axis + 1), axis=axis)
for k in range(self.size)], axis=0)
* np.sum([np.expand_dims(tools.get_subarray(ct_phi, k, axis + 1), axis=axis)
* np.expand_dims(tools.get_subarray(rho, k, axis), axis=axis + 1)
for k in range(self.size)], axis=0))
return 3 / (8 * np.pi) * np.sum(np.sum(sr, axis=axis), axis=axis).real
[docs] def plot(self, indices: array_like = None, draw_bounds: bool = False, show: bool = True):
"""
Plot a term scheme of the atom.
:param indices: The indices of the states to be drawn. If None, all states are drawn.
:param draw_bounds: Whether to draw the upper vertical bounds of the states.
:param show: Whether to show the plot.
:returns: The x and y positions of the states as well as the distance constant d.
"""
if indices is None:
indices = np.argsort([state.freq for state in self.states])
d = 2
y_i = 0
y_dict = {}
m_max = max(s.m for s in self.states)
ret_x = {}
ret_y = {}
for i in indices:
s = self.states[i]
key = (s.label, s.s, s.l, s.j, s.i, s.f)
if key not in y_dict.keys():
if key[:-1] not in [key_1[:-1] for key_1 in y_dict.keys()]:
y_i += 3
y_dict[key] = y_i * d
y_i += 1
if draw_bounds:
plt.hlines([y_dict[key] + d * 0.05, y_dict[key] + d * 0.95][1:],
xmin=-m_max - 0.5, xmax=m_max + 0.5, ls='--', colors='grey')
x = np.array([s.m - 0.45, s.m + 0.45]) # + x_off[key[:-1]]
y = np.array([y_dict[key], y_dict[key]])
ret_x[i] = x
ret_y[i] = y
plt.plot(x, y, 'k-')
x_ticks = np.linspace(-m_max, m_max, int(2 * m_max + 1), dtype=float)
plt.xticks(x_ticks)
y_ticks = np.array([[_y[1], (_y[0][0], _y[0][-1])] for _y in y_dict.items()], dtype=object)
order = np.argsort(y_ticks, axis=0)[:, 0]
y_ticks = y_ticks[order]
plt.yticks(y_ticks[:, 0].astype(float), [str(y_l) for y_l in y_ticks[:, 1]])
plt.xlabel(r'$m$')
plt.ylabel(r'$(\mathrm{label}, F)$')
if show:
plt.tight_layout()
plt.show()
return ret_x, ret_y, d
def _cast_delta(delta: array_like, m: Optional[int], size: int) -> ndarray:
"""
:param delta: An array of frequency shifts for the laser(s). 'delta' must be a scalar or a 1d- or 2d-array
with shapes (., ) or (., #lasers), respectively.
:param m: The index of the shifted laser. If delta is a 2d-array, 'm' ist omitted.
:param size: The number of available lasers.
:returns: An array of vectors with size 'size' containing frequency shifts for the lasers.
"""
if delta is None:
return np.zeros((1, size))
delta = np.array(delta, dtype=float, order='C')
if len(delta.shape) != 2 and not -size <= m < size:
raise IndexError('Laser index \'m\' is out of bounds. Must be {} <= m < {} or None but is {}.'
.format(-size, size, m))
error = False
if len(delta.shape) > 2:
error = True
elif len(delta.shape) == 0:
if m is None:
delta = np.full((1, size), delta, dtype=float)
else:
delta = np.expand_dims(tools.unit_vector(m, size, dtype=float) * delta, axis=0)
elif len(delta.shape) == 1:
if m is None:
delta = delta[:, None] + np.expand_dims(np.zeros(size), axis=0)
else:
delta = tools.unit_vector(np.full(delta.size, m, dtype=int), size, dtype=float) \
* np.expand_dims(delta, axis=1)
elif delta.shape[1] != size:
error = True
if error:
raise ValueError('\'delta\' must be a scalar or a 1d- or 2d-array with shapes '
'(., ) or (., #lasers), respectively.')
return delta
def _cast_y0(y0: Optional[array_like], i_solver: int, atom: Atom):
"""
:param y0: The initial state of an ensemble of atoms. Depending on the solver, this must be (...).
:param i_solver: The index of the solver.
:param atom: The atom.
:returns: The correctly shaped 'y0' for the chosen solver and its C type.
"""
size = atom.size
gs = atom.gs
if i_solver == 0: # Rate equations.
if y0 is None:
y0 = np.zeros(size, dtype=float)
y0[gs] = 1 / gs.size
return y0, c_double_p
y0 = np.array(y0, dtype=float, order='C')
if not y0.shape or y0.shape[-1] != size:
raise ValueError('\'y0\' must have size {} in the last axis but has shape {}.'.format(size, y0.shape))
y0 /= np.expand_dims(np.sum(y0, axis=-1), axis=-1)
return y0, c_double_p
elif i_solver == 1: # Schroedinger equation.
if y0 is None:
y0 = np.zeros(size, dtype=complex)
y0[gs[0]] = 1
return y0, c_complex_p
y0 = np.array(y0, dtype=complex, order='C')
if not y0.shape or y0.shape[-1] != size:
raise ValueError('\'y0\' must have size {} in the last axis but has shape {}.'.format(size, y0.shape))
y0 /= np.expand_dims(tools.absolute_complex(y0, axis=-1), axis=-1)
return y0, c_complex_p
elif i_solver == 2: # Master equation.
if y0 is None:
y0 = np.zeros(size, dtype=complex)
y0[gs] = 1 / gs.size
return np.diag(y0), c_complex_p
y0 = np.array(y0, dtype=complex, order='C')
if not y0.shape or (len(y0.shape) == 1 and y0.shape[-1] != size) \
or (len(y0.shape) > 1 and y0.shape[-2:] != (size, size)):
raise ValueError('\'y0\' must have a total shape of {}, or shape {} in the last two axes but has shape {}.'
.format((size, ), (size, size), y0.shape))
if len(y0.shape) > 1:
y0 /= np.sum(np.diagonal(y0, axis1=-2, axis2=-1), axis=-1)[:, None, None]
else:
y0 = np.diag(y0 / np.sum(y0))
return y0, c_complex_p
def _cast_v(v: Optional[array_like]):
"""
:param v: Atom velocities. Must be a scalar or have shape (n, ) or (n, 3). In the first two cases,
the velocity vector(s) is assumed to be aligned with the x-axis.
:returns: The correctly shaped velocities with shape (n, 3).
:raises ValueError: If 'v' has the wrong shape.
"""
if v is None:
return np.array([[0, 0, 0]], dtype=float)
v = np.array(v, dtype=float, order='C')
if len(v.shape) == 0:
return np.array([[v, 0, 0]], dtype=float)
elif len(v.shape) == 1:
ret = np.zeros((v.size, 3), dtype=float)
ret[:, 0] = v
return ret
elif len(v.shape) == 2:
if v.shape[1] == 3:
return v
raise ValueError('\'v\' must be a scalar or have shape (n, ) or (n, 3) but has shape {}.'.format(v.shape))
[docs]class Interaction:
"""
Class representing an Interaction between lasers and an atom.
"""
def __init__(self, atom: Atom = None, lasers: Iterable[Laser] = None, delta_max: scalar = 1e3,
controlled: bool = False, instance=None):
"""
:param atom: The atom.
:param lasers: The lasers.
:param delta_max: The maximum absolute difference between a laser and a transition frequency
for that transition to be considered laser-driven (MHz). The default value is 1 GHz.
:param controlled: Whether the ODE solver uses an error controlled stepper or a fixed step size.
Setting this to True is particularly useful for dynamics where a changing resolution is required.
However, this comes at the cost of computing time.
:param instance: A pointer to an existing Interaction instance.
If this is specified, the other parameters are omitted.
"""
self.instance = instance
if self.instance is None:
self.instance = dll.interaction_construct()
self._environment = self._get_environemnt()
if atom is None:
atom = Atom()
if lasers is None:
lasers = []
self.atom = atom
self.lasers = list(lasers)
self.delta_max = delta_max
self.controlled = controlled
self.update()
else:
self._environment = self._get_environemnt()
self._atom = self._get_atom()
self._lasers = self._get_lasers()
self.update()
def __del__(self):
dll.interaction_destruct(self.instance)
def _get_environemnt(self):
"""
:return: The environment used in the C++ class.
"""
return Environment(instance=dll.interaction_get_environment(self.instance))
def _get_atom(self):
"""
:return: The atom used in the C++ class.
"""
return Atom(instance=dll.interaction_get_atom(self.instance))
def _get_lasers(self):
"""
:returns: The lasers used in the C++ class.
"""
return [Laser(0, instance=dll.interaction_get_laser(self.instance, m))
for m in range(dll.interaction_get_lasers_size(self.instance))]
[docs] def update(self):
"""
Updates the Interaction.
:returns: None.
"""
dll.interaction_update(self.instance)
[docs] def resonance_info(self):
"""
Prints the detunings of the base frequencies of the lasers in the given atomic system.
In particular useful for systems with a hyperfine structure. Here
.. math:: \\Delta = \\nu_0 - \\nu_L.
:returns: None.
"""
print('Resonance info:') # \n<label>(S, L, J, I, F, m) -> <label\'>(S\', L\', J\', I\', F\', m\')')
for k, (laser, laser_m) in enumerate(zip(self.lasers, self.get_rabi())):
n = 0
print('Laser {} @ {} MHz:'.format(k, laser.freq))
for i, state_i in enumerate(self.atom):
for j, state_j in enumerate(self.atom):
if np.abs(laser_m)[i, j] != 0 and i < j:
if state_i.freq < state_j.freq:
print('{} -> {}: {} MHz'.format(repr(state_i), repr(state_j),
state_j.freq - state_i.freq - laser.freq))
else:
print('{} -> {}: {} MHz'.format(repr(state_j), repr(state_i),
state_i.freq - state_j.freq - laser.freq))
n += 1
if n == 0:
print('No resonances!')
print()
@property
def environment(self):
"""
:returns: The environment of the interaction.
"""
return self._environment
@environment.setter
def environment(self, value: Environment):
"""
:param value: The new environment of the interaction.
:returns: None.
"""
self._environment = value
dll.interaction_set_environment(self.instance, value.instance)
@property
def atom(self):
"""
:returns: The atom of the interaction.
"""
return self._atom
@atom.setter
def atom(self, value: Atom):
"""
:param value: The new atom of the interaction.
:returns: None.
"""
self._atom = value
dll.interaction_set_atom(self.instance, value.instance)
@property
def lasers(self):
"""
:returns: The lasers of the interaction.
"""
return self._lasers
@lasers.setter
def lasers(self, value: Iterable[Laser]):
"""
:param value: The new lasers of the interaction.
:returns: None.
"""
if value is None:
value = []
self._lasers = list(value)
dll.interaction_clear_lasers(self.instance)
for laser in self.lasers:
dll.interaction_add_laser(self.instance, laser.instance)
@property
def delta_max(self):
"""
:returns: The maximum absolute difference between a laser and a transition frequency
for that transition to be considered laser-driven (MHz). The default value is 1 GHz.
"""
return dll.interaction_get_delta_max(self.instance)
@delta_max.setter
def delta_max(self, value: scalar):
"""
:param value: The new maximum absolute difference between a laser and a transition frequency
for that transition to be considered laser-driven (MHz). The default value is 1 GHz.
:returns: None.
"""
dll.interaction_set_delta_max(self.instance, c_double(value))
@property
def controlled(self):
"""
:returns: Whether the ODE solver uses an error controlled stepper or a fixed step size.
Setting this to True is particularly useful for dynamics where a changing resolution is required.
However, this comes at the cost of computing time.
"""
return dll.interaction_get_controlled(self.instance)
@controlled.setter
def controlled(self, value: bool):
"""
:param value: Whether the ODE solver uses an error controlled stepper or a fixed step size.
Setting this to True is particularly useful for dynamics where a changing resolution is required.
However, this comes at the cost of computing time.
:returns: None.
"""
dll.interaction_set_controlled(self.instance, c_bool(value))
@property
def dt(self):
"""
:returns: The maximum step size of the solver and the rough time spacing of generated results.
However, this comes at the cost of computing time.
"""
return dll.interaction_get_dt(self.instance)
@dt.setter
def dt(self, value: scalar):
"""
:param value: The maximum step size of the solver and the rough time spacing of generated results.
:returns: None.
"""
dll.interaction_set_dt(self.instance, c_double(value))
@property
def loop(self):
"""
:returns: Whether there are loops formed by the lasers in the atom.
"""
return dll.interaction_get_loop(self.instance)
@property
def time_dependent(self):
"""
:returns: Whether the system hamiltonian is allowed to be time dependent.
"""
return dll.interaction_get_time_dependent(self.instance)
@time_dependent.setter
def time_dependent(self, value: bool):
"""
:param value: Set whether the system hamiltonian is allowed to be time dependent.
:returns: None.
"""
dll.interaction_set_time_dependent(self.instance, c_bool(value))
[docs] def get_rabi(self, m: int = None):
"""
:param m: The laser number 'm'. If None, the Rabi frequencies are returned for all lasers
as an array with shape (#lasers, atom.size, atom.size).
:returns: The Rabi frequencies (generated by the laser 'm').
"""
matrix_cd_p = np.ctypeslib.ndpointer(dtype=np.complex128, shape=(self.atom.size, self.atom.size))
set_restype(dll.interaction_get_rabi, matrix_cd_p)
if m is None:
return np.array([dll.interaction_get_rabi(self.instance, c_size_t(_m))
for _m in range(len(self.lasers))], dtype=complex)
return dll.interaction_get_rabi(self.instance, c_size_t(m))
@property
def summap(self):
"""
:returns: A (atom.size x atom.size)-matrix indicating the states which are laser-connected.
"""
matrix_i_p = np.ctypeslib.ndpointer(dtype=int, shape=(self.atom.size, self.atom.size))
set_restype(dll.interaction_get_summap, matrix_i_p)
return dll.interaction_get_summap(self.instance)
@property
def atommap(self):
"""
:returns: A projection matrix A mapping the state frequencies onto the diagonal of the Hamiltonian.
It holds diag(H)_i <- sum_j(A_ij * state_j.freq).
"""
matrix_d_p = np.ctypeslib.ndpointer(dtype=float, shape=(self.atom.size, self.atom.size))
set_restype(dll.interaction_get_atommap, matrix_d_p)
return dll.interaction_get_atommap(self.instance).T
@property
def deltamap(self):
"""
:returns: A projection matrix B mapping the laser frequencies onto the diagonal of the Hamiltonian.
It holds diag(H)_i <- sum_j(B_im * laser_m.freq).
"""
matrix_d_p = np.ctypeslib.ndpointer(dtype=float, shape=(len(self.lasers), self.atom.size))
set_restype(dll.interaction_get_deltamap, matrix_d_p)
return dll.interaction_get_deltamap(self.instance).T
[docs] def get_delta(self):
vector_d_p = np.ctypeslib.ndpointer(dtype=float, shape=(self.atom.size, ))
set_restype(dll.interaction_get_delta, vector_d_p)
return dll.interaction_get_delta(self.instance)
@property
def history_size(self):
"""
:returns: The length of the history of states visited during the generation of the diagonal maps.
"""
return dll.interaction_get_n_history(self.instance)
@property
def history(self):
"""
:returns: The history of states visited during the generation of the diagonal maps.
"""
vector_i_p = np.ctypeslib.ndpointer(dtype=c_size_t, shape=(self.history_size, ))
set_restype(dll.interaction_get_history, vector_i_p)
return dll.interaction_get_history(self.instance)
[docs] def rates(self, t: array_like, delta: array_like = None, m: Optional[int] = 0, v: array_like = None,
y0: array_like = None):
"""
:param t: The times when to compute the solution.
:param delta: An array of frequency shifts for the laser(s). 'delta' must be a scalar or a 1d- or 2d-array
with shapes (n, ) or (n, #lasers), respectively.
:param m: The index of the shifted laser. If delta is a 2d-array, 'm' ist omitted.
:param v: Atom velocities. Must be a scalar or have shape (n, ) or (n, 3). In the first two cases,
the velocity vector(s) is assumed to be aligned with the x-axis.
:param y0: The initial state of the atom. This must be None or have shape (n, #states).
If None, the ground states are populated equally.
:returns: The integrated rate equations as a real-valued array of shape (n, #states, #times).
"""
t = np.asarray(t, dtype=float).flatten()
t.sort()
t_size = t.size
if isinstance(delta, np.ndarray) and isinstance(v, np.ndarray) and isinstance(y0, np.ndarray):
if delta.shape == v.shape == y0.shape:
if delta.flags.f_contiguous:
delta = np.ascontiguousarray(delta)
if v.flags.f_contiguous:
v = np.ascontiguousarray(v)
if y0.flags.f_contiguous:
y0 = np.ascontiguousarray(y0)
sample_size = delta.shape[0]
else:
delta = _cast_delta(delta, m, len(self.lasers))
v = _cast_v(v)
y0, ctype = _cast_y0(y0, 0, self.atom)
sample_size = max([delta.shape[0], v.shape[0], 1])
delta = np.array(np.broadcast_to(delta, (sample_size, len(self.lasers))), dtype=float, order='C')
v = np.array(np.broadcast_to(v, (sample_size, 3)), dtype=float, order='C')
y0 = np.array(np.broadcast_to(y0, (sample_size, self.atom.size)), dtype=float, order='C')
results = np.zeros((sample_size, self.atom.size, t_size), dtype=float)
dll.interaction_rates(self.instance, t.ctypes.data_as(c_double_p), delta.ctypes.data_as(c_double_p),
v.ctypes.data_as(c_double_p), y0.ctypes.data_as(c_double_p),
results.ctypes.data_as(c_double_p), c_size_t(t_size), c_size_t(sample_size))
return results
[docs] def schroedinger(self, t: array_like, delta: array_like = None, m: Optional[int] = 0, v: array_like = None,
y0: array_like = None):
"""
:param t: The times when to compute the solution.
:param delta: An array of frequency shifts for the laser(s). 'delta' must be a scalar or a 1d- or 2d-array
with shapes (n, ) or (n, #lasers), respectively.
:param m: The index of the shifted laser. If delta is a 2d-array, 'm' ist omitted.
:param v: Atom velocities. Must be a scalar or have shape (n, ) or (n, 3). In the first two cases,
the velocity vector(s) are assumed to be aligned with the x-axis.
:param y0: The initial state of the atom. This must be None or have shape (n, #states).
If None, only the first ground state is populated.
:returns: The integrated Schroedinger equation as a complex-valued array of shape (n, #states, #times).
"""
t = np.asarray(t, dtype=float).flatten()
t.sort()
t_size = t.size
if isinstance(delta, np.ndarray) and isinstance(v, np.ndarray) and isinstance(y0, np.ndarray):
if delta.shape == v.shape == y0.shape:
if delta.flags.f_contiguous:
delta = np.ascontiguousarray(delta)
if v.flags.f_contiguous:
v = np.ascontiguousarray(v)
if y0.flags.f_contiguous:
y0 = np.ascontiguousarray(y0)
sample_size = delta.shape[0]
else:
delta = _cast_delta(delta, m, len(self.lasers))
v = _cast_v(v)
y0, ctype = _cast_y0(y0, 1, self.atom)
sample_size = max([delta.shape[0], v.shape[0], y0.shape[0], 1])
delta = np.array(np.broadcast_to(delta, (sample_size, len(self.lasers))), dtype=float, order='C')
v = np.array(np.broadcast_to(v, (sample_size, 3)), dtype=float, order='C')
y0 = np.array(np.broadcast_to(y0, (sample_size, self.atom.size)), dtype=complex, order='C')
results = np.zeros((sample_size, self.atom.size, t_size), dtype=complex)
dll.interaction_schroedinger(self.instance, t.ctypes.data_as(c_double_p), delta.ctypes.data_as(c_double_p),
v.ctypes.data_as(c_double_p), y0.ctypes.data_as(c_complex_p),
results.ctypes.data_as(c_complex_p), c_size_t(t_size), c_size_t(sample_size))
return results
[docs] def master(self, t: array_like, delta: array_like = None, m: Optional[int] = 0, v: array_like = None,
y0: array_like = None):
"""
:param t: The times when to compute the solution.
:param delta: An array of frequency shifts for the laser(s). 'delta' must be a scalar or a 1d- or 2d-array
with shapes (n, ) or (n, #lasers), respectively.
:param m: The index of the shifted laser. If delta is a 2d-array, 'm' ist omitted.
:param v: Atom velocities. Must be a scalar or have shape (n, ) or (n, 3). In the first two cases,
the velocity vector(s) are assumed to be aligned with the x-axis.
:param y0: The initial state / density matrix of the atom.
This must be None or have shape (#states, ) or (n, #states, #states).
If None, the ground states are populated equally.
:returns: The integrated master equation as a complex-valued array of shape (n, #states, #states, #times).
"""
t = np.asarray(t, dtype=float).flatten()
t.sort()
t_size = t.size
if isinstance(delta, np.ndarray) and isinstance(v, np.ndarray) and isinstance(y0, np.ndarray):
if delta.shape == v.shape == y0.shape:
if delta.flags.f_contiguous:
delta = np.ascontiguousarray(delta)
if v.flags.f_contiguous:
v = np.ascontiguousarray(v)
if y0.flags.f_contiguous:
y0 = np.ascontiguousarray(y0)
sample_size = delta.shape[0]
else:
delta = _cast_delta(delta, m, len(self.lasers))
v = _cast_v(v)
y0, ctype = _cast_y0(y0, 2, self.atom)
sample_size = max([delta.shape[0], v.shape[0], 1])
delta = np.array(np.broadcast_to(delta, (sample_size, len(self.lasers))), dtype=float, order='C')
v = np.array(np.broadcast_to(v, (sample_size, 3)), dtype=float, order='C')
y0 = np.array(np.broadcast_to(y0, (sample_size, self.atom.size, self.atom.size)), dtype=complex, order='C')
results = np.zeros((sample_size, self.atom.size, self.atom.size, t_size), dtype=complex)
dll.interaction_master(self.instance, t.ctypes.data_as(c_double_p), delta.ctypes.data_as(c_double_p),
v.ctypes.data_as(c_double_p), y0.ctypes.data_as(c_complex_p),
results.ctypes.data_as(c_complex_p), c_size_t(t_size), c_size_t(sample_size))
return results
[docs] def mc_master(self, t: array_like, delta: array_like = None, m: Optional[int] = 0, v: array_like = None,
y0: array_like = None, dynamics: bool = False, ntraj: int = 500):
"""
:param t: The times when to compute the solution.
:param delta: An array of frequency shifts for the laser(s). 'delta' must be a scalar or a 1d- or 2d-array
with shapes (n, ) or (n, #lasers), respectively.
:param m: The index of the shifted laser. If delta is a 2d-array, 'm' ist omitted.
:param v: Atom velocities. Must be a scalar or have shape (n, ) or (n, 3). In the first two cases,
the velocity vector(s) are assumed to be aligned with the x-axis.
:param y0: The initial state of the atom. This must be None or have shape (n, #states).
If None, only the first ground state is populated.
:param dynamics: Whether to compute the dynamics of the photon-atom interactions.
:param ntraj: The number of samples to compute if no samples were given with 'delta', 'v', or 'y0'.
:returns: The integrated MC-Schroedinger equation as a complex-valued array of shape (n, #states, #times).
"""
if self.controlled:
raise ValueError('Controlled steppers are not supported with \'master_mc\' yet.'
' Decrease the step size if necessary.')
if dynamics and self.atom.mass <= 0:
raise ValueError('To simulate mechanical dynamics, the mass of the atom must be specified.')
t = np.asarray(t, dtype=float).flatten()
t.sort()
t_size = t.size
if isinstance(delta, np.ndarray) and isinstance(v, np.ndarray) and isinstance(y0, np.ndarray):
if delta.shape == v.shape == y0.shape:
if delta.flags.f_contiguous:
delta = np.ascontiguousarray(delta)
if v.flags.f_contiguous:
v = np.ascontiguousarray(v)
if y0.flags.f_contiguous:
y0 = np.ascontiguousarray(y0)
sample_size = delta.shape[0]
else:
delta = _cast_delta(delta, m, len(self.lasers))
v = _cast_v(v)
y0, ctype = _cast_y0(y0, 1, self.atom)
sample_size = max([delta.shape[0], v.shape[0], y0.shape[0], 1])
if sample_size == 1:
sample_size = ntraj
delta = np.array(np.broadcast_to(delta, (sample_size, len(self.lasers))), dtype=float, order='C')
v = np.array(np.broadcast_to(v, (sample_size, 3)), dtype=float, order='C')
y0 = np.array(np.broadcast_to(y0, (sample_size, self.atom.size)), dtype=complex, order='C')
results = np.zeros((sample_size, self.atom.size, t_size), dtype=complex)
dll.interaction_mc_master(self.instance, t.ctypes.data_as(c_double_p), delta.ctypes.data_as(c_double_p),
v.ctypes.data_as(c_double_p), y0.ctypes.data_as(c_complex_p), c_bool(dynamics),
results.ctypes.data_as(c_complex_p), c_size_t(t_size), c_size_t(sample_size))
return results, v
[docs] def scattering_rate(self, rho: array_like, theta: array_like = None, phi: array_like = None,
i: array_like = None, j: array_like = None, axis: int = 1):
"""
:param rho: The density matrix of the atom. Must have the same size as the atom
along the specified 'axis' and 'axis' + 1.
:param theta: The elevation angle of detection.
:param phi: The azimuthal angle of detection.
:param i: The initially excited state indexes to consider for spontaneous decay.
If None, all states are considered.
:param j: The final decayed state indexes to consider for spontaneous decay. If None, all states are considered.
:param axis: The axis along which the population is aligned in 'rho'.
:returns: The scattering rate of the atom given the population 'rho' (MHz or Events / s).
:raises ValueError: 'rho' must have the same size as the atom along the specified 'axis'.
"""
return self.atom.scattering_rate(rho, theta=theta, phi=phi, i=i, j=j, axis=axis)
def _define_colors(n: int, label_map: dict, colormap: str = None):
"""
:param n: The size of the system.
:param label_map: A dictionary with state labels as keys
and an array of the indices of the states with the labels as values.
:param colormap: A matplotlib colormap.
:returns: A list of colors with size n.
"""
cmap = cm.get_cmap(colormap)
labels = [(k, v) for k, v in label_map.items()]
labels = sorted(labels, key=lambda kv: min(kv[1]))
colors = ['', ] * n
for i, (label, indices) in enumerate(labels):
for index in indices:
if colormap is None:
colors[index] = tools.COLORS.PYPLOT[i % 10]
else:
colors[index] = cmap(i / (len(labels) - 1))
return colors