def _parallel_mesolve(n, N, H, tlist, c_op_list, args, options):
col_idx, row_idx = np.unravel_index(n, (N, N))
rho0 = Qobj(
sp.csr_matrix(([1], ([row_idx], [col_idx])),
shape=(N, N),
dtype=complex))
output = mesolve(H,
rho0,
tlist,
c_op_list, [],
args,
options,
_safe_mode=False)
if config.tdname:
_cython_build_cleanup(config.tdname)
return output
def _correlation_me_2t(H, state0, tlist, taulist, c_ops, a_op, b_op, c_op,
args={}, options=Options()):
"""
Internal function for calculating the three-operator two-time
correlation function:
<A(t)B(t+tau)C(t)>
using a master equation solver.
"""
# the solvers only work for positive time differences and the correlators
# require positive tau
if state0 is None:
rho0 = steadystate(H, c_ops)
tlist = [0]
elif isket(state0):
rho0 = ket2dm(state0)
else:
rho0 = state0
if debug:
print(inspect.stack()[0][3])
rho_t = mesolve(H, rho0, tlist, c_ops, [],
args=args, options=options).states
corr_mat = np.zeros([np.size(tlist), np.size(taulist)], dtype=complex)
H_shifted, c_ops_shifted, _args = _transform_L_t_shift_new(H, c_ops, args)
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
for t_idx, rho in enumerate(rho_t):
if not isinstance(H, Qobj):
_args["_t0"] = tlist[t_idx]
corr_mat[t_idx, :] = mesolve(
H_shifted, c_op * rho * a_op, taulist, c_ops_shifted,
[b_op], args=_args, options=options
).expect[0]
if t_idx == 1:
options.rhs_reuse = True
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
return corr_mat
def _correlation_me_2t(H, state0, tlist, taulist, c_ops, a_op, b_op, c_op,
args={}, options=Options()):
"""
Internal function for calculating the three-operator two-time
correlation function:
<A(t)B(t+tau)C(t)>
using a master equation solver.
"""
# the solvers only work for positive time differences and the correlators
# require positive tau
if state0 is None:
rho0 = steadystate(H, c_ops)
tlist = [0]
elif isket(state0):
rho0 = ket2dm(state0)
else:
rho0 = state0
if debug:
print(inspect.stack()[0][3])
rho_t = mesolve(H, rho0, tlist, c_ops, [],
args=args, options=options).states
corr_mat = np.zeros([np.size(tlist), np.size(taulist)], dtype=complex)
H_shifted, c_ops_shifted, _args = _transform_L_t_shift(H, c_ops, args)
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
for t_idx, rho in enumerate(rho_t):
if not isinstance(H, Qobj):
_args["_t0"] = tlist[t_idx]
corr_mat[t_idx, :] = mesolve(
H_shifted, c_op * rho * a_op, taulist, c_ops_shifted,
[b_op], args=_args, options=options
).expect[0]
if t_idx == 1:
options.rhs_reuse = True
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
return corr_mat
def _correlation_2t(H, state0, tlist, taulist, c_ops, a_op, b_op, c_op,
solver="me", args={}, options=Options()):
"""
Internal function for calling solvers in order to calculate the
three-operator two-time correlation function:
<A(t)B(t+tau)C(t)>
"""
# Note: the current form of the correlator is sufficient for all possible
# two-time correlations (incuding those with 2ops vs 3). Ex: to compute a
# correlation of the form <A(t+tau)B(t)>: a_op = identity, b_op = A,
# and c_op = B.
if debug:
print(inspect.stack()[0][3])
if min(tlist) != 0:
raise TypeError("tlist must be positive and contain the element 0.")
if min(taulist) != 0:
raise TypeError("taulist must be positive and contain the element 0.")
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
H, c_ops, args = _td_wrap_array_str(H, c_ops, args, tlist)
if solver == "me":
return _correlation_me_2t(H, state0, tlist, taulist,
c_ops, a_op, b_op, c_op,
args=args, options=options)
elif solver == "mc":
return _correlation_mc_2t(H, state0, tlist, taulist,
c_ops, a_op, b_op, c_op,
args=args, options=options)
elif solver == "es":
return _correlation_es_2t(H, state0, tlist, taulist,
c_ops, a_op, b_op, c_op)
else:
raise ValueError("Unrecognized choice of solver" +
"%s (use me, mc, or es)." % solver)
def _correlation_2t(H, state0, tlist, taulist, c_ops, a_op, b_op, c_op,
solver="me", args={}, options=Options()):
"""
Internal function for calling solvers in order to calculate the
three-operator two-time correlation function:
<A(t)B(t+tau)C(t)>
"""
# Note: the current form of the correlator is sufficient for all possible
# two-time correlations (incuding those with 2ops vs 3). Ex: to compute a
# correlation of the form <A(t+tau)B(t)>: a_op = identity, b_op = A,
# and c_op = B.
if debug:
print(inspect.stack()[0][3])
if min(tlist) != 0:
raise TypeError("tlist must be positive and contain the element 0.")
if min(taulist) != 0:
raise TypeError("taulist must be positive and contain the element 0.")
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
H, c_ops, args = _td_wrap_array_str(H, c_ops, args, tlist)
if solver == "me":
return _correlation_me_2t(H, state0, tlist, taulist,
c_ops, a_op, b_op, c_op,
args=args, options=options)
elif solver == "mc":
return _correlation_mc_2t(H, state0, tlist, taulist,
c_ops, a_op, b_op, c_op,
args=args, options=options)
elif solver == "es":
return _correlation_es_2t(H, state0, tlist, taulist,
c_ops, a_op, b_op, c_op)
else:
raise ValueError("Unrecognized choice of solver" +
"%s (use me, mc, or es)." % solver)
def _generic_ode_solve(r, rho0, tlist, e_ops, opt, progress_bar):
"""
Internal function for solving ME. Solve an ODE which solver parameters
already setup (r). Calculate the required expectation values or invoke
callback function at each time step.
"""
#
# prepare output array
#
n_tsteps = len(tlist)
e_sops_data = []
output = Result()
output.solver = "mesolve"
output.times = tlist
if opt.store_states:
output.states = []
if isinstance(e_ops, types.FunctionType):
n_expt_op = 0
expt_callback = True
elif isinstance(e_ops, list):
n_expt_op = len(e_ops)
expt_callback = False
if n_expt_op == 0:
# fall back on storing states
output.states = []
opt.store_states = True
else:
output.expect = []
output.num_expect = n_expt_op
for op in e_ops:
e_sops_data.append(spre(op).data)
if op.isherm and rho0.isherm:
output.expect.append(np.zeros(n_tsteps))
else:
output.expect.append(np.zeros(n_tsteps, dtype=complex))
else:
raise TypeError("Expectation parameter must be a list or a function")
#
# start evolution
#
progress_bar.start(n_tsteps)
rho = Qobj(rho0)
dt = np.diff(tlist)
for t_idx, t in enumerate(tlist):
progress_bar.update(t_idx)
if not r.successful():
raise Exception("ODE integration error: Try to increase "
"the allowed number of substeps by increasing "
"the nsteps parameter in the Options class.")
if opt.store_states or expt_callback:
rho.data = dense2D_to_fastcsr_fmode(vec2mat(r.y), rho.shape[0], rho.shape[1])
if opt.store_states:
output.states.append(Qobj(rho, isherm=True))
if expt_callback:
# use callback method
e_ops(t, rho)
for m in range(n_expt_op):
if output.expect[m].dtype == complex:
output.expect[m][t_idx] = expect_rho_vec(e_sops_data[m],
r.y, 0)
else:
output.expect[m][t_idx] = expect_rho_vec(e_sops_data[m],
r.y, 1)
if t_idx < n_tsteps - 1:
r.integrate(r.t + dt[t_idx])
progress_bar.finished()
if (not opt.rhs_reuse) and (config.tdname is not None):
_cython_build_cleanup(config.tdname)
if opt.store_final_state:
rho.data = dense2D_to_fastcsr_fmode(vec2mat(r.y), rho.shape[0], rho.shape[1])
output.final_state = Qobj(rho, dims=rho0.dims, isherm=True)
return output
def _td_brmesolve(H, psi0, tlist, a_ops=[], e_ops=[], c_ops=[], args={},
use_secular=True, sec_cutoff=0.1,
tol=qset.atol, options=None,
progress_bar=None,_safe_mode=True,
verbose=False,
_prep_time=0):
if isket(psi0):
rho0 = ket2dm(psi0)
else:
rho0 = psi0
nrows = rho0.shape[0]
H_terms = []
H_td_terms = []
H_obj = []
A_terms = []
A_td_terms = []
C_terms = []
C_td_terms = []
CA_obj = []
spline_count = [0,0]
coupled_ops = []
coupled_lengths = []
coupled_spectra = []
if isinstance(H, Qobj):
H_terms.append(H.full('f'))
H_td_terms.append('1')
else:
for kk, h in enumerate(H):
if isinstance(h, Qobj):
H_terms.append(h.full('f'))
H_td_terms.append('1')
elif isinstance(h, list):
H_terms.append(h[0].full('f'))
if isinstance(h[1], Cubic_Spline):
H_obj.append(h[1].coeffs)
spline_count[0] += 1
H_td_terms.append(h[1])
else:
raise Exception('Invalid Hamiltonian specification.')
for kk, c in enumerate(c_ops):
if isinstance(c, Qobj):
C_terms.append(c.full('f'))
C_td_terms.append('1')
elif isinstance(c, list):
C_terms.append(c[0].full('f'))
if isinstance(c[1], Cubic_Spline):
CA_obj.append(c[1].coeffs)
spline_count[0] += 1
C_td_terms.append(c[1])
else:
raise Exception('Invalid collapse operator specification.')
coupled_offset = 0
for kk, a in enumerate(a_ops):
if isinstance(a, list):
if isinstance(a[0], Qobj):
A_terms.append(a[0].full('f'))
A_td_terms.append(a[1])
if isinstance(a[1], tuple):
if not len(a[1])==2:
raise Exception('Tuple must be len=2.')
if isinstance(a[1][0],Cubic_Spline):
spline_count[1] += 1
if isinstance(a[1][1],Cubic_Spline):
spline_count[1] += 1
elif isinstance(a[0], tuple):
if not isinstance(a[1], tuple):
raise Exception('Invalid bath-coupling specification.')
if (len(a[0])+1) != len(a[1]):
raise Exception('BR a_ops tuple lengths not compatible.')
coupled_ops.append(kk+coupled_offset)
coupled_lengths.append(len(a[0]))
coupled_spectra.append(a[1][0])
coupled_offset += len(a[0])-1
if isinstance(a[1][0],Cubic_Spline):
spline_count[1] += 1
for nn, _a in enumerate(a[0]):
A_terms.append(_a.full('f'))
A_td_terms.append(a[1][nn+1])
if isinstance(a[1][nn+1],Cubic_Spline):
CA_obj.append(a[1][nn+1].coeffs)
spline_count[1] += 1
else:
raise Exception('Invalid bath-coupling specification.')
string_list = []
for kk,_ in enumerate(H_td_terms):
string_list.append("H_terms[{0}]".format(kk))
for kk,_ in enumerate(H_obj):
string_list.append("H_obj[{0}]".format(kk))
for kk,_ in enumerate(C_td_terms):
string_list.append("C_terms[{0}]".format(kk))
for kk,_ in enumerate(CA_obj):
string_list.append("CA_obj[{0}]".format(kk))
for kk,_ in enumerate(A_td_terms):
string_list.append("A_terms[{0}]".format(kk))
#Add nrows to parameters
string_list.append('nrows')
for name, value in args.items():
if isinstance(value, np.ndarray):
raise TypeError('NumPy arrays not valid args for BR solver.')
else:
string_list.append(str(value))
parameter_string = ",".join(string_list)
if verbose:
print('BR prep time:', time.time()-_prep_time)
#
# generate and compile new cython code if necessary
#
if not options.rhs_reuse or config.tdfunc is None:
if options.rhs_filename is None:
config.tdname = "rhs" + str(os.getpid()) + str(config.cgen_num)
else:
config.tdname = opt.rhs_filename
if verbose:
_st = time.time()
cgen = BR_Codegen(h_terms=len(H_terms),
h_td_terms=H_td_terms, h_obj=H_obj,
c_terms=len(C_terms),
c_td_terms=C_td_terms, c_obj=CA_obj,
a_terms=len(A_terms), a_td_terms=A_td_terms,
spline_count=spline_count,
coupled_ops = coupled_ops,
coupled_lengths = coupled_lengths,
coupled_spectra = coupled_spectra,
config=config, sparse=False,
use_secular = use_secular,
sec_cutoff = sec_cutoff,
args=args,
use_openmp=options.use_openmp,
omp_thresh=qset.openmp_thresh if qset.has_openmp else None,
omp_threads=options.num_cpus,
atol=tol)
cgen.generate(config.tdname + ".pyx")
code = compile('from ' + config.tdname + ' import cy_td_ode_rhs',
'<string>', 'exec')
exec(code, globals())
config.tdfunc = cy_td_ode_rhs
if verbose:
print('BR compile time:', time.time()-_st)
initial_vector = mat2vec(rho0.full()).ravel()
_ode = scipy.integrate.ode(config.tdfunc)
code = compile('_ode.set_f_params(' + parameter_string + ')',
'<string>', 'exec')
_ode.set_integrator('zvode', method=options.method,
order=options.order, atol=options.atol,
rtol=options.rtol, nsteps=options.nsteps,
first_step=options.first_step,
min_step=options.min_step,
max_step=options.max_step)
_ode.set_initial_value(initial_vector, tlist[0])
exec(code, locals())
#
# prepare output array
#
n_tsteps = len(tlist)
e_sops_data = []
output = Result()
output.solver = "brmesolve"
output.times = tlist
if options.store_states:
output.states = []
if isinstance(e_ops, types.FunctionType):
n_expt_op = 0
expt_callback = True
elif isinstance(e_ops, list):
n_expt_op = len(e_ops)
expt_callback = False
if n_expt_op == 0:
# fall back on storing states
output.states = []
options.store_states = True
else:
output.expect = []
output.num_expect = n_expt_op
for op in e_ops:
e_sops_data.append(spre(op).data)
if op.isherm:
output.expect.append(np.zeros(n_tsteps))
else:
output.expect.append(np.zeros(n_tsteps, dtype=complex))
else:
raise TypeError("Expectation parameter must be a list or a function")
#
# start evolution
#
if type(progress_bar)==BaseProgressBar and verbose:
_run_time = time.time()
progress_bar.start(n_tsteps)
rho = Qobj(rho0)
dt = np.diff(tlist)
for t_idx, t in enumerate(tlist):
progress_bar.update(t_idx)
if not _ode.successful():
raise Exception("ODE integration error: Try to increase "
"the allowed number of substeps by increasing "
"the nsteps parameter in the Options class.")
if options.store_states or expt_callback:
rho.data = dense2D_to_fastcsr_fmode(vec2mat(_ode.y), rho.shape[0], rho.shape[1])
if options.store_states:
output.states.append(Qobj(rho, isherm=True))
if expt_callback:
# use callback method
e_ops(t, rho)
for m in range(n_expt_op):
if output.expect[m].dtype == complex:
output.expect[m][t_idx] = expect_rho_vec(e_sops_data[m],
_ode.y, 0)
else:
output.expect[m][t_idx] = expect_rho_vec(e_sops_data[m],
_ode.y, 1)
if t_idx < n_tsteps - 1:
_ode.integrate(_ode.t + dt[t_idx])
progress_bar.finished()
if type(progress_bar)==BaseProgressBar and verbose:
print('BR runtime:', time.time()-_run_time)
if (not options.rhs_reuse) and (config.tdname is not None):
_cython_build_cleanup(config.tdname)
if options.store_final_state:
rho.data = dense2D_to_fastcsr_fmode(vec2mat(_ode.y), rho.shape[0], rho.shape[1])
output.final_state = Qobj(rho, dims=rho0.dims, isherm=True)
return output
def _generic_ode_solve(r, rho0, tlist, e_ops, opt, progress_bar):
"""
Internal function for solving ME. Solve an ODE which solver parameters
already setup (r). Calculate the required expectation values or invoke
callback function at each time step.
"""
#
# prepare output array
#
n_tsteps = len(tlist)
e_sops_data = []
output = Result()
output.solver = "mesolve"
output.times = tlist
if opt.store_states:
output.states = []
if isinstance(e_ops, types.FunctionType):
n_expt_op = 0
expt_callback = True
elif isinstance(e_ops, list):
n_expt_op = len(e_ops)
expt_callback = False
if n_expt_op == 0:
# fall back on storing states
output.states = []
opt.store_states = True
else:
output.expect = []
output.num_expect = n_expt_op
for op in e_ops:
e_sops_data.append(spre(op).data)
if op.isherm and rho0.isherm:
output.expect.append(np.zeros(n_tsteps))
else:
output.expect.append(np.zeros(n_tsteps, dtype=complex))
else:
raise TypeError("Expectation parameter must be a list or a function")
#
# start evolution
#
progress_bar.start(n_tsteps)
rho = Qobj(rho0)
dt = np.diff(tlist)
for t_idx, t in enumerate(tlist):
progress_bar.update(t_idx)
if not r.successful():
raise Exception("ODE integration error: Try to increase "
"the allowed number of substeps by increasing "
"the nsteps parameter in the Options class.")
if opt.store_states or expt_callback:
rho.data = dense2D_to_fastcsr_fmode(vec2mat(r.y), rho.shape[0],
rho.shape[1])
if opt.store_states:
output.states.append(Qobj(rho, isherm=True))
if expt_callback:
# use callback method
e_ops(t, rho)
for m in range(n_expt_op):
if output.expect[m].dtype == complex:
output.expect[m][t_idx] = expect_rho_vec(
e_sops_data[m], r.y, 0)
else:
output.expect[m][t_idx] = expect_rho_vec(
e_sops_data[m], r.y, 1)
if t_idx < n_tsteps - 1:
r.integrate(r.t + dt[t_idx])
progress_bar.finished()
if (not opt.rhs_reuse) and (config.tdname is not None):
_cython_build_cleanup(config.tdname)
if opt.store_final_state:
rho.data = dense2D_to_fastcsr_fmode(vec2mat(r.y), rho.shape[0],
rho.shape[1])
output.final_state = Qobj(rho, dims=rho0.dims, isherm=True)
return output
def _sesolve_list_str_td(H_list, psi0, tlist, e_ops, args, opt, progress_bar):
"""
Internal function for solving the master equation. See mesolve for usage.
"""
if debug:
print(inspect.stack()[0][3])
#
# check initial state: must be a density matrix
#
if not isket(psi0):
raise TypeError("The unitary solver requires a ket as initial state")
#
# construct liouvillian
#
Ldata = []
Linds = []
Lptrs = []
Lcoeff = []
Lobj = []
# loop over all hamiltonian terms, convert to superoperator form and
# add the data of sparse matrix representation to h_coeff
for h_spec in H_list:
if isinstance(h_spec, Qobj):
h = h_spec
h_coeff = "1.0"
elif isinstance(h_spec, list):
h = h_spec[0]
h_coeff = h_spec[1]
else:
raise TypeError("Incorrect specification of time-dependent " +
"Hamiltonian (expected string format)")
L = -1j * h
Ldata.append(L.data.data)
Linds.append(L.data.indices)
Lptrs.append(L.data.indptr)
if isinstance(h_coeff, Cubic_Spline):
Lobj.append(h_coeff.coeffs)
Lcoeff.append(h_coeff)
# the total number of liouvillian terms (hamiltonian terms +
# collapse operators)
n_L_terms = len(Ldata)
# Check which components should use OPENMP
omp_components = None
if qset.has_openmp:
if opt.use_openmp:
omp_components = openmp_components(Lptrs)
#
# setup ode args string: we expand the list Ldata, Linds and Lptrs into
# and explicit list of parameters
#
string_list = []
for k in range(n_L_terms):
string_list.append("Ldata[%d], Linds[%d], Lptrs[%d]" % (k, k, k))
# Add object terms to end of ode args string
for k in range(len(Lobj)):
string_list.append("Lobj[%d]" % k)
for name, value in args.items():
if isinstance(value, np.ndarray):
string_list.append(name)
else:
string_list.append(str(value))
parameter_string = ",".join(string_list)
#
# generate and compile new cython code if necessary
#
if not opt.rhs_reuse or config.tdfunc is None:
if opt.rhs_filename is None:
config.tdname = "rhs" + str(os.getpid()) + str(config.cgen_num)
else:
config.tdname = opt.rhs_filename
cgen = Codegen(h_terms=n_L_terms,
h_tdterms=Lcoeff,
args=args,
config=config,
use_openmp=opt.use_openmp,
omp_components=omp_components,
omp_threads=opt.openmp_threads)
cgen.generate(config.tdname + ".pyx")
code = compile('from ' + config.tdname + ' import cy_td_ode_rhs',
'<string>', 'exec')
exec(code, globals())
config.tdfunc = cy_td_ode_rhs
#
# setup integrator
#
initial_vector = psi0.full().ravel()
r = scipy.integrate.ode(config.tdfunc)
r.set_integrator('zvode',
method=opt.method,
order=opt.order,
atol=opt.atol,
rtol=opt.rtol,
nsteps=opt.nsteps,
first_step=opt.first_step,
min_step=opt.min_step,
max_step=opt.max_step)
r.set_initial_value(initial_vector, tlist[0])
code = compile('r.set_f_params(' + parameter_string + ')', '<string>',
'exec')
exec(code, locals(), args)
# Remove RHS cython file if necessary
if not opt.rhs_reuse and config.tdname:
_cython_build_cleanup(config.tdname)
#
# call generic ODE code
#
return _generic_ode_solve(r,
psi0,
tlist,
e_ops,
opt,
progress_bar,
dims=psi0.dims)
def _sesolve_list_str_td(H_list, psi0, tlist, e_ops, args, opt,
progress_bar):
"""
Internal function for solving the master equation. See mesolve for usage.
"""
if debug:
print(inspect.stack()[0][3])
if psi0.isket:
oper_evo = False
elif psi0.isunitary:
oper_evo = True
else:
raise TypeError("The unitary solver requires psi0 to be"
" a ket as initial state"
" or a unitary as initial operator.")
#
# construct dynamics generator
#
Ldata = []
Linds = []
Lptrs = []
Lcoeff = []
Lobj = []
# loop over all hamiltonian terms, convert to superoperator form and
# add the data of sparse matrix representation to h_coeff
for h_spec in H_list:
if isinstance(h_spec, Qobj):
h = h_spec
h_coeff = "1.0"
elif isinstance(h_spec, list):
h = h_spec[0]
h_coeff = h_spec[1]
else:
raise TypeError("Incorrect specification of time-dependent " +
"Hamiltonian (expected string format)")
L = -1j * h
Ldata.append(L.data.data)
Linds.append(L.data.indices)
Lptrs.append(L.data.indptr)
if isinstance(h_coeff, Cubic_Spline):
Lobj.append(h_coeff.coeffs)
Lcoeff.append(h_coeff)
# the total number of Hamiltonian terms
n_L_terms = len(Ldata)
# Check which components should use OPENMP
omp_components = None
if qset.has_openmp:
if opt.use_openmp:
omp_components = openmp_components(Lptrs)
#
# setup ode args string: we expand the list Ldata, Linds and Lptrs into
# and explicit list of parameters
#
string_list = []
for k in range(n_L_terms):
string_list.append("Ldata[%d], Linds[%d], Lptrs[%d]" % (k, k, k))
# Add object terms to end of ode args string
for k in range(len(Lobj)):
string_list.append("Lobj[%d]" % k)
for name, value in args.items():
if isinstance(value, np.ndarray):
string_list.append(name)
else:
string_list.append(str(value))
parameter_string = ",".join(string_list)
#
# generate and compile new cython code if necessary
#
if not opt.rhs_reuse or config.tdfunc is None:
if opt.rhs_filename is None:
config.tdname = "rhs" + str(os.getpid()) + str(config.cgen_num)
else:
config.tdname = opt.rhs_filename
cgen = Codegen(h_terms=n_L_terms, h_tdterms=Lcoeff, args=args,
config=config, use_openmp=opt.use_openmp,
omp_components=omp_components,
omp_threads=opt.openmp_threads)
cgen.generate(config.tdname + ".pyx")
code = compile('from ' + config.tdname + ' import cy_td_ode_rhs',
'<string>', 'exec')
exec(code, globals())
config.tdfunc = cy_td_ode_rhs
#
# setup integrator
#
if oper_evo:
initial_vector = psi0.full().ravel('F')
r = scipy.integrate.ode(_td_ode_rhs_oper)
code = compile('r.set_f_params([' + parameter_string + '])',
'<string>', 'exec')
else:
initial_vector = psi0.full().ravel()
r = scipy.integrate.ode(config.tdfunc)
code = compile('r.set_f_params(' + parameter_string + ')',
'<string>', 'exec')
r.set_integrator('zvode', method=opt.method, order=opt.order,
atol=opt.atol, rtol=opt.rtol, nsteps=opt.nsteps,
first_step=opt.first_step, min_step=opt.min_step,
max_step=opt.max_step)
r.set_initial_value(initial_vector, tlist[0])
exec(code, locals(), args)
# Remove RHS cython file if necessary
if not opt.rhs_reuse and config.tdname:
_cython_build_cleanup(config.tdname)
#
# call generic ODE code
#
return _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar,
dims=psi0.dims)
def _correlation_mc_2t(H, state0, tlist, taulist, c_ops, a_op, b_op, c_op,
args={}, options=Options()):
"""
Internal function for calculating the three-operator two-time
correlation function:
<A(t)B(t+tau)C(t)>
using a Monte Carlo solver.
"""
if not c_ops:
raise TypeError("If no collapse operators are required, use the `me`" +
"or `es` solvers")
# the solvers only work for positive time differences and the correlators
# require positive tau
if state0 is None:
raise NotImplementedError("steady state not implemented for " +
"mc solver, please use `es` or `me`")
elif not isket(state0):
raise TypeError("state0 must be a state vector.")
psi0 = state0
if debug:
print(inspect.stack()[0][3])
psi_t_mat = mcsolve(
H, psi0, tlist, c_ops, [],
args=args, ntraj=options.ntraj[0], options=options, progress_bar=None
).states
corr_mat = np.zeros([np.size(tlist), np.size(taulist)], dtype=complex)
H_shifted, c_ops_shifted, _args = _transform_L_t_shift(H, c_ops, args)
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
# calculation of <A(t)B(t+tau)C(t)> from only knowledge of psi0 requires
# averaging over both t and tau
for t_idx in range(np.size(tlist)):
if not isinstance(H, Qobj):
_args["_t0"] = tlist[t_idx]
for trial_idx in range(options.ntraj[0]):
if isinstance(a_op, Qobj) and isinstance(c_op, Qobj):
if a_op.dag() == c_op:
# A shortcut here, requires only 1/4 the trials
chi_0 = (options.mc_corr_eps + c_op) * \
psi_t_mat[trial_idx, t_idx]
# evolve these states and calculate expectation value of B
c_tau = chi_0.norm()**2 * mcsolve(
H_shifted, chi_0/chi_0.norm(), taulist, c_ops_shifted,
[b_op],
args=_args, ntraj=options.ntraj[1], options=options,
progress_bar=None
).expect[0]
# final correlation vector computed by combining the
# averages
corr_mat[t_idx, :] += c_tau/options.ntraj[1]
else:
# otherwise, need four trial wavefunctions
# (Ad+C)*psi_t, (Ad+iC)*psi_t, (Ad-C)*psi_t, (Ad-iC)*psi_t
if isinstance(a_op, Qobj):
a_op_dag = a_op.dag()
else:
# assume this is a number, ex. i.e. a_op = 1
# if this is not correct, the over-loaded addition
# operation will raise errors
a_op_dag = a_op
chi_0 = [(options.mc_corr_eps + a_op_dag +
np.exp(1j*x*np.pi/2)*c_op) *
psi_t_mat[trial_idx, t_idx]
for x in range(4)]
# evolve these states and calculate expectation value of B
c_tau = [
chi.norm()**2 * mcsolve(
H_shifted, chi/chi.norm(), taulist, c_ops_shifted,
[b_op],
args=_args, ntraj=options.ntraj[1], options=options,
progress_bar=None
).expect[0]
for chi in chi_0
]
# final correlation vector computed by combining the averages
corr_mat_add = np.asarray(
1.0 / (4*options.ntraj[0]) *
(c_tau[0] - c_tau[2] - 1j*c_tau[1] + 1j*c_tau[3]),
dtype=corr_mat.dtype
)
corr_mat[t_idx, :] += corr_mat_add
if t_idx == 1:
options.rhs_reuse = True
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
return corr_mat
def _parallel_sesolve(n, N, H, tlist, args, options):
psi0 = basis(N, n)
output = sesolve(H, psi0, tlist, [], args, options, _safe_mode=False)
if config.tdname:
_cython_build_cleanup(config.tdname)
return output
def generic_ode_solve_checkpoint(r, rho0, tlist, e_ops, opt, progress_bar,
save, subdir):
"""
Internal function for solving ME. Solve an ODE which solver parameters
already setup (r). Calculate the required expectation values or invoke
callback function at each time step.
"""
#
# prepare output array
#
n_tsteps = len(tlist)
e_sops_data = []
output = Result()
output.solver = "mesolve"
output.times = tlist
if opt.store_states:
output.states = []
e_ops_dict = e_ops
e_ops = [e for e in e_ops_dict.values()]
headings = [key for key in e_ops_dict.keys()]
if isinstance(e_ops, types.FunctionType):
n_expt_op = 0
expt_callback = True
elif isinstance(e_ops, list):
n_expt_op = len(e_ops)
expt_callback = False
if n_expt_op == 0:
# fall back on storing states
output.states = []
opt.store_states = True
else:
output.expect = []
output.num_expect = n_expt_op
for op in e_ops:
e_sops_data.append(spre(op).data)
if op.isherm and rho0.isherm:
output.expect.append(np.zeros(n_tsteps))
else:
output.expect.append(np.zeros(n_tsteps, dtype=complex))
else:
raise TypeError("Expectation parameter must be a list or a function")
results_row = np.zeros(n_expt_op)
#
# start evolution
#
progress_bar.start(n_tsteps)
rho = Qobj(rho0)
dims = rho.dims
dt = np.diff(tlist)
end_time = tlist[-1]
for t_idx, t in tqdm(enumerate(tlist)):
progress_bar.update(t_idx)
if not r.successful():
raise Exception("ODE integration error: Try to increase "
"the allowed number of substeps by increasing "
"the nsteps parameter in the Options class.")
if opt.store_states or expt_callback:
rho.data = vec2mat(r.y)
if opt.store_states:
output.states.append(Qobj(rho))
if expt_callback:
# use callback method
e_ops(t, rho)
for m in range(n_expt_op):
if output.expect[m].dtype == complex:
output.expect[m][t_idx] = expect_rho_vec(
e_sops_data[m], r.y, 0)
results_row[m] = output.expect[m][t_idx]
else:
output.expect[m][t_idx] = expect_rho_vec(
e_sops_data[m], r.y, 1)
results_row[m] = output.expect[m][t_idx]
results = pd.DataFrame(results_row).T
results.columns = headings
results.index = [t]
results.index.name = 'times'
if t == 0:
first_row = True
else:
first_row = False
if save:
rho_checkpoint = Qobj(vec2mat(r.y))
rho_checkpoint.dims = dims
if t_idx % 200 == 0:
rho_c = rho_checkpoint.ptrace(0)
with open('./cavity_states.pkl', 'ab') as f:
pickle.dump(rho_c, f)
with open('./results.csv', 'a') as file:
results.to_csv(file, header=first_row, float_format='%.15f')
qsave(rho_checkpoint, './state_checkpoint')
save = True
if t_idx < n_tsteps - 1:
r.integrate(r.t + dt[t_idx])
progress_bar.finished()
if not opt.rhs_reuse and config.tdname is not None:
_cython_build_cleanup(config.tdname)
return output
def _correlation_mc_2t(H, state0, tlist, taulist, c_ops, a_op, b_op, c_op,
args={}, options=Options()):
"""
Internal function for calculating the three-operator two-time
correlation function:
<A(t)B(t+tau)C(t)>
using a Monte Carlo solver.
"""
if not c_ops:
raise TypeError("If no collapse operators are required, use the `me`" +
"or `es` solvers")
# the solvers only work for positive time differences and the correlators
# require positive tau
if state0 is None:
raise NotImplementedError("steady state not implemented for " +
"mc solver, please use `es` or `me`")
elif not isket(state0):
raise TypeError("state0 must be a state vector.")
psi0 = state0
if debug:
print(inspect.stack()[0][3])
psi_t_mat = mcsolve(
H, psi0, tlist, c_ops, [],
args=args, ntraj=options.ntraj[0], options=options, progress_bar=None
).states
corr_mat = np.zeros([np.size(tlist), np.size(taulist)], dtype=complex)
H_shifted, c_ops_shifted, _args = _transform_L_t_shift_new(H, c_ops, args)
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
# calculation of <A(t)B(t+tau)C(t)> from only knowledge of psi0 requires
# averaging over both t and tau
for t_idx in range(np.size(tlist)):
if not isinstance(H, Qobj):
_args["_t0"] = tlist[t_idx]
for trial_idx in range(options.ntraj[0]):
if isinstance(a_op, Qobj) and isinstance(c_op, Qobj):
if a_op.dag() == c_op:
# A shortcut here, requires only 1/4 the trials
chi_0 = (options.mc_corr_eps + c_op) * \
psi_t_mat[trial_idx, t_idx]
# evolve these states and calculate expectation value of B
c_tau = chi_0.norm()**2 * mcsolve(
H_shifted, chi_0/chi_0.norm(), taulist, c_ops_shifted,
[b_op],
args=_args, ntraj=options.ntraj[1], options=options,
progress_bar=None
).expect[0]
# final correlation vector computed by combining the
# averages
corr_mat[t_idx, :] += c_tau/options.ntraj[1]
else:
# otherwise, need four trial wavefunctions
# (Ad+C)*psi_t, (Ad+iC)*psi_t, (Ad-C)*psi_t, (Ad-iC)*psi_t
if isinstance(a_op, Qobj):
a_op_dag = a_op.dag()
else:
# assume this is a number, ex. i.e. a_op = 1
# if this is not correct, the over-loaded addition
# operation will raise errors
a_op_dag = a_op
chi_0 = [(options.mc_corr_eps + a_op_dag +
np.exp(1j*x*np.pi/2)*c_op) *
psi_t_mat[trial_idx, t_idx]
for x in range(4)]
# evolve these states and calculate expectation value of B
c_tau = [
chi.norm()**2 * mcsolve(
H_shifted, chi/chi.norm(), taulist, c_ops_shifted,
[b_op],
args=_args, ntraj=options.ntraj[1], options=options,
progress_bar=None
).expect[0]
for chi in chi_0
]
# final correlation vector computed by combining the averages
corr_mat_add = np.asarray(
1.0 / (4*options.ntraj[0]) *
(c_tau[0] - c_tau[2] - 1j*c_tau[1] + 1j*c_tau[3]),
dtype=corr_mat.dtype
)
corr_mat[t_idx, :] += corr_mat_add
if t_idx == 1:
options.rhs_reuse = True
if config.tdname:
_cython_build_cleanup(config.tdname)
rhs_clear()
return corr_mat
def propagator(H,
t,
c_op_list=[],
args={},
options=None,
unitary_mode='batch',
parallel=False,
progress_bar=None,
**kwargs):
"""
Calculate the propagator U(t) for the density matrix or wave function such
that :math:`\psi(t) = U(t)\psi(0)` or
:math:`\\rho_{\mathrm vec}(t) = U(t) \\rho_{\mathrm vec}(0)`
where :math:`\\rho_{\mathrm vec}` is the vector representation of the
density matrix.
Parameters
----------
H : qobj or list
Hamiltonian as a Qobj instance of a nested list of Qobjs and
coefficients in the list-string or list-function format for
time-dependent Hamiltonians (see description in :func:`qutip.mesolve`).
t : float or array-like
Time or list of times for which to evaluate the propagator.
c_op_list : list
List of qobj collapse operators.
args : list/array/dictionary
Parameters to callback functions for time-dependent Hamiltonians and
collapse operators.
options : :class:`qutip.Options`
with options for the ODE solver.
unitary_mode = str ('batch', 'single')
Solve all basis vectors simulaneously ('batch') or individually
('single').
parallel : bool {False, True}
Run the propagator in parallel mode. This will override the
unitary_mode settings if set to True.
progress_bar: BaseProgressBar
Optional instance of BaseProgressBar, or a subclass thereof, for
showing the progress of the simulation. By default no progress bar
is used, and if set to True a TextProgressBar will be used.
Returns
-------
a : qobj
Instance representing the propagator :math:`U(t)`.
"""
kw = _default_kwargs()
if 'num_cpus' in kwargs:
num_cpus = kwargs['num_cpus']
else:
num_cpus = kw['num_cpus']
if progress_bar is None:
progress_bar = BaseProgressBar()
elif progress_bar is True:
progress_bar = TextProgressBar()
if options is None:
options = Options()
options.rhs_reuse = True
rhs_clear()
if isinstance(t, (int, float, np.integer, np.floating)):
tlist = [0, t]
else:
tlist = t
td_type = _td_format_check(H, c_op_list, solver='me')
if isinstance(
H,
(types.FunctionType, types.BuiltinFunctionType, functools.partial)):
H0 = H(0.0, args)
elif isinstance(H, list):
H0 = H[0][0] if isinstance(H[0], list) else H[0]
else:
H0 = H
if len(c_op_list) == 0 and H0.isoper:
# calculate propagator for the wave function
N = H0.shape[0]
dims = H0.dims
if parallel:
unitary_mode = 'single'
u = np.zeros([N, N, len(tlist)], dtype=complex)
output = parallel_map(_parallel_sesolve,
range(N),
task_args=(N, H, tlist, args, options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N):
for k, t in enumerate(tlist):
u[:, n, k] = output[n].states[k].full().T
else:
if unitary_mode == 'single':
u = np.zeros([N, N, len(tlist)], dtype=complex)
progress_bar.start(N)
for n in range(0, N):
progress_bar.update(n)
psi0 = basis(N, n)
output = sesolve(H,
psi0,
tlist, [],
args,
options,
_safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = output.states[k].full().T
progress_bar.finished()
if config.tdname:
_cython_build_cleanup(config.tdname)
elif unitary_mode == 'batch':
u = np.zeros(len(tlist), dtype=object)
_rows = np.array([(N + 1) * m for m in range(N)])
_cols = np.zeros_like(_rows)
_data = np.ones_like(_rows, dtype=complex)
psi0 = Qobj(sp.coo_matrix((_data, (_rows, _cols))).tocsr())
if td_type[1] > 0 or td_type[2] > 0:
H2 = []
for k in range(len(H)):
if isinstance(H[k], list):
H2.append([tensor(qeye(N), H[k][0]), H[k][1]])
else:
H2.append(tensor(qeye(N), H[k]))
else:
H2 = tensor(qeye(N), H)
output = sesolve(H2,
psi0,
tlist, [],
args=args,
_safe_mode=False,
options=Options(normalize_output=False))
for k, t in enumerate(tlist):
u[k] = sp_reshape(output.states[k].data, (N, N))
unit_row_norm(u[k].data, u[k].indptr, u[k].shape[0])
u[k] = u[k].T.tocsr()
if config.tdname:
_cython_build_cleanup(config.tdname)
else:
raise Exception('Invalid unitary mode.')
elif len(c_op_list) == 0 and H0.issuper:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
unitary_mode = 'single'
N = H0.shape[0]
sqrt_N = int(np.sqrt(N))
dims = H0.dims
u = np.zeros([N, N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,
range(N * N),
task_args=(sqrt_N, H, tlist, c_op_list, args,
options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N)
for n in range(0, N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n, (sqrt_N, sqrt_N))
rho0 = Qobj(
sp.csr_matrix(([1], ([row_idx], [col_idx])),
shape=(sqrt_N, sqrt_N),
dtype=complex))
output = mesolve(H,
rho0,
tlist, [], [],
args,
options,
_safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
else:
# calculate the propagator for the vector representation of the
# density matrix (a superoperator propagator)
unitary_mode = 'single'
N = H0.shape[0]
dims = [H0.dims, H0.dims]
u = np.zeros([N * N, N * N, len(tlist)], dtype=complex)
if parallel:
output = parallel_map(_parallel_mesolve,
range(N * N),
task_args=(N, H, tlist, c_op_list, args,
options),
progress_bar=progress_bar,
num_cpus=num_cpus)
for n in range(N * N):
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output[n].states[k].full()).T
else:
progress_bar.start(N * N)
for n in range(N * N):
progress_bar.update(n)
col_idx, row_idx = np.unravel_index(n, (N, N))
rho0 = Qobj(
sp.csr_matrix(([1], ([row_idx], [col_idx])),
shape=(N, N),
dtype=complex))
output = mesolve(H,
rho0,
tlist,
c_op_list, [],
args,
options,
_safe_mode=False)
for k, t in enumerate(tlist):
u[:, n, k] = mat2vec(output.states[k].full()).T
progress_bar.finished()
if len(tlist) == 2:
if unitary_mode == 'batch':
return Qobj(u[-1], dims=dims)
else:
return Qobj(u[:, :, 1], dims=dims)
else:
if unitary_mode == 'batch':
return np.array([Qobj(u[k], dims=dims) for k in range(len(tlist))],
dtype=object)
else:
return np.array(
[Qobj(u[:, :, k], dims=dims) for k in range(len(tlist))],
dtype=object)
def mcsolve(H, psi0, tlist, c_ops, e_ops, ntraj=None,
args={}, options=None, progress_bar=True,
map_func=None, map_kwargs=None):
"""Monte Carlo evolution of a state vector :math:`|\psi \\rangle` for a
given Hamiltonian and sets of collapse operators, and possibly, operators
for calculating expectation values. Options for the underlying ODE solver
are given by the Options class.
mcsolve supports time-dependent Hamiltonians and collapse operators using
either Python functions of strings to represent time-dependent
coefficients. Note that, the system Hamiltonian MUST have at least one
constant term.
As an example of a time-dependent problem, consider a Hamiltonian with two
terms ``H0`` and ``H1``, where ``H1`` is time-dependent with coefficient
``sin(w*t)``, and collapse operators ``C0`` and ``C1``, where ``C1`` is
time-dependent with coeffcient ``exp(-a*t)``. Here, w and a are constant
arguments with values ``W`` and ``A``.
Using the Python function time-dependent format requires two Python
functions, one for each collapse coefficient. Therefore, this problem could
be expressed as::
def H1_coeff(t,args):
return sin(args['w']*t)
def C1_coeff(t,args):
return exp(-args['a']*t)
H = [H0, [H1, H1_coeff]]
c_ops = [C0, [C1, C1_coeff]]
args={'a': A, 'w': W}
or in String (Cython) format we could write::
H = [H0, [H1, 'sin(w*t)']]
c_ops = [C0, [C1, 'exp(-a*t)']]
args={'a': A, 'w': W}
Constant terms are preferably placed first in the Hamiltonian and collapse
operator lists.
Parameters
----------
H : :class:`qutip.Qobj`
System Hamiltonian.
psi0 : :class:`qutip.Qobj`
Initial state vector
tlist : array_like
Times at which results are recorded.
ntraj : int
Number of trajectories to run.
c_ops : array_like
single collapse operator or ``list`` or ``array`` of collapse
operators.
e_ops : array_like
single operator or ``list`` or ``array`` of operators for calculating
expectation values.
args : dict
Arguments for time-dependent Hamiltonian and collapse operator terms.
options : Options
Instance of ODE solver options.
progress_bar: BaseProgressBar
Optional instance of BaseProgressBar, or a subclass thereof, for
showing the progress of the simulation. Set to None to disable the
progress bar.
map_func: function
A map function for managing the calls to the single-trajactory solver.
map_kwargs: dictionary
Optional keyword arguments to the map_func function.
Returns
-------
results : :class:`qutip.solver.Result`
Object storing all results from the simulation.
.. note::
It is possible to reuse the random number seeds from a previous run
of the mcsolver by passing the output Result object seeds via the
Options class, i.e. Options(seeds=prev_result.seeds).
"""
if debug:
print(inspect.stack()[0][3])
if options is None:
options = Options()
if ntraj is None:
ntraj = options.ntraj
config.map_func = map_func if map_func is not None else parallel_map
config.map_kwargs = map_kwargs if map_kwargs is not None else {}
if not psi0.isket:
raise Exception("Initial state must be a state vector.")
if isinstance(c_ops, Qobj):
c_ops = [c_ops]
if isinstance(e_ops, Qobj):
e_ops = [e_ops]
if isinstance(e_ops, dict):
e_ops_dict = e_ops
e_ops = [e for e in e_ops.values()]
else:
e_ops_dict = None
config.options = options
if progress_bar:
if progress_bar is True:
config.progress_bar = TextProgressBar()
else:
config.progress_bar = progress_bar
else:
config.progress_bar = BaseProgressBar()
# set num_cpus to the value given in qutip.settings if none in Options
if not config.options.num_cpus:
config.options.num_cpus = qutip.settings.num_cpus
if config.options.num_cpus == 1:
# fallback on serial_map if num_cpu == 1, since there is no
# benefit of starting multiprocessing in this case
config.map_func = serial_map
# set initial value data
if options.tidy:
config.psi0 = psi0.tidyup(options.atol).full().ravel()
else:
config.psi0 = psi0.full().ravel()
config.psi0_dims = psi0.dims
config.psi0_shape = psi0.shape
# set options on ouput states
if config.options.steady_state_average:
config.options.average_states = True
# set general items
config.tlist = tlist
if isinstance(ntraj, (list, np.ndarray)):
config.ntraj = np.sort(ntraj)[-1]
else:
config.ntraj = ntraj
# set norm finding constants
config.norm_tol = options.norm_tol
config.norm_steps = options.norm_steps
# convert array based time-dependence to string format
H, c_ops, args = _td_wrap_array_str(H, c_ops, args, tlist)
# SETUP ODE DATA IF NONE EXISTS OR NOT REUSING
# --------------------------------------------
if not options.rhs_reuse or not config.tdfunc:
# reset config collapse and time-dependence flags to default values
config.soft_reset()
# check for type of time-dependence (if any)
time_type, h_stuff, c_stuff = _td_format_check(H, c_ops, 'mc')
c_terms = len(c_stuff[0]) + len(c_stuff[1]) + len(c_stuff[2])
# set time_type for use in multiprocessing
config.tflag = time_type
# check for collapse operators
if c_terms > 0:
config.cflag = 1
else:
config.cflag = 0
# Configure data
_mc_data_config(H, psi0, h_stuff, c_ops, c_stuff, args, e_ops,
options, config)
# compile and load cython functions if necessary
_mc_func_load(config)
else:
# setup args for new parameters when rhs_reuse=True and tdfunc is given
# string based
if config.tflag in [1, 10, 11]:
if any(args):
config.c_args = []
arg_items = list(args.items())
for k in range(len(arg_items)):
config.c_args.append(arg_items[k][1])
# function based
elif config.tflag in [2, 3, 20, 22]:
config.h_func_args = args
# load monte carlo class
mc = _MC(config)
# Run the simulation
mc.run()
# Remove RHS cython file if necessary
if not options.rhs_reuse and config.tdname:
_cython_build_cleanup(config.tdname)
# AFTER MCSOLVER IS DONE
# ----------------------
# Store results in the Result object
output = Result()
output.solver = 'mcsolve'
output.seeds = config.options.seeds
# state vectors
if (mc.psi_out is not None and config.options.average_states
and config.cflag and ntraj != 1):
output.states = parfor(_mc_dm_avg, mc.psi_out.T)
elif mc.psi_out is not None:
output.states = mc.psi_out
# expectation values
if (mc.expect_out is not None and config.cflag
and config.options.average_expect):
# averaging if multiple trajectories
if isinstance(ntraj, int):
output.expect = [np.mean(np.array([mc.expect_out[nt][op]
for nt in range(ntraj)],
dtype=object),
axis=0)
for op in range(config.e_num)]
elif isinstance(ntraj, (list, np.ndarray)):
output.expect = []
for num in ntraj:
expt_data = np.mean(mc.expect_out[:num], axis=0)
data_list = []
if any([not op.isherm for op in e_ops]):
for k in range(len(e_ops)):
if e_ops[k].isherm:
data_list.append(np.real(expt_data[k]))
else:
data_list.append(expt_data[k])
else:
data_list = [data for data in expt_data]
output.expect.append(data_list)
else:
# no averaging for single trajectory or if average_expect flag
# (Options) is off
if mc.expect_out is not None:
output.expect = mc.expect_out
# simulation parameters
output.times = config.tlist
output.num_expect = config.e_num
output.num_collapse = config.c_num
output.ntraj = config.ntraj
output.col_times = mc.collapse_times_out
output.col_which = mc.which_op_out
if e_ops_dict:
output.expect = {e: output.expect[n]
for n, e in enumerate(e_ops_dict.keys())}
return output
def _td_brmesolve(H,
psi0,
tlist,
a_ops=[],
e_ops=[],
c_ops=[],
use_secular=True,
tol=qset.atol,
options=None,
progress_bar=None,
_safe_mode=True):
if isket(psi0):
rho0 = ket2dm(psi0)
else:
rho0 = psi0
nrows = rho0.shape[0]
H_terms = []
H_td_terms = []
H_obj = []
A_terms = []
A_td_terms = []
C_terms = []
C_td_terms = []
C_obj = []
spline_count = [0, 0]
if isinstance(H, Qobj):
H_terms.append(H.full('f'))
H_td_terms.append('1')
else:
for kk, h in enumerate(H):
if isinstance(h, Qobj):
H_terms.append(h.full('f'))
H_td_terms.append('1')
elif isinstance(h, list):
H_terms.append(h[0].full('f'))
if isinstance(h[1], Cubic_Spline):
H_obj.append(h[1].coeffs)
spline_count[0] += 1
H_td_terms.append(h[1])
else:
raise Exception('Invalid Hamiltonian specifiction.')
for kk, c in enumerate(c_ops):
if isinstance(c, Qobj):
C_terms.append(c.full('f'))
C_td_terms.append('1')
elif isinstance(c, list):
C_terms.append(c[0].full('f'))
if isinstance(c[1], Cubic_Spline):
C_obj.append(c[1].coeffs)
spline_count[0] += 1
C_td_terms.append(c[1])
else:
raise Exception('Invalid collape operator specifiction.')
for kk, a in enumerate(a_ops):
if isinstance(a, list):
A_terms.append(a[0].full('f'))
A_td_terms.append(a[1])
if isinstance(a[1], tuple):
if not len(a[1]) == 2:
raise Exception('Tuple must be len=2.')
if isinstance(a[1][0], Cubic_Spline):
spline_count[1] += 1
if isinstance(a[1][1], Cubic_Spline):
spline_count[1] += 1
else:
raise Exception('Invalid bath-coupling specifiction.')
string_list = []
for kk, _ in enumerate(H_td_terms):
string_list.append("H_terms[{0}]".format(kk))
for kk, _ in enumerate(H_obj):
string_list.append("H_obj[{0}]".format(kk))
for kk, _ in enumerate(C_td_terms):
string_list.append("C_terms[{0}]".format(kk))
for kk, _ in enumerate(C_obj):
string_list.append("C_obj[{0}]".format(kk))
for kk, _ in enumerate(A_td_terms):
string_list.append("A_terms[{0}]".format(kk))
#Add nrows to parameters
string_list.append('nrows')
parameter_string = ",".join(string_list)
#
# generate and compile new cython code if necessary
#
if not options.rhs_reuse or config.tdfunc is None:
if options.rhs_filename is None:
config.tdname = "rhs" + str(os.getpid()) + str(config.cgen_num)
else:
config.tdname = opt.rhs_filename
cgen = BR_Codegen(
h_terms=len(H_terms),
h_td_terms=H_td_terms,
h_obj=H_obj,
c_terms=len(C_terms),
c_td_terms=C_td_terms,
c_obj=C_obj,
a_terms=len(A_terms),
a_td_terms=A_td_terms,
spline_count=spline_count,
config=config,
sparse=False,
use_secular=use_secular,
use_openmp=options.use_openmp,
omp_thresh=qset.openmp_thresh if qset.has_openmp else None,
omp_threads=options.num_cpus,
atol=tol)
cgen.generate(config.tdname + ".pyx")
code = compile('from ' + config.tdname + ' import cy_td_ode_rhs',
'<string>', 'exec')
exec(code, globals())
config.tdfunc = cy_td_ode_rhs
initial_vector = mat2vec(rho0.full()).ravel()
_ode = scipy.integrate.ode(config.tdfunc)
code = compile('_ode.set_f_params(' + parameter_string + ')', '<string>',
'exec')
_ode.set_integrator('zvode',
method=options.method,
order=options.order,
atol=options.atol,
rtol=options.rtol,
nsteps=options.nsteps,
first_step=options.first_step,
min_step=options.min_step,
max_step=options.max_step)
_ode.set_initial_value(initial_vector, tlist[0])
exec(code, locals())
#
# prepare output array
#
n_tsteps = len(tlist)
e_sops_data = []
output = Result()
output.solver = "brmesolve"
output.times = tlist
if options.store_states:
output.states = []
if isinstance(e_ops, types.FunctionType):
n_expt_op = 0
expt_callback = True
elif isinstance(e_ops, list):
n_expt_op = len(e_ops)
expt_callback = False
if n_expt_op == 0:
# fall back on storing states
output.states = []
options.store_states = True
else:
output.expect = []
output.num_expect = n_expt_op
for op in e_ops:
e_sops_data.append(spre(op).data)
if op.isherm:
output.expect.append(np.zeros(n_tsteps))
else:
output.expect.append(np.zeros(n_tsteps, dtype=complex))
else:
raise TypeError("Expectation parameter must be a list or a function")
#
# start evolution
#
progress_bar.start(n_tsteps)
rho = Qobj(rho0)
dt = np.diff(tlist)
for t_idx, t in enumerate(tlist):
progress_bar.update(t_idx)
if not _ode.successful():
raise Exception("ODE integration error: Try to increase "
"the allowed number of substeps by increasing "
"the nsteps parameter in the Options class.")
if options.store_states or expt_callback:
rho.data = dense2D_to_fastcsr_fmode(vec2mat(_ode.y), rho.shape[0],
rho.shape[1])
if options.store_states:
output.states.append(Qobj(rho, isherm=True))
if expt_callback:
# use callback method
e_ops(t, rho)
for m in range(n_expt_op):
if output.expect[m].dtype == complex:
output.expect[m][t_idx] = expect_rho_vec(
e_sops_data[m], _ode.y, 0)
else:
output.expect[m][t_idx] = expect_rho_vec(
e_sops_data[m], _ode.y, 1)
if t_idx < n_tsteps - 1:
_ode.integrate(_ode.t + dt[t_idx])
progress_bar.finished()
if (not options.rhs_reuse) and (config.tdname is not None):
_cython_build_cleanup(config.tdname)
if options.store_final_state:
rho.data = dense2D_to_fastcsr_fmode(vec2mat(_ode.y), rho.shape[0],
rho.shape[1])
output.final_state = Qobj(rho, dims=rho0.dims, isherm=True)
return output
def mcsolve(H, psi0, tlist, c_ops, e_ops, ntraj=None,
args={}, options=None, progress_bar=True,
map_func=None, map_kwargs=None):
"""Monte Carlo evolution of a state vector :math:`|\psi \\rangle` for a
given Hamiltonian and sets of collapse operators, and possibly, operators
for calculating expectation values. Options for the underlying ODE solver
are given by the Options class.
mcsolve supports time-dependent Hamiltonians and collapse operators using
either Python functions of strings to represent time-dependent
coefficients. Note that, the system Hamiltonian MUST have at least one
constant term.
As an example of a time-dependent problem, consider a Hamiltonian with two
terms ``H0`` and ``H1``, where ``H1`` is time-dependent with coefficient
``sin(w*t)``, and collapse operators ``C0`` and ``C1``, where ``C1`` is
time-dependent with coeffcient ``exp(-a*t)``. Here, w and a are constant
arguments with values ``W`` and ``A``.
Using the Python function time-dependent format requires two Python
functions, one for each collapse coefficient. Therefore, this problem could
be expressed as::
def H1_coeff(t,args):
return sin(args['w']*t)
def C1_coeff(t,args):
return exp(-args['a']*t)
H = [H0, [H1, H1_coeff]]
c_ops = [C0, [C1, C1_coeff]]
args={'a': A, 'w': W}
or in String (Cython) format we could write::
H = [H0, [H1, 'sin(w*t)']]
c_ops = [C0, [C1, 'exp(-a*t)']]
args={'a': A, 'w': W}
Constant terms are preferably placed first in the Hamiltonian and collapse
operator lists.
Parameters
----------
H : :class:`qutip.Qobj`
System Hamiltonian.
psi0 : :class:`qutip.Qobj`
Initial state vector
tlist : array_like
Times at which results are recorded.
ntraj : int
Number of trajectories to run.
c_ops : array_like
single collapse operator or ``list`` or ``array`` of collapse
operators.
e_ops : array_like
single operator or ``list`` or ``array`` of operators for calculating
expectation values.
args : dict
Arguments for time-dependent Hamiltonian and collapse operator terms.
options : Options
Instance of ODE solver options.
progress_bar: BaseProgressBar
Optional instance of BaseProgressBar, or a subclass thereof, for
showing the progress of the simulation. Set to None to disable the
progress bar.
map_func: function
A map function for managing the calls to the single-trajactory solver.
map_kwargs: dictionary
Optional keyword arguments to the map_func function.
Returns
-------
results : :class:`qutip.solver.Result`
Object storing all results from the simulation.
.. note::
It is possible to reuse the random number seeds from a previous run
of the mcsolver by passing the output Result object seeds via the
Options class, i.e. Options(seeds=prev_result.seeds).
"""
if debug:
print(inspect.stack()[0][3])
if options is None:
options = Options()
if ntraj is None:
ntraj = options.ntraj
config.map_func = map_func if map_func is not None else parallel_map
config.map_kwargs = map_kwargs if map_kwargs is not None else {}
if not psi0.isket:
raise Exception("Initial state must be a state vector.")
if isinstance(c_ops, Qobj):
c_ops = [c_ops]
if isinstance(e_ops, Qobj):
e_ops = [e_ops]
if isinstance(e_ops, dict):
e_ops_dict = e_ops
e_ops = [e for e in e_ops.values()]
else:
e_ops_dict = None
config.options = options
if progress_bar:
if progress_bar is True:
config.progress_bar = TextProgressBar()
else:
config.progress_bar = progress_bar
else:
config.progress_bar = BaseProgressBar()
# set num_cpus to the value given in qutip.settings if none in Options
if not config.options.num_cpus:
config.options.num_cpus = qutip.settings.num_cpus
if config.options.num_cpus == 1:
# fallback on serial_map if num_cpu == 1, since there is no
# benefit of starting multiprocessing in this case
config.map_func = serial_map
# set initial value data
if options.tidy:
config.psi0 = psi0.tidyup(options.atol).full().ravel()
else:
config.psi0 = psi0.full().ravel()
config.psi0_dims = psi0.dims
config.psi0_shape = psi0.shape
# set options on ouput states
if config.options.steady_state_average:
config.options.average_states = True
# set general items
config.tlist = tlist
if isinstance(ntraj, (list, np.ndarray)):
config.ntraj = np.sort(ntraj)[-1]
else:
config.ntraj = ntraj
# set norm finding constants
config.norm_tol = options.norm_tol
config.norm_steps = options.norm_steps
# convert array based time-dependence to string format
H, c_ops, args = _td_wrap_array_str(H, c_ops, args, tlist)
# SETUP ODE DATA IF NONE EXISTS OR NOT REUSING
# --------------------------------------------
if not options.rhs_reuse or not config.tdfunc:
# reset config collapse and time-dependence flags to default values
config.soft_reset()
# check for type of time-dependence (if any)
time_type, h_stuff, c_stuff = _td_format_check(H, c_ops, 'mc')
c_terms = len(c_stuff[0]) + len(c_stuff[1]) + len(c_stuff[2])
# set time_type for use in multiprocessing
config.tflag = time_type
# check for collapse operators
if c_terms > 0:
config.cflag = 1
else:
config.cflag = 0
# Configure data
_mc_data_config(H, psi0, h_stuff, c_ops, c_stuff, args, e_ops,
options, config)
# compile and load cython functions if necessary
_mc_func_load(config)
else:
# setup args for new parameters when rhs_reuse=True and tdfunc is given
# string based
if config.tflag in [1, 10, 11]:
if any(args):
config.c_args = []
arg_items = list(args.items())
for k in range(len(arg_items)):
config.c_args.append(arg_items[k][1])
# function based
elif config.tflag in [2, 3, 20, 22]:
config.h_func_args = args
# load monte carlo class
mc = _MC(config)
# Run the simulation
mc.run()
# Remove RHS cython file if necessary
if not options.rhs_reuse and config.tdname:
_cython_build_cleanup(config.tdname)
# AFTER MCSOLVER IS DONE
# ----------------------
# Store results in the Result object
output = Result()
output.solver = 'mcsolve'
output.seeds = config.options.seeds
# state vectors
if (mc.psi_out is not None and config.options.average_states
and config.cflag and ntraj != 1):
output.states = parfor(_mc_dm_avg, mc.psi_out.T)
elif mc.psi_out is not None:
output.states = mc.psi_out
# expectation values
if (mc.expect_out is not None and config.cflag
and config.options.average_expect):
# averaging if multiple trajectories
if isinstance(ntraj, int):
output.expect = [np.mean(np.array([mc.expect_out[nt][op]
for nt in range(ntraj)],
dtype=object),
axis=0)
for op in range(config.e_num)]
elif isinstance(ntraj, (list, np.ndarray)):
output.expect = []
for num in ntraj:
expt_data = np.mean(mc.expect_out[:num], axis=0)
data_list = []
if any([not op.isherm for op in e_ops]):
for k in range(len(e_ops)):
if e_ops[k].isherm:
data_list.append(np.real(expt_data[k]))
else:
data_list.append(expt_data[k])
else:
data_list = [data for data in expt_data]
output.expect.append(data_list)
else:
# no averaging for single trajectory or if average_expect flag
# (Options) is off
if mc.expect_out is not None:
output.expect = mc.expect_out
# simulation parameters
output.times = config.tlist
output.num_expect = config.e_num
output.num_collapse = config.c_num
output.ntraj = config.ntraj
output.col_times = mc.collapse_times_out
output.col_which = mc.which_op_out
if e_ops_dict:
output.expect = {e: output.expect[n]
for n, e in enumerate(e_ops_dict.keys())}
return output
