def d(self, p):
"""
Lazily estimate d_p(A) ~= || A^p ||^(1/p) where ||.|| is the 1-norm.
"""
if p not in self._d:
est = onenormest((self._a * aslinearoperator(self._A))**p)
self._d[p] = est**(1.0 / p)
return self._d[p]
def d(self, p):
"""
Lazily estimate d_p(A) ~= || A^p ||^(1/p) where ||.|| is the 1-norm.
"""
if p not in self._d:
matvec = lambda v: self._a * (self._A.dot(v) - self._mu * v)
rmatvec = lambda v: _np.conj(self._a) * (self._A.H.dot(v) - _np.
conj(self._mu) * v)
LO = LinearOperator(self._A.shape,
dtype=self._dtype,
matvec=matvec,
rmatvec=rmatvec)
est = onenormest(LO**p)
# est = onenormest((self._a * aslinearoperator(self._A))**p)
self._d[p] = est**(1.0 / p)
return self._d[p]
def _onenormest_m1_power(A, p, t=2, itmax=5, compute_v=False, compute_w=False):
"""
Efficiently estimate the 1-norm of (A - I)^p.
Parameters
----------
A : ndarray
Matrix whose 1-norm of a power is to be computed.
p : int
Non-negative integer power.
t : int, optional
A positive parameter controlling the tradeoff between
accuracy versus time and memory usage.
Larger values take longer and use more memory
but give more accurate output.
itmax : int, optional
Use at most this many iterations.
compute_v : bool, optional
Request a norm-maximizing linear operator input vector if True.
compute_w : bool, optional
Request a norm-maximizing linear operator output vector if True.
Returns
-------
est : float
An underestimate of the 1-norm of the sparse matrix.
v : ndarray, optional
The vector such that ||Av||_1 == est*||v||_1.
It can be thought of as an input to the linear operator
that gives an output with particularly large norm.
w : ndarray, optional
The vector Av which has relatively large 1-norm.
It can be thought of as an output of the linear operator
that is relatively large in norm compared to the input.
"""
return onenormest(
_MatrixM1PowerOperator(A, p),
t=t,
itmax=itmax,
compute_v=compute_v,
compute_w=compute_w)
def _onenormest_m1_power(A, p, t=2, itmax=5, compute_v=False, compute_w=False):
"""
Efficiently estimate the 1-norm of (A - I)^p.
Parameters
----------
A : ndarray
Matrix whose 1-norm of a power is to be computed.
p : int
Non-negative integer power.
t : int, optional
A positive parameter controlling the tradeoff between
accuracy versus time and memory usage.
Larger values take longer and use more memory
but give more accurate output.
itmax : int, optional
Use at most this many iterations.
compute_v : bool, optional
Request a norm-maximizing linear operator input vector if True.
compute_w : bool, optional
Request a norm-maximizing linear operator output vector if True.
Returns
-------
est : float
An underestimate of the 1-norm of the sparse matrix.
v : ndarray, optional
The vector such that ||Av||_1 == est*||v||_1.
It can be thought of as an input to the linear operator
that gives an output with particularly large norm.
w : ndarray, optional
The vector Av which has relatively large 1-norm.
It can be thought of as an output of the linear operator
that is relatively large in norm compared to the input.
"""
return onenormest(_MatrixM1PowerOperator(A, p),
t=t,
itmax=itmax,
compute_v=compute_v,
compute_w=compute_w)
def py_zhexpv(v,
A,
anorm=None,
wsp=None,
iwsp=None,
m=20,
t=1.0,
tol=0.0,
return_work=False):
A = aslinearoperator(A)
v = v.astype(np.complex128, casting="safe", copy=False).ravel()
n = v.shape[0]
if A.shape[1] != A.shape[0]:
raise ValueError("Expecting square LinearOperator.")
if A.shape[1] != v.shape[0]:
raise ValueError(
"Dimension mismatch between LinearOperator and input vector.")
if anorm is None:
anorm = onenormest(A)
if wsp is None:
wsp = np.zeros(7 + n * (m + 2) + 5 * (m + 2) * (m + 2),
dtype=np.complex128)
if iwsp is None:
iwsp = np.zeros(m + 2, dtype=np.int32)
if return_work:
return dict(anorm=anorm, wsp=wsp, iwsp=iwsp, m=m, return_work=False)
u, tol0, iflag0 = zhexpv(m, t, v, tol, anorm, wsp, iwsp, A.matvec, 0)
if iflag0 > 0:
raise ExpokitError(messages[iflag0])
elif iflag0 < 0:
raise ExpokitError("bad input arguments")
return u
def residual_gevp(A,B,la,evec):
r"""
:param A: left matrix
:param B: right matrix
:param la: eigenvalues
:param evec: eigenvectors
:return: res - residual
"""
n = la.shape[0]
res = numpy.zeros(n,numpy.complex_)
for i in range(0, n):
vec = evec[:,i:i+1]
AB = A - la[i]*B
res[i] = linalg.norm(AB.dot(vec))
#calculate relative residual
if(sparse.issparse(AB)):
AB_norm=norm.onenormest(AB, t=3, itmax=5, compute_v=False, compute_w=False)
else:
AB_norm=linalg.norm(AB)
res[i] = res[i]/(AB_norm*linalg.norm(vec))
return res;
def residual_qevp(M,C,K,la,evec):
r"""
:param M: Mass Matrix
:param C: Damping Matrix
:param K: Stiffness Matrix
:param la: eigenvalues
:param evec: eigenvectors
:return: res - residual
"""
n = la.shape[0]
res = numpy.zeros(n,numpy.complex_)
for i in range(0, n):
vec = evec[:,i:i+1]
A = la[i]*la[i]*M + la[i]*C + K
res[i] = linalg.norm(A.dot(vec))
#calculate relative residual
if(sparse.issparse(A)):
A_norm=norm.onenormest(A, t=3, itmax=5, compute_v=False, compute_w=False)
else:
A_norm=linalg.norm(A)
res[i] = res[i]/(A_norm*linalg.norm(vec))
return res;
def shift_matrices(obj, m, c, k, tau):
r"""
:param obj: object of the class ``BrakeClass``
:param m: Mass Matrix
:param c: Damping Matrix
:param k: Stiffness Matrix
:param tau: Shift
:return: M, C, K - Shifted Mass, Damping and Stiffness Matrix respectively
Procedure::
The M , C , K are obtained as follows:
M = m
C = 2 * tau * m + c
K = tau_squared * m + tau * c + k
"""
#object attributes used in the function
LOG_LEVEL = obj.log_level
logger_t = obj.logger_t
logger_i = obj.logger_i
if (LOG_LEVEL == 10): #Debug Mode
logger_i.debug("\n" + "\n" + "\n" +
'In shift matrices (shift_matrices.py)')
logger_i.debug(
'----------------------------------------------------------------')
logger_i.debug('Shifting about the point ' + str(tau))
M = m
C = 2 * tau * m + c
tau_squared = tau * tau
K = tau_squared * m + tau * c + k
if (LOG_LEVEL == 10): #Debug Mode
m_norm = norm.onenormest(M,
t=3,
itmax=5,
compute_v=False,
compute_w=False)
c_norm = norm.onenormest(C,
t=3,
itmax=5,
compute_v=False,
compute_w=False)
k_norm = norm.onenormest(K,
t=3,
itmax=5,
compute_v=False,
compute_w=False)
logger_i.debug('Properties of shifted matrices')
logger_i.debug("\n" + 'M ' + ' ' + 'No of nonzeros = ' + str(M.nnz) +
' 1-Norm = ' + str(m_norm))
logger_i.debug('C ' + ' ' + 'No of nonzeros = ' + str(C.nnz) +
' 1-Norm = ' + str(c_norm))
logger_i.debug('K ' + ' ' + 'No of nonzeros = ' + str(K.nnz) +
' 1-Norm = ' + str(k_norm))
return M, C, K
def GetConditionNumber(self,A):
self.matrix_condition_number = onenormest(K_b)
return self.matrix_condition_number
def create_MCK(obj, sparse_list, omega):
r"""
:param obj: object of the class ``BrakeClass``
:param sparse_list: a python list of matrices in Compressed Sparse Column format
of type '<type 'numpy.float64'>',
:param omega: angular frequency
:return: M - Mass Matrix, C - Damping Matrix, K - Stiffness Matrix
:raises: Assemble_BadInputError, When a matrix in the list is not sparse
:raises: Assemble_BadInputError, When a matrix in the list is not square
:raises: Assemble_BadInputError, When the matrix are not of the same size
Procedure::
The M , C , K are assembled as follows:
- M = M1
- C = D1+DR*(omegaRef/omega)+DG*(omega/omegaRef)
- K = K1+KR+KGeo*math.pow((omega/omegaRef),2)
"""
#object attributes used in the function
LOG_LEVEL = obj.log_level
logger_t = obj.logger_t
logger_i = obj.logger_i
omegaRef = obj.omegaRef
fRef = obj.fRef
#unittesting
if len(sparse_list) != 8:
raise Assemble_BadInputError('The sparse list is not of length 8')
else:
for i in range(0,len(sparse_list)):
if(not scipy.sparse.issparse(sparse_list[i])):
raise Assemble_BadInputError('The list is not sparse')
if (sparse_list[i].shape[0]!=sparse_list[i].shape[1]):
raise Assemble_BadInputError('The matrix is not square')
if (sparse_list[i].shape[0]!=sparse_list[0].shape[0]):
raise Assemble_BadInputError('The matrix are not of the same size')
if(LOG_LEVEL==10): #Debug Mode
logger_i.debug("\n"+"\n"+"\n"+'In assemble matrices (assemble.py)')
logger_i.debug('----------------------------------------------------------------')
#converting csc to csr format
M1 = sparse_list[0].tocsr()
D1 = sparse_list[1].tocsr()
DG = sparse_list[2].tocsr()
DR = sparse_list[3].tocsr()
D4 = sparse_list[4].tocsr()
K1 = sparse_list[5].tocsr()
KR = sparse_list[6].tocsr()
KGeo = sparse_list[7].tocsr()
if(LOG_LEVEL==10): #Debug Mode
logger_i.debug('Matrices in CSC format converted to CSR')
for i in range(0,len(sparse_list)):
componentMatrix = sparse_list[i]
normMatrix = norm.onenormest(componentMatrix.tocsr(), t=3, itmax=5, compute_v=False, compute_w=False)
logger_i.debug(obj.data_file_list[i]+' '+obj.data_file_name[i]+' Nonzeros = '+str(componentMatrix.nnz)+' 1-Norm = '+str(normMatrix))
#Old implementation from sarosh's code
'''
M = m
C = c1+c2*(omega/omegaRef)+c3*((omegaRef/omega)-1)+c4/(2*math.pi*fRef)
K = k1+k2+k3*(math.pow((omega/omegaRef),2)-1);
'''
M = M1
C = D1+DR*(omegaRef/omega)+DG*(omega/omegaRef)
K = K1+KR+KGeo*math.pow((omega/omegaRef),2)
if(LOG_LEVEL==10): #Debug Mode
m_norm = norm.onenormest(M, t=3, itmax=5, compute_v=False, compute_w=False)
c_norm = norm.onenormest(C, t=3, itmax=5, compute_v=False, compute_w=False)
k_norm = norm.onenormest(K, t=3, itmax=5, compute_v=False, compute_w=False)
logger_i.debug("\n"+"\n"+'Properties of assembled matrices')
logger_i.debug('M '+' '+'Nonzeros = '+str(M.nnz)+' 1-Norm = '+str(m_norm))
logger_i.debug('C '+' '+'Nonzeros = '+str(C.nnz)+' 1-Norm = '+str(c_norm))
logger_i.debug('K '+' '+'Nonzeros = '+str(K.nnz)+' 1-Norm = '+str(k_norm))
return M, C, K;
print "\n" + "\n" + 'Beginning Data Analysis'
sparse_list = load.load_matrices(obj)
obj.logger_i.info("\n" + 'Matrices in CSC format converted to CSR')
obj.logger_i.info("\n\n" + 'Properties of various component matrices' + "\n")
for i in range(0, len(sparse_list)):
componentMatrix = sparse_list[i]
csrForm = componentMatrix.tocsr()
normMatrix = onenormest(csrForm,
t=3,
itmax=5,
compute_v=False,
compute_w=False)
#diffFlag = numpy.allclose(csrForm.data, csrForm.transpose().data)
#symmFlag = 'symmetric' if diffFlag else 'not symmetric'
eps = pow(10, -9)
rank = estimate_rank(aslinearoperator(componentMatrix), eps)
#print 'Approximate rank with relative error of(', eps, ')for numerical rank definition = ',rank
print obj.data_file_list[i], obj.data_file_name[
i], componentMatrix.shape, ' NonZeros = ', componentMatrix.nnz, ' 1-Norm = ', normMatrix, ' rank ', rank
obj.logger_i.info(obj.data_file_list[i] + ' ' + obj.data_file_name[i] +
str(componentMatrix.shape) + ' NonZeros = ' +
str(componentMatrix.nnz) + ' 1-Norm = ' +
def test_onenormest(B):
C = spla.onenormest(B)
npt.assert_allclose(C, np.linalg.norm(B.todense(), 1))
def scale_matrices(obj, m, c, k):
r"""
:param obj: object of the class ``BrakeClass``
:param m: Mass Matrix
:param c: Damping Matrix
:param k: Stiffness Matrix
:return: M, C, K - Scaled Mass, Damping and Stiffness Matrix respectively
:return: scaling parameters, - gamma and delta.
Procedure::
The M , C , K, gamma, delta are obtained as follows:
M = gamma*gamma*delta*m;
C = gamma*delta*c;
K = delta*k;
gamma = math.sqrt(k_norm/m_norm);
delta = 2/(k_norm+c_norm*gamma);
Note::
scipy.sparse.linalg.onenormest - Computes a lower bound of the 1-norm of a sparse matrix.
In the disk brake modelling theory spectral norm has been used but since I could not
find a better(less computational cost) way to obtain this in python, I have used 1-norm
approximation
"""
#object attributes used in the function
LOG_LEVEL = obj.log_level
logger_t = obj.logger_t
logger_i = obj.logger_i
if (LOG_LEVEL == 10): #Debug Mode
logger_i.debug("\n" + "\n" + "\n" +
'In scale matrices (scale_matrices.py)')
logger_i.debug(
'----------------------------------------------------------------')
m_norm = norm.onenormest(m, t=3, itmax=5, compute_v=False, compute_w=False)
c_norm = norm.onenormest(c, t=3, itmax=5, compute_v=False, compute_w=False)
k_norm = norm.onenormest(k, t=3, itmax=5, compute_v=False, compute_w=False)
gamma = math.sqrt(k_norm / m_norm)
delta = 2 / (k_norm + c_norm * gamma)
M = gamma * gamma * delta * m
C = gamma * delta * c
K = delta * k
if (LOG_LEVEL == 10): #Debug Mode
m_scaled_norm = gamma * gamma * delta * m_norm
c_scaled_norm = gamma * delta * c_norm
k_scaled_norm = delta * k_norm
rho = max(m_scaled_norm, c_scaled_norm, k_scaled_norm) / min(
m_scaled_norm, k_scaled_norm)
logger_i.debug('Scaling factors gamma = ' + str(gamma) + ' delta = ' +
str(delta))
logger_i.debug('The scaled problem has rho = ' + str(rho))
return M, C, K, gamma
def test_onenormest(matrices):
A_dense, A_sparse, b = matrices
est0 = splin.onenormest(A_dense)
est = splin.onenormest(A_sparse)
assert_allclose(est, est0)
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
#Putting the matrix in CSR form
csr = csr_matrix((value, index, pointer))
#Some tests on the matrix:
#Norm
norm = sp.norm(csr)
print('Frobenius norm =', norm)
#Lower bound of 1-norm
onenorm = sp.onenormest(csr)
print('Lower bound of 1-norm =', onenorm)
#Solving the system
#Note that sp.spsolve doesn't seem to work. It says that 'colind and rwoptr must be of type cint', so use lsmr instead
soln = sp.lsmr(csr, RHS) #Least squares approx
print('Least squares solution =', soln[0])
#Condition number
print('Condition number =', soln[6])
from _expokit import dgexpv, dsexpv, zgexpv, zhexpv
import numpy as np
from scipy.sparse import random
from scipy.sparse.linalg import onenormest, expm_multiply, aslinearoperator
n = 100
m = 20
time = 10.0
iflag = np.array([1])
tol = 0.0
v = np.zeros(n, dtype=complex)
v[0] = 1
v /= np.linalg.norm(v)
A = random(n, n, format="csr")
A = -1j * ((A.T + A) / 2)
anorm = onenormest(A)
wsp = np.zeros(7 + n * (m + 2) + 5 * (m + 2) * (m + 2), dtype=complex)
iwsp = np.zeros(m + 2, dtype=int)
A_op = aslinearoperator(A)
output_vec, tol0, iflag0 = zgexpv(m, time, v, tol, anorm, wsp, iwsp,
A_op.matvec, 0)
exact_vec = expm_multiply(time * A, v)
print iflag0
print np.linalg.norm(output_vec - exact_vec)
import meshes import FEM import time import scipy.sparse.linalg as spla meshwidth = 0.05 processes = 8 p,t=meshes.grid_square(1,meshwidth) timer=time.time() seqM=FEM.mass(p,t) seqTime=time.time()-timer timer=time.time() parM,processes=FEM.massParallel(p,t,processes) parTime=time.time()-timer print "difference of matrices: ", spla.onenormest(seqM-parM) print "time for sequential assembling:", seqTime print "time for parallel assembling: ", parTime print "speed-up: ", seqTime/parTime print "efficiency:", (seqTime/parTime)/processes