def _init_ptrace_stuff(sel):
psi0 = Qobj(odeconfig.psi0,
dims=odeconfig.psi0_dims,
shape=odeconfig.psi0_shape)
qtf90.qutraj_run.init_ptrace_stuff(odeconfig.psi0_dims[0],
np.array(sel) + 1,
psi0.ptrace(sel).shape[0])
def _init_ptrace_stuff(sel):
psi0 = Qobj(config.psi0,
dims=config.psi0_dims,
shape=config.psi0_shape)
qtf90.qutraj_run.init_ptrace_stuff(config.psi0_dims[0],
np.array(sel) + 1,
psi0.ptrace(sel).shape[0])
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
