# -*- coding: utf-8 -*-
"""
This module contains the main class :class:`~MoleCool.System.System` which provides all
information about the lasers light fields, the atomic or molecular level structure,
and the magnetic field to carry out simulation calculations, e.g.
via the rate equations or Optical Bloch equations (OBEs).
"""
import numpy as np
from scipy.integrate import solve_ivp, cumulative_trapezoid
from scipy.constants import c,h,hbar,pi,g,physical_constants
from scipy.constants import k as k_B
from scipy.constants import u as u_mass
from sympy.physics.wigner import clebsch_gordan,wigner_3j,wigner_6j
from MoleCool.Lasersystem import *
from MoleCool.Levelsystem import *
from MoleCool import tools
from MoleCool.tools import save_object, open_object, ODEs, return_fun_default
from MoleCool.Bfield import Bfield
import time
import sys, os
from copy import deepcopy
import multiprocessing
from tqdm import tqdm
import warnings
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.pylab as pl
import matplotlib.ticker as mtick
np.set_printoptions(precision=4,suppress=True)
#%%
[docs]
class System:
def __init__(self, description=None,load_constants='',verbose=True):
"""An instance of this class is the starting point for simulating any
atomic or molecular dynamics simulations. Specficially, it defines an
object to not only store all important information about the system
but also to calculate any time evolution.
Parameters
----------
description : str, optional
A short description of this System can be added. If not provided,
the attribute is set to the name of the respective executed main
python file.
load_constants : str, optional
File name of a certain molecule, atom or more general system whose
respective level constants to be loaded or imported by the class
:class:`~MoleCool.Levelsystem.Levelsystem` via the constants
defined in the .json file. The default is ''.
Example
-------
After initiating a :class:`~MoleCool.System.System` object, the
instances of :class:`~MoleCool.Lasersystem.Lasersystem`,
:class:`~MoleCool.Levelsystem.Levelsystem`,
and :class:`~MoleCool.Bfield.Bfield` can be accessed via::
system = System()
print(system.lasers)
print(system.levels)
print(system.Bfield)
"""
self.lasers = Lasersystem()
self.levels = Levelsystem(load_constants=load_constants,verbose=verbose)
self.Bfield = Bfield()
#self.particles = Particlesystem()
if description == None:
self.description = os.path.basename(sys.argv[0])[:-3]
else:
self.description = description
self.reset_N0() # initial population of all levels
self.v0 = np.array([0.,0.,0.]) #: initial velocity of the particle
self.r0 = np.array([0.,0.,0.]) #: initial position of the particle
"""dictionary for parameters specifying the steady state conditions."""
self.steadystate = {'t_ini' : None,
'maxiters' : 100,
'condition' : [0.1,50],
'period' : None}
"""dictionary for parameters specifying the multiprocessing of calculations."""
self.multiprocessing = {'processes' : int(multiprocessing.cpu_count()*0.95),#None
'maxtasksperchild' : None,
'show_progressbar' : True,
'savetofile' : True}
if verbose:
print("System is created with description: {}".format(self.description))
[docs]
def calc_rateeqs(self,t_int=20e-6,t_start=0.,dt=None,t_eval = [],
magn_remixing=False, magn_strength=8,
position_dep=False, trajectory=False,
verbose=True, return_fun=return_fun_default,
**kwargs):
"""Calculate the time evolution of the single level populations with
rate equations.
Parameters
----------
t_int : float, optional
interaction time in which the molecule is exposed to the lasers.
The default is 20e-6.
t_start : float, optional
starting time when the ode_solver starts the calculation. Useful
for the situation when e.g. all cooling lasers are shut off at a
specific time t1, so that a new calculation with another laser
configuration (e.g. including only a probe laser) can be started
at t_start=t1 to continue the simulation. The default is 0.0.
dt : float, optional
time step of the output data. So in this case the ODE solver will
decide at which time points to calculate the solution.
The default is None.
t_eval : list or numpy array, optional
If it desired to get the solution of the ode solver only at
specific time points, the `t_eval` argument can be used to specify
these points. If `_eval` is given, the `dt` argument is ignored.
The default is [].
magn_remixing : bool, optional
if True, the adjacent ground hyperfine levels are perfectly mixed
by a magnetic field. The default is False.
magn_strength : float, optional
measure of the magnetic field strength (i.e. the magnetic remixing
matrix is multiplied by 10^magn_strength). Reasonable values are
between 6 and 9. The default is 8.
position_dep : bool, optional
determines whether to take the Gaussian intensity distribution of
the laser beams into account. The default is False.
trajectory : bool, optional
determines whether a trajectory of the molecule is calculated using
simple equations of motion. This yields the additional time-dependent
parameters ``v`` and ``r`` for the velocity and position. So, the force
which is acting on the molecule changes the velocity which in turn can
alter the Doppler shift. Further, the `position_dep` parameter determines
if either a uniform unitensity or complex intensity distribution due to
the Gaussian beam shapes is assumed through which the particle is
propagated. The default is False.
verbose : bool, optional
whether to print additional information like execution time or the
scattered photon number. The default is True.
return_fun : function-type, optional
if `mp` == True, the returned dictionary of this function determines
the quantities which are save for every single parameter configuration.
The default is None.
**kwargs : keyword arguments, optional
other options of the `solve_ivp` scipy function can be specified
(see homepage of scipy for further information).
Note
----
function creates attributes
* ``N`` : solution of the time dependent populations N,
* ``Nscatt`` : time dependent scattering number,
* ``Nscattrate`` : time dependent scattering rate,
* ``photons``: totally scattered photons
* ``args``: input arguments of the call of this function
* ``t`` : times at which the solution was calculated
* ``v`` : calculated velocities of the molecule at times ``t``
(only given if `trajectory` == True)
* ``r`` : calculated positions of the molecule at times ``t``
(only given if `trajectory` == True)
"""
self.calcmethod = 'rateeqs'
#___input arguments of this called function
self.args = locals()
self.check_config(raise_Error=True)
#___parameters belonging to the levels
self.levels.calc_all()
#___start multiprocessing if desired only after calculating the levels
# properties so that they don't have to be re-calculated every time.
if self._identify_iter_params() or self.lasers._identify_iter_params()[0]:
return self._start_mp()
#___parameters belonging to the lasers (and partially to the levels)
# wave vector k
self.k = self.lasers.getarr('k')*self.lasers.getarr('kabs')[:,None] #no unit vectors
self.delta = self.lasers.getarr('omega')[None,None,:] - self.levels.calc_freq()[:,:,None]
# saturation parameter of intensity (lNum,uNum,pNum)
self.sp = self.lasers.getarr('I')[None,None,:]/(self.levels.Isat[:,:,None])
#polarization switching time
tswitch = 1/self.lasers.freq_pol_switch
self.rx1 = np.abs(np.dot(self.levels.calc_dMat(),self.lasers.getarr('f_q').T))**2
# for polarization switching:
if np.any([la.pol_switching for la in self.lasers]):
self.rx2 = np.abs(np.dot(self.levels.calc_dMat(),self.lasers.getarr('f_q2').T))**2
else:
self.rx2 = self.rx1.copy()
#___magnetic remixing of the ground states. An empty array is left for no B-field
if magn_remixing:
self.M = self.Bfield.get_remix_matrix(self.levels.grstates[0],remix_strength=magn_strength)
else:
self.M = np.array([[],[]])
#___specify the initial (normalized) occupations of the levels
self.initialize_N0()
#___determine the time points at which the ODE solver should evaluate the equations
if len(t_eval) != 0: self.t_eval = np.array(t_eval)
else:
if dt != None and dt < t_int:
self.t_eval = np.linspace(t_start,t_start+t_int,int(t_int/dt)+1)
else:
self.t_eval = None
#___depenending on the position dependence two different ODE evaluation functions are called
if trajectory:
self.y0 = np.array([*self.N0, *self.v0, *self.r0])
else:
# position dependent intensity due to Gaussian shape of Laserbeam:
if position_dep:
self.sp *= self.lasers.I_tot(self.r0,sum_lasers=False,use_jit=False)[None,None,:]
self.R1 = self.levels.calc_Gamma()[None,:,None]/2*self.rx1*self.sp / (
1+4*(self.delta-np.dot(self.k,self.v0)[None,None,:])**2/self.levels.calc_Gamma()[None,:,None]**2)
self.R2 = self.levels.calc_Gamma()[None,:,None]/2*self.rx2*self.sp / (
1+4*(self.delta-np.dot(self.k,self.v0)[None,None,:])**2/self.levels.calc_Gamma()[None,:,None]**2)
#sum R1 & R2 over pNum:
self.R1sum, self.R2sum = np.sum(self.R1,axis=2), np.sum(self.R2,axis=2)
self.y0 = self.N0
#number of ground, excited states and lasers
lNum,uNum,pNum = self.levels.lNum, self.levels.uNum, self.lasers.pNum
# ---------------Ordinary Differential Equation solver----------------
#solve initial value problem of the ordinary first order differential equation with scipy
if verbose: print('Solving ode with rate equations...', end=' ')
start_time = time.perf_counter()
if not trajectory:
sol = solve_ivp(ODEs.ode0_rateeqs_jit, (t_start,t_start+t_int), self.y0,
t_eval=self.t_eval, **self.args['kwargs'],
args=(lNum,uNum,pNum,self.levels.calc_Gamma(),self.levels.calc_branratios(),
self.R1sum,self.R2sum,tswitch,self.M))
else:
sol = solve_ivp(ODEs.ode1_rateeqs_jit, (t_start,t_start+t_int), self.y0,
t_eval=self.t_eval, **self.args['kwargs'],
args=(lNum,uNum,pNum,np.reshape(self.levels.calc_Gamma(),(1,-1,1)),self.levels.calc_branratios(),
self.rx1,self.rx2,self.delta,self.sp,
self.lasers.getarr('w'),self.lasers.getarr('_w_cylind'),
self.k,self.lasers.getarr('kabs'),self.lasers.getarr('r_k'),
self.lasers.getarr('_r_cylind_trunc'),self.lasers.getarr('_dir_cylind'), #unit vectors
self.levels.mass,tswitch,self.M,position_dep,self.lasers.getarr('beta')))
#velocity v and position r
self.v = sol.y[-6:-3]
self.r = sol.y[-3:]
#: execution time for the ODE solving
self.exectime = time.perf_counter()-start_time
#: array of the times at which the solutions are calculated
self.t = sol.t
#: solution of the time dependent populations N
self.N = sol.y[:lNum+uNum]
self._verify_calculation()
if return_fun == True: return {'system':self}
elif return_fun: return return_fun(self)
#%%
[docs]
def plot_all(self):
self.plot_N(); self.plot_Nsum(); self.plot_dt()
self.plot_Nscatt(); self.plot_Nscattrate(); self.plot_F()
if ('trajectory' in self.args) and self.args['trajectory']:
self.plot_r(); self.plot_v()
[docs]
def plot_N(self,figname=None,figsize=(12,5),smallspacing=0):
"""plot populations of all levels over time."""
if figname == None:
plt.figure('N ({}): {}, {}, {}'.format(
self.calcmethod,self.description,self.levels.description,
self.lasers.description), figsize=figsize)
else: plt.figure(figname,figsize=figsize)
N_sum = 0
for i,ElSt in enumerate(self.levels.electronic_states):
if 'v' in ElSt.states[0].QuNrs:
alphas = np.linspace(1,0,ElSt.v_max+2)[:-1] #####
j_v0 = np.array([jj for jj,st in enumerate(ElSt) if st.v == 0])
N_red = len(j_v0)
else:
alpha = 1.0
N = ElSt.N
for j, state in enumerate(ElSt.states):
ls = ['solid','dashed','dashdot','dotted','dashdotdotted',
'loosely dotted','loosely dashed','loosely dashdotted'][i]
if N > 10:
if 'v' in state.QuNrs:
if j in j_v0:
color = pl.cm.jet(np.argwhere(j_v0==j)[0][0]/(N_red-1))
else:
color = 'grey'
else:
color = pl.cm.jet(j/(N-1))
else:
color = 'C' + str(j)
if 'v' in state.QuNrs:
if state.is_lossstate:
color = 'k'
alpha = 1.0
else:
alpha = alphas[state.v]
label = str(state).split(f'{ElSt.gs_exs}=')[-1]
plt.plot(self.t*1e6,(self.N[j+N_sum,:]+smallspacing*i)*1e2,
label=label,ls=ls,color=color,alpha=alpha)
N_sum += N
plt.xlabel('time $t$ in $\mu$s')
plt.ylabel('Populations $N$ in %')
leg = plt.legend(title='States:',loc='center left',labelspacing=-0.0,
bbox_to_anchor=(1, 0.5),fontsize='x-small')
# set the linewidth of each legend object
for legobj in leg.legend_handles:
legobj.set_linewidth(1.4)
[docs]
def plot_Nscatt(self,sum_over_ElSts=False):
"""plot the scattered photon number over time (integral of `Nscattrate`)."""
plt.figure('Nscatt: {}, {}, {}'.format(self.description,self.levels.description,self.lasers.description))
plt.xlabel('time $t$ in $\mu$s')
plt.ylabel('Totally scattered photons')
Nscatt = self.get_Nscatt(sum_over_ElSts=sum_over_ElSts)
for i,ElSt in enumerate(self.levels.exstates):
plt.plot(self.t*1e6, Nscatt[i,:], '-',label=ElSt.label)
if Nscatt.shape[0]>1:
plt.plot(self.t*1e6, np.sum(Nscatt,axis=0), '-',label='Sum')
plt.legend()
[docs]
def plot_Nscattrate(self,sum_over_ElSts=False):
"""plot the photon scattering rate over time (derivative of `Nscatt`)."""
plt.figure('Nscattrate: {}, {}, {}'.format(self.description,self.levels.description,self.lasers.description))
plt.xlabel('time $t$ in $\mu$s')
plt.ylabel('Photon scattering rate $\gamma\prime$ in MHz')
Nscattrate = self.get_Nscattrate(sum_over_ElSts=sum_over_ElSts)
for i,ElSt in enumerate(self.levels.exstates):
plt.plot(self.t*1e6, Nscattrate[i,:]*1e-6, '-',label=ElSt.label)
if Nscattrate.shape[0]>1:
plt.plot(self.t*1e6, np.sum(Nscattrate,axis=0)*1e-6, '-',label='Sum')
plt.legend()
[docs]
def plot_Nsum(self):
"""plot the population sum of all levels over time to ensure a small
numerical deviation of the ODE solver."""
plt.figure('Nsum: {}, {}, {}'.format(self.description,self.levels.description,self.lasers.description))
plt.xlabel('time $t$ in $\mu$s')
plt.ylabel('Population sum $\sum N_i$')
plt.plot(self.t*1e6, np.sum(self.N,axis=0), '-')
[docs]
def plot_dt(self):
"""plot the time steps at which the populations are calculated. If no `dt`
argument is given for the calulations they are chosen from the ODE solver."""
if 'method' in self.args['kwargs']: method = self.args['kwargs']['method']
else: method = 'RK45'
plt.figure('dt: {}, {}, {}'.format(self.description,self.levels.description,self.lasers.description))
plt.xlabel('time $t$ in $\mu$s')
plt.ylabel('timestep d$t$ in s')
plt.plot(self.t[:-1]*1e6,np.diff(self.t),label=method)
plt.yscale('log')
plt.legend()
[docs]
def plot_v(self):
"""plot the velocity over time for all three axes 'x','y', and'z'."""
plt.figure('v: {}, {}, {}'.format(self.description,self.levels.description,self.lasers.description))
plt.xlabel('time $t$ in $\mu$s')
plt.ylabel('velocity $v$ in m/s')
ls_arr = ['-','--','-.']
for i,axis in enumerate(['x','y','z']):
plt.plot(self.t*1e6,self.v[i,:],label='$v_{}$'.format(axis),ls=ls_arr[i])
plt.legend()
[docs]
def plot_r(self):
"""plot the position over time for all three axes 'x','y', and'z'."""
plt.figure('r: {}, {}, {}'.format(self.description,self.levels.description,self.lasers.description))
plt.xlabel('time $t$ in $\mu$s')
plt.ylabel('position $r$ in m')
ls_arr = ['-','--','-.']
for i,axis in enumerate(['x','y','z']):
plt.plot(self.t*1e6,self.r[i,:],label='$r_{}$'.format(axis),ls=ls_arr[i])
plt.legend()
[docs]
def plot_FFT(self,only_sum=True,start_time=0.0):
"""plot the fast Fourier transform (FFT) of the time-dependent populations.
Parameters
----------
only_sum : bool, optional
if True the sum of the FFTs of all populations is plottet. Otherwise
the distinct FFTs for all levels are shown. The default is True.
start_time : float between 0 and 1, optional
starting time in units of the interaction time `t_int` at which the
FFT is calculated. The default is 0.0.
"""
FT_sum = 0
t_int = self.args['t_int']
for i,st in enumerate(self.levels):
FT = (np.fft.rfft(self.N[i,int(self.t.size*start_time):]).real)**2
mean_zero = FT[int(FT.size/4):].mean()
start = np.where(np.diff(FT)>0)[0][0]
if start < 3: start = 3
FT[:start] = mean_zero
if not only_sum:
FT[np.where(FT<2*mean_zero)[0]] = mean_zero
plt.plot(np.arange(FT.size)/(t_int*(1-start_time))*1e-6,FT*1.**i)
else:
FT_sum += FT
if only_sum:
self.FT_sum = FT_sum
plt.plot(np.arange(FT.size)/(t_int*(1-start_time))*1e-6,self.FT_sum)
plt.yscale('log')
plt.xlabel('Frequency $f$ in MHz')
plt.ylabel('Power spectrum of the FFT')
[docs]
def calc_Rabi_freqs(self,position_dep=False):
"""Calculate the (angular) pure Rabi frequencies for each transition
and each laser component (with 2*pi included).
Parameters
----------
position_dep : bool, optional
Whether the intensity for the calculation is evaluated at a certain
position within the Gaussian laser beam profiles (True) or at the
maximum (False). The default is False.
Returns
-------
np.ndarray((lNum,uNum,pNum))
Angular Rabi frequencies for each combination of laser, excited states,
and ground state.
"""
# saturation parameter of intensity (lNum,uNum,pNum)
self.sp = self.lasers.getarr('I')[None,None,:]/(self.levels.Isat[:,:,None])
self.Rabi_freqs = self.levels.calc_Gamma()[None,:,None]*np.sqrt(self.sp/2) \
* np.dot(self.levels.calc_dMat(), self.lasers.getarr('f_q').T)
# position dependent intensity due to Gaussian shape of Laserbeam:
if position_dep:
self.Rabi_freqs *= np.sqrt(self.lasers.I_tot(self.r0,sum_lasers=False,use_jit=False))[None,:]
# for simple levelsystems, the Rabi frequency can be set via the laser instances
Rabi_set = self.lasers.getarr('freq_Rabi') #this feature must be tested!?
if not np.all(np.isnan(Rabi_set)):
Rabi_freqs_max = np.abs(self.Rabi_freqs).max(axis=(0,1)) #shape (pNum)
for p,la in enumerate(self.lasers):
if not np.isnan(Rabi_set[p]):
ratio = Rabi_set[p]/Rabi_freqs_max[p]
self.Rabi_freqs[:,:,p] *= ratio
la.I *= ratio**2
return self.Rabi_freqs
[docs]
def plot_F(self,figname=None,axes=['x','y','z']):
"""plot the Force over time for all three axes 'x','y', and'z'."""
if figname == None:
plt.figure('F ({}): {}, {}, {}'.format(
self.calcmethod,self.description,self.levels.description,
self.lasers.description))
else: plt.figure(figname)
ls_arr = ['-','--','-.']
for axis in axes:
i = {'x':0,'y':1,'z':2}[axis]
plt.plot(self.t * 1e6, self.F[i,:] / self.hbarkG2,
label='$F_{}$'.format(axis),ls=ls_arr[i])
plt.xlabel('time $t$ in $\mu$s')
plt.ylabel('Force $F$ in $\hbar k \Gamma_{}/2$'.format(self.levels.exstates_labels[0]))
plt.legend()
[docs]
def plot_spectrum(self, wavelengths=[], lasers=True, transitions=True,
unit = 'MHz',
exs = [],
relative_to_wavelengths = False,
axs = [],
subplot_sep = 200,
laser_spectrum_kwargs = dict(),
transitions_kwargs = dict(),
):
"""Plot the spectrum of :class:`~.Lasersystem.Laser` objects
and transition spectra with their respective intensities.
This method cleverly combines :meth:`~.Lasersystem.plot_spectrum`
and :meth:`.Levelsystem.plot_transition_spectrum` methods.
Parameters
----------
wavelengths : list, optional
wavelengths that should be plotted within the range ``subplot_sep``.
By default all available laser wavelengths are used.
lasers : bool, optional
Whether to include the laser spectrum. The default is True.
transitions : bool, optional
Whether to include the transition spectrum. The default is True.
unit : str, optional
Unit of the x-axis to be plotted.
Can be one of ``['GHz','MHz','kHz','Hz','Gamma']``. Default is 'MHz'.
exs : list(str), optional
See ``exs`` in :meth:`.Levelsystem.plot_transition_spectrum`.
relative_to_wavelengths : bool, optional
Whether the x-axis should be plotted in absolute frequency units or
relative to ``wavelengths``.
axs : list of ``matplotlib.pyplot.axis`` objects, optional
axis objects to put the plot(s) on. The default is [].
subplot_sep : float, optional
Defines the range of the plotted x-axis and the separation for the
automatic inclusion of all wavlengths (see parameter ``wavelengths``)
in units of ``ElectronicState.Gamma``. Default is 200.
laser_spectrum_kwargs : kwargs, optional
Additional keyword arguments
(see :meth:`.Lasersystem.plot_spectrum`). The default is dict().
transitions_kwargs : kwargs, optional
Additional keyword arguments
(see :meth:`.Levelsystem.plot_transition_spectrum`). The default is dict().
Returns
-------
axs : list of ``matplotlib.pyplot.axis``.
Axes of the subplot(s).
"""
if (not lasers) and (not transitions):
raise Exception("Either one of <lasers> or <transitions> must be True!")
# exctract Gamma
if not 'exs' in transitions_kwargs:
transitions_kwargs['exs'] = [self.levels.exstates_labels[0]]
ExSt = transitions_kwargs['exs'][0]
Gamma = self.levels[ExSt].Gamma*1e6
if lasers and self.lasers.entries:
subplot_sep_las = subplot_sep*Gamma
if not wavelengths:
wavelengths = self.lasers._get_wavelength_regimes(subplot_sep_las)
std = laser_spectrum_kwargs.pop('std', Gamma/2)
axs_lasers = self.lasers.plot_spectrum(
axs = axs,
wavelengths = wavelengths,
unit = unit,
relative_to_wavelengths = relative_to_wavelengths,
subplot_sep = subplot_sep_las,
std = std,
**laser_spectrum_kwargs,
)
if transitions:
axs = [ax.twinx() for ax in axs_lasers]
for ax in axs_lasers:
ax.tick_params(axis='y', labelcolor='grey')
ax.yaxis.label.set_color('grey')
else:
axs = axs_lasers
if transitions and self.levels.exstates and self.levels.grstates and self.levels.states:
kwargs_sum = transitions_kwargs.pop(
'kwargs_sum', dict(color='k',alpha=0.7,ls='--'))
self.levels.plot_transition_spectrum(
ax = axs,
wavelengths = wavelengths,
E_unit = unit,
relative_to_wavelengths = relative_to_wavelengths,
subplot_sep = subplot_sep,
kwargs_sum = kwargs_sum,
**transitions_kwargs,
)
return axs
@property
def hbarkG2(self):
"""Calculate :math:`\hbar k \Gamma /2` using the wave vector defined
in the laser system and natural lifetime defined in the first excited
electronic state."""
lambs = self.lasers.getarr('lamb')
dev = 1e-3
diff = lambs.std()/lambs.mean()
Gamma = self.levels.exstates[0].Gamma*2*pi*1e6
if len(self.levels.exstates) > 1:
warnings.warn("For calculation of hbar * k * Gamma/2, Gamma of the first ElSt is used")
if diff > dev:
warnings.warn(
("For calculation of hbar * k * Gamma/2, the wavelengths "
f"of the lasers' wavelengths differ by {diff*1e2:.2f} % "
f"(which is more than {dev*1e2:.2f} %). The "
"returned unit might thus me inappropriate.")
)
return hbar*2*pi/lambs.mean()*Gamma/2
@property
def F(self):
"""calculate the force over time.
Returns
-------
F : np.ndarray, shape(3,ntimes)
Force array for all <ntimes> time points and three axes 'x','y','z'.
"""
if self.calcmethod == 'rateeqs':
if not self.args['trajectory']:
lNum,uNum = self.levels.lNum, self.levels.uNum
N_lu = self.N[:lNum,:][:,None,:] - self.N[lNum:lNum+uNum,:][None,:,:]
F = hbar * np.sum( np.dot(self.R1,self.k)[:,:,:,None] * N_lu[:,:,None,:], axis=(0,1)) #+ g
else:
F = np.zeros((3,self.t.size))
F[:,1:] = np.diff(self.v)/np.diff(self.t)*self.levels.mass
if self.calcmethod == 'OBEs':
T = self.freq_unit*self.t
T_size = T.size
size = self.levels.lNum*self.levels.uNum*self.lasers.pNum*3*T_size #total size of the force array below
T_step = T_size//(size//int(10e6) + 1)
F = np.zeros((3,T_size))
for i,t1 in enumerate(range(0,T_size,T_step)): #ensure that not too much memory is needed for huge np arrays
t2 = t1 + T_step
if t2 > T_size: t2 = T_size
F[:,t1:t2] += np.real(np.transpose(self.ymat[self.levels.lNum:,:self.levels.lNum,t1:t2],axes=(1,0,2))[:,:,None,None,:] \
*self.Gfd[:,:,:,None,None]*np.exp(1j*T[None,None,None,None,t1:t2]*self.om_gek[:,:,:,None,None]) \
*self.k[None,None,:,:,None] ).sum(axis=(0,1,2))
F *= 2*hbar*self.freq_unit
return F
#%%
[docs]
def calc_trajectory(self,F_profile,t_int=20e-6,t_start=0.,dt=None,t_eval=None,
verbose=True,force_axis=None,
interpol_kind='linear',save_scipy_sols=False,**kwargs):
"""Calculate Monte Carlo simulations of classical particles
which are propagated through a provided pre-calculated force profile
to be used as interpolated function"""
if 'method' not in kwargs:
kwargs['method'] = 'LSODA' # setting default solver method for ODE
# Checking F_profile argument for datatype, etc.:
if not isinstance(F_profile, dict):
if np.all([isinstance(dic,dict) for dic in F_profile]):
F_profile = {k:v for dic in F_profile for k,v in dic.items()}
else:
raise ValueError(f"F_profile must be dictionary or iterable of dictionaries!")
if ('a' not in F_profile) and ('F' in F_profile):
F_profile['a'] = F_profile['F']/self.levels.mass
print(f'Converting force to acceleration with mass {self.levels.mass}')
else:
raise ValueError("Either 'F' or 'a' must be included in F_profile!")
# attribute for storing the solutions of the trajectory simulations:
self.trajectory_results = dict(kwargs=locals(), sols=[])
v0_arr = np.atleast_2d(self.v0)
r0_arr = np.atleast_2d(self.r0)
if isinstance(t_int, float):
t_int = np.ones(v0_arr.shape[0])*t_int
self.trajectory_results['final_values'] = dict(photons = np.zeros(len(v0_arr)),
v = np.zeros((len(v0_arr),3)),
r = np.zeros((len(v0_arr),3)))
if 'v' in F_profile or 'v0' in F_profile:
v = F_profile['v'] if 'v' in F_profile else F_profile['v0']
else:
raise ValueError("Either 'v' or 'v0' must be included in F_profile!")
if 'I' in F_profile:
position_dep = True
I = F_profile['I']
I_tot = self.lasers.get_intensity_func()
from scipy.interpolate import RegularGridInterpolator as interp
a_intp = interp((v,I), F_profile['a'], method=interpol_kind,
bounds_error=False,fill_value=None)
R_intp = interp((v,I), F_profile['Nscattrate'], method=interpol_kind,
bounds_error=False,fill_value=None)
def a(v,r): return a_intp(xi=(v,I_tot(r)))
def R(v,r): return R_intp(xi=(v,I_tot(r)))
if force_axis == '-v':
v0_arr2 = v0_arr.copy()
v0_arr2[:,0] *= 0
force_axis = -v0_arr2/np.linalg.norm(v0_arr2,axis=-1)[:,None]
elif isinstance(force_axis,(list,np.ndarray)):
force_axis = np.atleast_2d(np.array(force_axis)/np.linalg.norm(force_axis)) +v0_arr*0
else:
raise Exception('input argument <force axis> has to be given!')
else:
position_dep = False
from scipy.interpolate import interp1d
a = interp1d(v, F_profile['a'], kind=interpol_kind)
R = interp1d(v, F_profile['Nscattrate'], kind=interpol_kind)
force_axis = np.array(force_axis)/np.linalg.norm(force_axis)
# Differential equation calculation:
#__________________________________
def ode_MC1D(t,y,force_axis,position_dep):
dy = np.zeros(6+1)
v_proj = np.sum(y[:3]*force_axis)
if position_dep:
dy[:3] = a(v_proj,y[3:6])*force_axis
dy[-1] = R(v_proj,y[3:6])
else:
dy[:3] = a(v_proj)*force_axis
dy[-1] = R(v_proj)
dy[3:6] = y[:3]
return dy
#__________________________________
iterator = tqdm(v0_arr,smoothing=0.0) if verbose else v0_arr
for i,v0 in enumerate(iterator):
sol = solve_ivp(ode_MC1D, (0.,t_int[i]), np.array([*v0, *r0_arr[i], 0.0]),
t_eval=t_eval, args=(force_axis[i],position_dep),
**kwargs)
if save_scipy_sols:
self.trajectory_results['sols'].append(sol)
for key, indices in dict(photons=-1, r=[3,4,5], v=[0,1,2]).items():
self.trajectory_results['final_values'][key][i] = sol.y[indices,-1]
#%%
[docs]
def calc_OBEs(self, t_int=20e-6, t_start=0., dt=None, t_eval = [],
magn_remixing=False, freq_clip_TH=500, steadystate=False,
position_dep=False, rounded=False,
verbose=True, return_fun=return_fun_default,
**kwargs):
"""Calculate the time evolution of the single level populations with
the optical Bloch equations.
Parameters
----------
t_int : float, optional
interaction time in which the molecule is exposed to the lasers.
The default is 20e-6.
t_start : float, optional
starting time when the ode_solver starts the calculation. Useful
for the situation when e.g. all cooling lasers are shut off at a
specific time t1, so that a new calculation with another laser
configuration (e.g. including only a probe laser) can be started
at t_start=t1 to continue the simulation. The default is 0.0.
dt : float, optional
time step of the output data. So in this case the ODE solver will
decide at which time points to calculate the solution. If dt='auto',
an appropriate time step is chosen by using the smallest Rabi frequency
between a single transition and single laser component.
The default is None.
t_eval : list or numpy array, optional
If it desired to get the solution of the ode solver only at
specific time points, the `t_eval` argument can be used to specify
these points. If `_eval` is given, the `dt` argument is ignored.
The default is [].
magn_remixing : bool, optional
if True, the magnetic field, which is defined in the instance
:class:`~System.Bfield` contained in this class, is considered in the
calculation. Otherwise, the magnetic field is set to zero.
The default is False.
freq_clip_TH : float or string, optional
determines the threshold frequency at which the coupling of a single
transition detuned by a frequency from a specific laser component
is neglected. If a float is provided, only the transitions with
detunings smaller than `freq_clip_TH` times Gamma[0] are driven by the
light field. If `freq_clip_TH` == 'auto', the threshold frequencies for
all transitions are chosen seperately by considering the transition
strengths and intensities of each laser component.
The default is 500.
steadystate : bool, optional
determines whether the equations are propagated until a steady
state or periodic quasi-steady state is reached. The dictionary
``steadystate`` of this class specifies the steady state conditions.
The default is False.
position_dep : bool, optional
determines whether to take position of the particle in an Gaussian
intensity distribution of the laser beams into account.
The default is False.
rounded : float, optional
if specified, all frequencies and velocities are rounded to the frequency
`rounded` in units of max(Gamma).
The default is False.
verbose : bool, optional
whether to print additional information like execution time or the
scattered photon number. The default is True.
return_fun : function-type, optional
if `mp` == True, the returned dictionary of this function determines
the quantities which are save for every single parameter configuration.
The default is None.
**kwargs : keyword arguments, optional
other options of the `solve_ivp` scipy function can be specified
(see homepage of scipy for further information).
Note
----
function creates attributes
* ``N`` : solution of the time dependent populations N,
* ``Nscatt`` : time dependent scattering number,
* ``Nscattrate`` : time dependent scattering rate,
* ``photons``: totally scattered photons
* ``args``: input arguments of the call of this function
* ``t`` : times at which the solution was calculated
"""
self.calcmethod = 'OBEs'
#___input arguments of this called function
self.args = locals()
self.check_config(raise_Error=True)
#___parameters belonging to the levels
self.levels.calc_all()
# for dimensionless time units
freq_unit = self.levels.calc_Gamma().max() # choose the linewidth of the first electronic state as unit
self.freq_unit = freq_unit
#frequency differences between the ground and excited states (delta)
self.om_eg = self.levels.calc_freq()/freq_unit
if rounded:
self.om_eg = np.around(self.om_eg/rounded)*rounded
#___start multiprocessing if desired
if self._identify_iter_params() or self.lasers._identify_iter_params()[0]:
return self._start_mp()
#___parameters belonging to the lasers (and partially to the levels)
# Rabi frequency in dimensionless units (lNum,uNum,pNum)
self.calc_Rabi_freqs(position_dep=position_dep)
# wave vectors k (no unit vectors)
self.k = self.lasers.getarr('k')*self.lasers.getarr('kabs')[:,None]
# laser frequencies omega_k
if rounded:
self.om_k = np.around(self.lasers.getarr('omega')/freq_unit/rounded)*rounded \
- np.around(np.dot(self.k,self.v0)/freq_unit/rounded)*rounded
else:
self.om_k = (self.lasers.getarr('omega') - np.dot(self.k,self.v0))/freq_unit
#___magnetic remixing of the ground states and excited states
if magn_remixing:
betaB = self.Bfield.Bvec_sphbasis/(hbar*freq_unit/self.Bfield.mu_B)
else:
betaB = np.array([0.,0.,0.])
#coefficients h to neglect highly-oscillating terms of the OBEs (with frequency threshold freq_clip_TH)
self.om_gek = self.om_eg[:,:,None] - self.om_k[None,None,:]
if freq_clip_TH == 'auto':
FWHM = np.sqrt(self.levels.calc_Gamma()[None,:,None]**2 + 2*np.abs(self.Rabi_freqs)**2)/freq_unit #in dimensionless units
self.h_gek = np.where(np.abs(self.om_gek) < 8*FWHM/2, 1.0, 0.0)
self.h_gege = np.where(np.abs(self.om_eg[:,:,None,None]-self.om_eg[None,None,:,:])\
< 8*np.max(FWHM)/2, 1.0, 0.0)
else:
self.h_gek = np.where(np.abs(self.om_gek) < freq_clip_TH, 1.0, 0.0)
self.h_gege = np.where(np.abs(self.om_eg[:,:,None,None]-self.om_eg[None,None,:,:])\
< freq_clip_TH, 1.0, 0.0)
#___coefficients for new defined differential equations
self.Gfd = 1j/2*np.exp(1j*self.lasers.getarr('phi')[None,None,:])*self.h_gek*self.Rabi_freqs/freq_unit
self.betamu = tuple(1j* np.dot(self.levels.calc_muMat()[i], np.flip(betaB*np.array([-1,1,-1])))
for i in range(2) )
self.dd = self.h_gege*(self.levels.calc_dMat()[:,:,None,None,:]\
*self.levels.calc_dMat()[None,None,:,:,:]).sum(axis=-1)
self.ck_indices = (tuple(np.where(self.Gfd[i,:,:] != 0.0) for i in range(self.levels.lNum)),
tuple(np.where(self.Gfd[:,i,:] != 0.0) for i in range(self.levels.uNum)))
self.ck_indices = (tuple( np.array([i[0],i[1]]) for i in self.ck_indices[0] ),
tuple( np.array([i[0],i[1]]) for i in self.ck_indices[1] ))
#___specify the initial (normalized) occupations of the levels
# and transform these values into the density matrix elements N0mat
N0mat = self.initialize_N0(return_densitymatrix=True)
if verbose: print('Solving ode with OBEs...', end='')
start_time = time.perf_counter()
#___if steady state is wanted, multiple calculation steps of the OBEs
#___have to be performed while the occupations between this steps are compared
if not steadystate:
self._evaluate(t_start, t_int, dt, N0mat)
self.step = 0
else:
#___initial propagation of the equations for reaching the equilibrium region
if self.steadystate['t_ini']:
self.args['t_eval'] = [t_start, self.steadystate['t_ini']] #only the start and end point are important to be calculated for initial period
self._evaluate(t_start, self.steadystate['t_ini'], dt, N0mat)
self.args['t_eval'] = []
t_start = self.t[-1]
N0mat = self.ymat[:,:,-1]
#___specifying interaction time for the next multiple iterations to compare
# if callable(self.steadystate['period']):
if isinstance(self.steadystate['period'],float):
t_int = self.steadystate['period']
elif self.args['rounded']:
t_int = 2*pi/(freq_unit*self.args['rounded'])
elif self.steadystate['period'] == 'standingwave':
if np.linalg.norm(self.v0) != 0: #if v0==0, then t_int is not changed and thus used for int time.
lambda_mean = (c/(self.om_eg*freq_unit/2/pi)).mean()
if np.any(np.abs(c/(self.om_eg*freq_unit/2/pi) /lambda_mean -1)>0.1e-2 ):#percental deviation from mean
print('WARNING: averaging over standing wave periods might not be accurate since the wavelengths differ.')
period = lambda_mean/np.linalg.norm(self.v0)#/2
t_int = period*(t_int//period+1) # int(t_int - t_int % period)
self._evaluate(t_start, t_int, dt, N0mat)
t_start = self.t[-1]
N0mat = self.ymat[:,:,-1]
m1 = self.N.mean(axis=1)
# if self.steadystate['period'] == None: t_int *= 0.1
con1, con2 = self.steadystate['condition']
for self.step in range(1,self.steadystate['maxiters']):
self._evaluate(t_start, t_int, dt, N0mat)
m2 = self.N.mean(axis=1)
# print('diff & prop',np.all(np.abs(m1-m2)*1e2<con1),np.all(np.abs(1-m1/m2)*1e2 <con2))
#___check if conditions for steady state are fulfilled
if np.all(np.abs(m1-m2)*1e2 < con1) and np.all(np.nan_to_num(np.abs(1-m1/m2)*1e2,posinf=0,neginf=0) < con2):
break
else:
m1 = m2
N0mat = self.ymat[:,:,-1]
t_start = self.t[-1]
if verbose: print(' calculation steps: ',self.step+1)
#: execution time for the ODE solving
self.exectime = time.perf_counter() - start_time
self._verify_calculation()
if return_fun == True: return {'system':self}
elif return_fun: return return_fun(self)
def _verify_calculation(self):
dev_TH = {'rateeqs':1e-8,'OBEs':1e-6}[self.calcmethod]
dev = abs(self.N[:,-1].sum() -1)
#: Variable success indicates of the calcualtion could be verified.
self.success = True
self.message = ""
if dev > dev_TH:
message = 'Sum of populations not stable! Deviation: {:.2e}.\n'.format(dev)
print('WARNING:', message)
self.message += message
self.success = False
if np.any(self.N < -1e-3):
message = 'Populations got negative!'
print('WARNING:', message)
self.message += message
self.success = False
# printing some information...
if self.args['verbose']:
print("Execution time: {:2.4f} seconds".format(self.exectime))
for i,Ex_label in enumerate(self.levels.exstates_labels):
print("Scattered Photons ({}): {:.6f}".format(Ex_label,self.photons[i]))
def _start_mp(self):
self.results = tools.multiproc(obj=deepcopy(self),kwargs=self.args)
if self.multiprocessing['savetofile']:# multiprocessing.cpu_count() > 16:
save_object(self)
try:
sys.path.append('../')
import subprocess, sending_email
hostname = subprocess.check_output('hostname').decode("utf-8")
sending_email.send_message('Calculation complete!','File {} at Server {}'.format(self.description,hostname))
except:
pass
return None
def _identify_iter_params(self):
"""Identify which parameters are iterative arrays to loop through
including magnetic field strength and direction, initial position and
velocity. This function is inevitable to determine whether multiple
evaluations of the OBEs and rate equations are efficiently conducted
using the multiprocessing package from python (see :meth:`tools.multiproc`
and :meth:`Lasersystem._identify_iter_params`).
Returns
-------
iters_dict : dict
Dictionary with all iterative parameters and their number of iterations.
"""
iters_dict = {}
for obj, ndim, label in [
(self.Bfield.strength, 0, 'strength'),
(self.Bfield.direction, 1, 'direction'),
(self.r0, 1, 'r0'),
(self.v0, 1, 'v0'),
]:
if np.array(obj).ndim != ndim:
iters_dict[label] = len(obj)
return iters_dict
[docs]
def reset_N0(self):
"""Reset (last created) initial population self.N0 and initial populations
for all electronic states"""
self.N0 = np.array([]) #: initial population of all levels
for ElSt in self.levels.electronic_states:
ElSt.N0 = []
[docs]
def initialize_N0(self,return_densitymatrix=False,random=False):
"""
Initialize initial populations as a starting point for the OBEs or
rate equations using pre-defined population attribute ``N0`` of the
electronic states :class:`~MoleCool.Levelsystem.ElectronicState`.
Parameters
----------
return_densitymatrix : bool, optional
Whether the populations should be transformed into density matrix
or one-dimensional population array. The default is False.
random : bool, optional
Whether the popultions are sampled from a random distirbution.
The default is False.
Returns
-------
N0mat : numpy.ndarray
density matrix with populations on the diagonal.
"""
#___specify the initial (normalized) occupations of the levels
N,iNum = self.levels.N, self.levels.iNum
if random:
self.N0 = np.random.rand(N)
else:
if np.any([len(ElSt.N0) for ElSt in self.levels.electronic_states]):
self.N0 = np.zeros(N)
for ElSt in self.levels.electronic_states:
if len(ElSt.N0) != 0:
i_ElSt = self.levels.index_ElSt(ElSt,include_Ngrs_for_exs=True)
self.N0[i_ElSt:(i_ElSt+ElSt.N)] = ElSt.N0
else:
self.N0 = np.array(self.N0, dtype=float)
if len(self.N0) == 0:
if 'v' in self.levels.grstates[0][0].QuNrs:
N0_indices = [i for i,st in enumerate(self.levels.grstates[0]) if st.v==0]
if len(N0_indices) == 0:
self.N0 = np.ones(N)
else:
self.N0 = np.zeros(N)
for i in N0_indices:
self.N0[i] = 1.0
else:
self.N0 = np.ones(N)
else:
if len(self.N0) != N:
if len(self.N0) == N+iNum:
self.N0 = self.N0[:N]
else:
raise ValueError('Wrong size of N0')
self.N0 /= self.N0.sum() #initial populations are always normalized
if iNum > 0:
self.N0 = np.array([*self.N0,*self.N0[self.levels.lNum-iNum:self.levels.lNum]])
if return_densitymatrix:
N = N +iNum
N0mat = np.zeros((N,N),dtype=np.complex64)
if random:
print('MUST be modified that the density matrix elements that are twice apparent are equal')
N0mat = (np.random.rand(N,N) + 1j*np.random.rand(N,N))/2
N0mat += np.conj(N0mat).T
#transform these initial values into the density matrix elements N0mat
N0mat[(np.arange(N),np.arange(N))] = self.N0
return N0mat
[docs]
def get_N(self, return_sum=True, **QuNrs):
"""
Retrieve the time-dependent populations as results after calculating
the dynamics. Can be used to either obtain the populations of all
individual levels or to conveniently combine the populations for only a
subset of levels with specific Quantum numbers.
Parameters
----------
return_sum : bool, optional
Whether to sum up of all levels. The default is True.
**QuNrs : kwargs
Keyword arguments as Quantum numbers can be provided for only a subset
of levels with specific Quantum numbers, e.g. v=0. If empty, all
levels are considered. If only a specific ground or excited
electronic state should be included, add e.g. `gs='X'` or `exs='A'`
as a first QuNr keyword.
Returns
-------
np.ndarrays
Time-dependent population(s).
"""
if not QuNrs:
return self.N
states = self.levels.states
inds = np.array([i for i,st in enumerate(states) if st.check_QuNrvals(**QuNrs)])
if return_sum:
return self.N[inds,:].sum(axis=0)
else:
return self.N[inds,:]
[docs]
def get_Nscattrate(self,sum_over_ElSts=False):
"""Calculate time dependent scattering rate.
Parameters
----------
sum_over_ElSts : bool, optional
Whether to sum over multiple electronic excited states.
The default is False.
Returns
-------
numpy.ndarray
Time dependent scattering rate.
"""
Nscattrate_arr = self.levels.calc_Gamma()[:,None]*self.N[self.levels.lNum:,:]
if not sum_over_ElSts:# and (len(self.levels.exstates_labels) > 1)
Nscattrate_summed = np.zeros((len(self.levels.exstates_labels),self.N.shape[1]))
N_states = 0
for i,ElSt in enumerate(self.levels.exstates):
N = ElSt.N
Nscattrate_summed[i,:] = np.sum(Nscattrate_arr[N_states:N_states+N,:],axis=0)
N_states += N
return Nscattrate_summed
else:
return np.sum(Nscattrate_arr,axis=0)
[docs]
def get_Nscatt(self,sum_over_ElSts=False):
"""Calculate time dependent scatterd photon number.
Parameters
----------
sum_over_ElSts : bool, optional
Whether to sum over multiple electronic excited states.
The default is False.
Returns
-------
numpy.ndarray
Time dependent scattered photon number.
"""
return cumulative_trapezoid(self.get_Nscattrate(sum_over_ElSts=sum_over_ElSts),
self.t, initial = 0.0, axis=-1)
[docs]
def get_photons(self,sum_over_ElSts=False):
"""Calculate totally scattered photon number.
Parameters
----------
sum_over_ElSts : bool, optional
Whether to sum over multiple electronic excited states.
The default is False.
Returns
-------
numpy.ndarray
Totally scattered photon number.
"""
return np.transpose(self.get_Nscatt(sum_over_ElSts=sum_over_ElSts))[-1]
#___compute several physical variables using the solution of the ODE
#: time dependent scattering rate
Nscattrate = property(get_Nscattrate)
#: time dependent scattering number
Nscatt = property(get_Nscatt)
#: totally scattered photons
photons = property(get_photons)
def _evaluate(self,t_start,t_int,dt,N0mat): #for OBEs only?!
freq_unit = self.freq_unit
#___determine the time points at which the ODE solver should evaluate the equations
if len(self.args['t_eval']) != 0:
self.t_eval = np.array(self.args['t_eval'])
else:
if dt == 'auto': #must be tested!?!
dt = 1/np.max(self.h_gek*np.sqrt(self.om_gek**2*freq_unit**2+np.abs(self.Rabi_freqs)**2)/2/pi)/8.11 #1/8 of one Rabi-oscillation
if dt != None and dt < t_int:
self.t_eval = np.linspace(t_start,t_start+t_int,int(t_int/dt)+1)
else:
self.t_eval, T_eval = None, None
if np.all(self.t_eval) != None:
T_eval = self.t_eval * freq_unit
#___transform the initial density matrix N0mat in a vector for the ode
self.y0 = self._density_mat2vec(N0mat)
# ---------------Ordinary Differential Equation solver----------------
# solve initial value problem of the ordinary first order differential equation with scipy
lNum,uNum,iNum,pNum = self.levels.lNum,self.levels.uNum,self.levels.iNum,self.lasers.pNum
kwargs = self.args['kwargs']
# sol = solve_ivp(ode0_OBEs, (t_start*freq_unit,(t_start+t_int)*freq_unit),
# self.y0, t_eval=T_eval, **kwargs,
# args=(lNum,uNum,pNum,self.G,self.f,self.om_eg,self.om_k,
# betaB,self.dMat,self.muMat,
# self.M_indices,h_gek,h_gege,self.phi)) # delete?
# sol = solve_ivp(ode1_OBEs, (t_start*freq_unit,(t_start+t_int)*freq_unit),
# self.y0, t_eval=T_eval, **kwargs,
# args=(lNum,uNum,pNum, self.M_indices,
# self.Gfd,self.om_gek,self.betamu,self.dd))
# sol = solve_ivp(ode1_OBEs_opt2, (t_start*freq_unit,(t_start+t_int)*freq_unit), #<- optimized form for one el. ex. state!
# self.y0, t_eval=T_eval, **kwargs,
# args=(lNum,uNum,pNum, self.M_indices,
# self.Gfd,self.om_gek,self.betamu,self.dd,self.ck_indices))
if iNum == 0:
sol = solve_ivp(ODEs.ode1_OBEs_opt3, (t_start*freq_unit,(t_start+t_int)*freq_unit), #<- can also handle two electr. states with diff. Gamma
self.y0, t_eval=T_eval, **kwargs,
args=(lNum,uNum,pNum, self.levels.calc_M_indices(),
self.Gfd,self.om_gek,self.betamu,self.dd,self.ck_indices,
self.levels.calc_Gamma()/freq_unit))
else:
sol = solve_ivp(ODEs.ode1_OBEs_opt4, (t_start*freq_unit,(t_start+t_int)*freq_unit), #<- can also handle two electr. states with diff. Gamma
self.y0, t_eval=T_eval, **kwargs,
args=(lNum,uNum,iNum,pNum, self.levels.calc_M_indices(),
self.Gfd,self.om_gek,self.betamu,self.dd,self.ck_indices,
self.levels.calc_Gamma()/freq_unit))
self.sol = sol
self.ymat = self._density_vec2mat(sol.y)
#: solution of the time dependent populations N
self.N = np.real(self.ymat[(np.arange(self.levels.N),np.arange(self.levels.N))])
#: array of the times at which the solutions are calculated
self.t = sol.t/freq_unit
def _density_mat2vec(self,mat):
"""Transform the time dependent density matrix from matrix form to
solution vector of the ODEs. Inverse operation to _density_vec2mat.
"""
if mat.ndim != 2:
raise Exception('Matrix must be 2-dimensional')
if mat.shape[0] != mat.shape[1]:
raise Exception('Matrix must be square matrix')
N = mat.shape[0]
vec = np.zeros( N*(N+1) )
count = 0
for i in range(N):
for j in range(i,N):
vec[count] = mat[i,j].real
vec[count+1] = mat[i,j].imag
count += 2
return vec
def _density_vec2mat(self,vec):
"""Transform the solution vector of the time dependent density matrix
into matrix form. Inverse operation to _density_mat2vec.
"""
N,iNum = self.levels.N, self.levels.iNum
if vec.shape[0] != (N+iNum)*(N+iNum+1):
raise Exception('Shape[0] of vector must have the length N+iNum')
mat = np.zeros((N+iNum,N+iNum,vec.shape[-1]),dtype=np.complex64)
count = 0
for i in range(N+iNum):
for j in range(i,N+iNum):
mat[i,j,:] = vec[count] + 1j* vec[count+1]
count += 2
mat = mat[:N,:N]
if np.any(np.abs(mat[(np.arange(N),np.arange(N))].imag) > 1e-13):
warnings.warn('Populations got an imaginary part > 1e-13')
mat += np.conj(np.transpose(mat,axes=(1,0,2))) #is diagonal remaining purely real or complex?
mat[(np.arange(N),np.arange(N))] *=0.5
return mat
[docs]
def check_config(self,raise_Error=False):
"""Check System configuration for simulating internal dynamics."""
if self.calcmethod == 'rateeqs':
#pre-defined kwargs for solve_ivp function
kwargs_default = {'method':'LSODA', 'max_step':10e-6}
self.args['kwargs'] = dict(kwargs_default,**self.args['kwargs'])
self.lasers.check_config(raise_Error=raise_Error)
self.levels.check_config(raise_Error=raise_Error)
if 'trajectory' in self.args:
if (self.args['trajectory']) and (self.levels.mass==0.0):
raise ValueError('No mass is provided for trajectory to be calculated')
# check if some lasers are completely off to some states or if some states are not addressed by any laser!?
def __set_v0r0(self, vec=[0,0,0], direction=[], label='v0'):
vec = np.array(vec, dtype=np.float64)
if isinstance(direction, str):
direction = {'x':[1,0,0],'y':[0,1,0],'z':[0,0,1],'':[]}[direction]
if np.any(direction):
direction = np.array(direction)/np.linalg.norm(direction)
if direction.shape != (3,):
raise Exception(f"direction vector must be of shape (3,) instead of {direction.shape}")
vec = direction[None,:]*vec[:,None]
if vec.ndim not in [1,2]:
raise Exception(f"{label} must have one or two dimensions instead of {vec.ndim}")
if vec.shape[-1] != 3:
raise Exception(f"Last dimension of {label} must be of length 3 instead of {vec.shape[-1]}")
return vec
[docs]
def set_v0(self, v0=[0,0,0], direction=[]):
"""Set initial velocity of the particle(s).
Parameters
----------
v0 : list or np.ndarray, optional
Velocity array. Can be either a one-dimensional array with three
entries for the x, y and z component, a two dimensional array
where the first dimension corresponds to different particles'
velocities, or a one dimensional array with absolute velocity
values while the argument ``direction`` defines the direction.
The default is [0,0,0].
direction : str, list or np.ndarray, optional
Direction of the absolute velocity values defined as the argument
``v0``. Can be either a str (x,y,z) or list / array of length 3.
The default is [].
Returns
-------
np.ndarray
velocity array (1 or 2-dimensional with the shape (number of
particles, 3) in cartesian coordinates.
"""
self.__v0 = self.__set_v0r0(vec=v0, direction=direction, label='v0')
return self.__v0
[docs]
def set_r0(self, r0=[0,0,0], direction=[]):
"""Set initial position of the particle(s).
Parameters
----------
r0 : list or numpy.ndarray, optional
Position array. Can be either a one-dimensional array with three
entries for the x, y and z component, a two dimensional array
where the first dimension corresponds to different particles'
positions, or a one dimensional array with absolute position
values while the argument ``direction`` defines the direction.
The default is [0,0,0].
direction : str, list or numpy.ndarray, optional
Direction of the absolute position values defined as the argument
``r0``. Can be either a str (x,y,z) or list / array of length 3.
The default is [].
Returns
-------
numpy.ndarray
position array (1 or 2-dimensional with the shape (number of
particles, 3) in cartesian coordinates.
"""
self.__r0 = self.__set_v0r0(vec=r0, direction=direction, label='r0')
return self.__r0
@property
def v0(self):
"""Initial velocity vector or array of vectors."""
if 'v0' in self.__dict__:
val = self.__dict__['v0']
else:
val = self.__v0
return val
@v0.setter
def v0(self,val):
self.set_v0(v0=val)
@property
def r0(self):
"""Initial position vector or array of vectors."""
if 'r0' in self.__dict__:
val = self.__dict__['r0']
else:
val = self.__r0
return val
@r0.setter
def r0(self,val):
self.set_r0(r0=val)
[docs]
def draw_levels(self,GrSts=None,ExSts=None,branratios=True,lasers=True,
QuNrs_sep=[], br_fun='identity', br_TH=0.01, # 1e-16 default
freq_clip_TH='auto', cmap='viridis',yaxis_unit='MHz'):
"""Draw all levels of certain Electronic states sorted
by certain Qunatum numbers. Additionally, the branching ratios and the
transitions addressed by the lasers can be added.
Parameters
----------
GrSts : list of str, optional
list of labels of ground electronic states to be displayed.
The default is None which corresponds to all ground states.
ExSts : list of str, optional
list of labels of excited electronic states to be displayed.
The default is None which corresponds to all excited states.
branratios : bool, optional
Whether to show the branching ratios. The default is True.
lasers : bool, optional
Whether to show the transitions addressed by the lasers.
By default it is set to True when lasers are defined.
QuNrs_sep : list of str or tuple of two lists of str, optional
Quantum numbers for separating all levels into subplots.
For example the levels can be grouped into subplots by the vibrational
Quantum number, i.e. ['v'] or (['v'],['v']) for ex. and gr. states.
br_fun : str or callable, optional
Function to be applied onto the branching ratios. Can be either
'identity', 'log10', 'sqrt', or a custom defined function.
The default is 'identity'.
br_TH : float, optional
Threshold for the branching ratios to be shown. The default is 0.01.
freq_clip_TH : TYPE, optional
Same argument as in OBE's calculation method ':func:`calc_OBEs`.
The default is 'auto'.
cmap : str, optional
Colormap to be applied to the branching ratios. The default is 'viridis'.
yaxis_unit : str or float, optional
Unit of the y-axis. Can be either 'MHz','1/cm', or 'Gamma' for the
natural linewidth. Alternatively, an arbitrary unit (in MHz) can be
given as float. Default is 'MHz'.
"""
# from mpl.patches import ConnectionPatch
def draw_line(subfig,axes,xys,arrowstyle='->',dxyB=[0.,0.],**kwargs):
con = mpl.patches.ConnectionPatch(
xyA=xys[0], coordsA=axes[0].transData,
xyB=xys[1]+dxyB, coordsB=axes[1].transData,
arrowstyle=arrowstyle,
shrinkB=1,
**kwargs
)
subfig.add_artist(con)
axes[1].plot(*(xys[1]+dxyB)) #invisible line for automatically adjusting axis limits
levels = self.levels
if GrSts == None:
GrSts = levels.grstates
else:
GrSts = [levels.__dict__[label] for label in GrSts]
if ExSts == None:
ExSts = levels.exstates
else:
ExSts = [levels.__dict__[label] for label in ExSts]
# create big figure, and nested subfigures and subplot axes
fig = plt.figure(constrained_layout=True)
# subfigures subfigs for dividing main axes content and axes for legends
subfigs = fig.subfigures(1, 2, hspace=0.0, width_ratios=[5,1.5])
# Two main subfigures for ground and excited electronic state
height_ratios_main = [1, 2]
subfigs_main = subfigs[0].subfigures(2, 1, wspace=0.0, height_ratios=height_ratios_main )
subfigs_ExSts= subfigs_main[0].subfigures(1, len(ExSts), hspace=0.0, width_ratios=[Exs.N for Exs in ExSts],squeeze=False )[0]
subfigs_GrSts= subfigs_main[1].subfigures(1, len(GrSts), hspace=0.0, width_ratios=[Grs.N for Grs in GrSts],squeeze=False )[0]
# Two subfigure instances for legend and colorbar with each a single subplot axis
subfigs_leg = subfigs[1].subfigures(2, 1, wspace=0.0, height_ratios=height_ratios_main )
axs_legend = subfigs_leg[1].subplots(1, 1)
ax_cbar = subfigs_leg[0].subplots(1, 1)
tools.make_axes_invisible(axs_legend)
#needed for drawing arrows on top over multiple figures
subfig_last = subfigs_GrSts[-1]
# draw only levels first
if isinstance(QuNrs_sep,tuple): QuNrs_sep_u, QuNrs_sep_l = QuNrs_sep
else: QuNrs_sep_u, QuNrs_sep_l = 2*[QuNrs_sep]
coords_u = [ElSt.draw_levels(fig=subfigs_ExSts[i], QuNrs_sep=QuNrs_sep_u,
yaxis_unit=yaxis_unit, ylabel=not bool(i),
xlabel_pos='top')
for i, ElSt in enumerate(ExSts)]
coords_l = [ElSt.draw_levels(fig=subfigs_GrSts[i], QuNrs_sep=QuNrs_sep_l,
yaxis_unit=yaxis_unit)
for i, ElSt in enumerate(GrSts)]
try:
get_cmap = mpl.cm.get_cmap
except AttributeError:
get_cmap = mpl.colormaps.get_cmap
cmap = get_cmap(cmap)
# map branching ratios onto a color using a certain function:
if branratios:
branratios = levels.calc_branratios()
from scipy.interpolate import interp1d
if not callable(br_fun):
br_fun_dict = {'log10':np.log10,'identity':lambda x:x,'sqrt':np.sqrt}
if not br_fun in br_fun_dict: raise Exception('Not valid argument given for br_fun!')
br_fun = br_fun_dict[br_fun]
# isolate all branratios over a certain threshold in a separate flattened array and apply function
br_flat = br_fun(branratios[branratios > br_TH])
# map branratio onto interval [0,1] for colormap
map_br = interp1d([br_flat.min(),br_flat.max()],[1,0])
# iterate over states and draw branratios
for GrSt,coords_l_ in zip(GrSts,coords_l):
for ExSt,coords_u_ in zip(ExSts,coords_u):
if GrSt == ExSt: continue
Ng = self.levels.index_ElSt(GrSt,gs_exs='gs')
Ne = self.levels.index_ElSt(ExSt,gs_exs='exs')
for l,u in np.argwhere(branratios[Ng:Ng+GrSt.N,Ne:Ne+ExSt.N] > br_TH): #does not work when ExSts are switched, e.g. ['B','A']??!
draw_line(subfig_last,
axes=(coords_l_['axes'][l], coords_u_['axes'][u]),
xys=(coords_l_['xy'][l], coords_u_['xy'][u]),
arrowstyle='-',dxyB=[-0.2,0],alpha=0.5,
linewidth=0.6,linestyle='--',
color= cmap(map_br(br_fun(branratios[Ng+l,Ne+u]))))
# draw colorbar
norm = mpl.colors.Normalize(vmax=br_flat.max(),vmin=br_flat.min())
bounds = np.array(ax_cbar.get_position().bounds)
ax_cbar.set_position(bounds*np.array([1,1,0.2,1])) #left bottom width height
ax_cbar.tick_params(labelsize='small')
fig.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap.reversed()),
ticks=np.linspace(br_flat.min(),br_flat.max(),5),
format='{:.2%}'.format,
cax=ax_cbar, orientation='vertical', label='bran. ratio')
else:
tools.make_axes_invisible(ax_cbar)
# draw lasers onto their addressing transitions as arrows:
if self.lasers.pNum == 0: lasers = False
if lasers:
ls_fun = lambda x: ['-','-.',':'][x//10]
self.calc_OBEs(t_int=1e-11, t_start=0., dt=None, t_eval = [],
magn_remixing=False, freq_clip_TH=freq_clip_TH, steadystate=False,
position_dep=False, rounded=False,
verbose=False, return_fun=None, method='RK45')
for GrSt,coords_l_ in zip(GrSts,coords_l):
for ExSt,coords_u_ in zip(ExSts,coords_u):
if GrSt == ExSt: continue
Ng = self.levels.index_ElSt(GrSt,gs_exs='gs')
Ne = self.levels.index_ElSt(ExSt,gs_exs='exs')
for l,u,k in np.argwhere(np.abs(self.Gfd[Ng:Ng+GrSt.N,Ne:Ne+ExSt.N,:])>0):
det = self.om_gek[Ng+l,Ne+u,k]*self.freq_unit/2/pi*1e-6/coords_u_['yaxis_unit']
draw_line(subfig_last,
axes=(coords_l_['axes'][l], coords_u_['axes'][u]),
xys =(coords_l_['xy'][l], coords_u_['xy'][u]),
arrowstyle='->',linestyle=ls_fun(k),dxyB=[+0.2,det],
alpha=0.8,color='C'+str(k))
# legend for lasers
import matplotlib.lines as mlines
handles = [mlines.Line2D([],[],lw=1,color='C'+str(k),ls=ls_fun(k),
label='{:d}: {:.0f}, {:.1f}'.format(
k, la.lamb*1e9, la.P*1e3)
)
for k,la in enumerate(self.lasers)]
legend = axs_legend.legend(handles=handles,loc='upper left',
bbox_to_anchor=(0, 1),fontsize='x-small',
title='i: $\lambda$ [nm], $P$ [mW]')
plt.setp(legend.get_title(),fontsize='x-small')
#%%
if __name__ == '__main__':
system = System(description='test_System_mod',load_constants='138BaF')
# Build level scheme
system.levels.add_all_levels(v_max=0)
system.levels.X.del_lossstate()
# Add laser configuration
system.lasers.add_sidebands(
lamb = 859.83e-9,
P = 20e-3,
offset_freq = 19e6,
mod_freq = 39.33e6,
sidebands = [-2, -1, 1, 2],
ratios = [0.8, 1, 1, 0.8]
)
# Magnetic field
system.Bfield.turnon(strength=5e-4, direction=[1, 1, 1])
# Dynamics simulations
system.calc_OBEs(t_int=8e-6, dt=1e-9,
magn_remixing=True)
system.calc_rateeqs(t_int=8e-6, magn_remixing=True,
position_dep=True)
system.plot_N()
