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 _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