def _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar, dims=None):
"""
Internal function for solving ODEs.
"""
#
# prepare output array
#
n_tsteps = len(tlist)
output = Result()
output.solver = "sesolve"
output.times = tlist
if psi0.isunitary:
oper_evo = True
oper_n = dims[0][0]
norm_dim_factor = np.sqrt(oper_n)
else:
oper_evo = False
norm_dim_factor = 1.0
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:
# fallback on storing states
output.states = []
opt.store_states = True
else:
output.expect = []
output.num_expect = n_expt_op
for op in e_ops:
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")
def get_curr_state_data():
if oper_evo:
return vec2mat(r.y)
else:
return r.y
#
# start evolution
#
progress_bar.start(n_tsteps)
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.")
# get the current state / oper data if needed
cdata = None
if opt.store_states or opt.normalize_output or n_expt_op > 0:
cdata = get_curr_state_data()
if opt.normalize_output:
# cdata *= _get_norm_factor(cdata, oper_evo)
cdata *= norm_dim_factor / la_norm(cdata)
if oper_evo:
r.set_initial_value(cdata.ravel('F'), r.t)
else:
r.set_initial_value(cdata, r.t)
if opt.store_states:
output.states.append(Qobj(cdata, dims=dims))
if expt_callback:
# use callback method
e_ops(t, Qobj(cdata, dims=dims))
for m in range(n_expt_op):
output.expect[m][t_idx] = cy_expect_psi(e_ops[m].data,
cdata, e_ops[m].isherm)
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:
try:
os.remove(config.tdname + ".pyx")
except:
pass
if opt.store_final_state:
cdata = get_curr_state_data()
if opt.normalize_output:
cdata *= norm_dim_factor / la_norm(cdata)
output.final_state = Qobj(cdata, dims=dims)
return output
def _generic_ode_solve(func, ode_args, psi0, tlist, e_ops, opt,
progress_bar, dims=None):
"""
Internal function for solving ODEs.
"""
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# This function is made similar to mesolve's one for futur merging in a
# solver class
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# prepare output array
n_tsteps = len(tlist)
output = Result()
output.solver = "sesolve"
output.times = tlist
if psi0.isunitary:
initial_vector = psi0.full().ravel('F')
oper_evo = True
size = psi0.shape[0]
# oper_n = dims[0][0]
# norm_dim_factor = np.sqrt(oper_n)
elif psi0.isket:
initial_vector = psi0.full().ravel()
oper_evo = False
# norm_dim_factor = 1.0
r = scipy.integrate.ode(func)
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)
if ode_args:
r.set_f_params(*ode_args)
r.set_initial_value(initial_vector, tlist[0])
e_ops_data = []
output.expect = []
if callable(e_ops):
n_expt_op = 0
expt_callback = True
output.num_expect = 1
elif isinstance(e_ops, list):
n_expt_op = len(e_ops)
expt_callback = False
output.num_expect = n_expt_op
if n_expt_op == 0:
# fallback on storing states
opt.store_states = True
else:
for op in e_ops:
if not isinstance(op, Qobj) and callable(op):
output.expect.append(np.zeros(n_tsteps, dtype=complex))
continue
if op.isherm:
output.expect.append(np.zeros(n_tsteps))
else:
output.expect.append(np.zeros(n_tsteps, dtype=complex))
if oper_evo:
for e in e_ops:
if isinstance(e, Qobj):
e_ops_data.append(e.dag().data)
continue
e_ops_data.append(e)
else:
for e in e_ops:
if isinstance(e, Qobj):
e_ops_data.append(e.data)
continue
e_ops_data.append(e)
else:
raise TypeError("Expectation parameter must be a list or a function")
if opt.store_states:
output.states = []
if oper_evo:
def get_curr_state_data(r):
return vec2mat(r.y)
else:
def get_curr_state_data(r):
return r.y
#
# start evolution
#
dt = np.diff(tlist)
cdata = None
progress_bar.start(n_tsteps)
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.")
# get the current state / oper data if needed
if opt.store_states or opt.normalize_output \
or n_expt_op > 0 or expt_callback:
cdata = get_curr_state_data(r)
if opt.normalize_output:
# normalize per column
if oper_evo:
cdata /= la_norm(cdata, axis=0)
#cdata *= norm_dim_factor / la_norm(cdata)
r.set_initial_value(cdata.ravel('F'), r.t)
else:
#cdata /= la_norm(cdata)
norm = normalize_inplace(cdata)
if norm > 1e-12:
# only reset the solver if state changed
r.set_initial_value(cdata, r.t)
else:
r._y = cdata
if opt.store_states:
if oper_evo:
fdata = dense2D_to_fastcsr_fmode(cdata, size, size)
output.states.append(Qobj(fdata, dims=dims))
else:
fdata = dense1D_to_fastcsr_ket(cdata)
output.states.append(Qobj(fdata, dims=dims, fast='mc'))
if expt_callback:
# use callback method
output.expect.append(e_ops(t, Qobj(cdata, dims=dims)))
if oper_evo:
for m in range(n_expt_op):
if callable(e_ops_data[m]):
func = e_ops_data[m]
output.expect[m][t_idx] = func(t, Qobj(cdata, dims=dims))
continue
output.expect[m][t_idx] = (e_ops_data[m] * cdata).trace()
else:
for m in range(n_expt_op):
if callable(e_ops_data[m]):
func = e_ops_data[m]
output.expect[m][t_idx] = func(t, Qobj(cdata, dims=dims))
continue
output.expect[m][t_idx] = cy_expect_psi(e_ops_data[m], cdata,
e_ops[m].isherm)
if t_idx < n_tsteps - 1:
r.integrate(r.t + dt[t_idx])
progress_bar.finished()
if opt.store_final_state:
cdata = get_curr_state_data(r)
if opt.normalize_output:
cdata /= la_norm(cdata, axis=0)
# cdata *= norm_dim_factor / la_norm(cdata)
output.final_state = Qobj(cdata, dims=dims)
return output