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