Source code for qutip.mcsolve

# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################

__all__ = ['mcsolve']

import os
from types import FunctionType
import numpy as np
from numpy.random import RandomState, randint
from scipy.integrate import ode
import scipy.sparse as sp
from scipy.integrate._ode import zvode
from scipy.linalg.blas import get_blas_funcs
from qutip.qobj import Qobj
from qutip.parallel import parfor, parallel_map, serial_map
from qutip.cy.spmatfuncs import cy_ode_rhs, cy_expect_psi_csr, spmv, spmv_csr
from qutip.cy.codegen import Codegen
from qutip.cy.utilities import _cython_build_cleanup
from qutip.cy.spconvert import dense2D_to_fastcsr_cmode
from qutip.solver import Options, Result, config, _solver_safety_check
from qutip.rhs_generate import _td_format_check, _td_wrap_array_str
from qutip.interpolate import Cubic_Spline
from qutip.settings import debug
from qutip.ui.progressbar import TextProgressBar, BaseProgressBar
from qutip.fastsparse import csr2fast
import qutip.settings

dznrm2 = get_blas_funcs("znrm2", dtype=np.float64)

if debug:
    import inspect

#
# Internal, global variables for storing references to dynamically loaded
# cython functions
#
_cy_col_spmv_func = None
_cy_col_expect_func = None
_cy_col_spmv_call_func = None
_cy_col_expect_call_func = None
_cy_rhs_func = None


class qutip_zvode(zvode):
    def step(self, *args):
        itask = self.call_args[2]
        self.rwork[0] = args[4]
        self.call_args[2] = 5
        r = self.run(*args)
        self.call_args[2] = itask
        return r


[docs]def mcsolve(H, psi0, tlist, c_ops=[], e_ops=[], ntraj=None, args={}, options=None, progress_bar=True, map_func=None, map_kwargs=None, _safe_mode=True): """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 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 if _safe_mode: _solver_safety_check(H, psi0, c_ops, e_ops, args) 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.") 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
# ----------------------------------------------------------------------------- # MONTE CARLO CLASS # ----------------------------------------------------------------------------- class _MC(): """ Private class for solving Monte Carlo evolution from mcsolve """ def __init__(self, config): self.config = config # set output variables, even if they are not used to simplify output # code. self.psi_out = None self.expect_out = [] self.collapse_times_out = None self.which_op_out = None # FOR EVOLUTION WITH COLLAPSE OPERATORS if config.c_num: # preallocate ntraj arrays for state vectors, collapse times, and # which operator self.collapse_times_out = np.zeros(config.ntraj, dtype=np.ndarray) self.which_op_out = np.zeros(config.ntraj, dtype=np.ndarray) if config.e_num == 0 or config.options.store_states: self.psi_out = [None] * config.ntraj if config.e_num > 0: self.expect_out = [None] * config.ntraj # setup seeds array if self.config.options.seeds is None: self.config.options.seeds = \ randint(1, 100000000.0 + 1, size=config.ntraj) else: # if ntraj was reduced but reusing seeds seed_length = len(config.options.seeds) if seed_length > config.ntraj: self.config.options.seeds = \ config.options.seeds[0:config.ntraj] # if ntraj was increased but reusing seeds elif seed_length < config.ntraj: newseeds = randint(1, 100000000.0 + 1, size=(config.ntraj - seed_length)) self.config.options.seeds = np.hstack( (config.options.seeds, newseeds)) def run(self): if debug: print(inspect.stack()[0][3]) if self.config.c_num == 0: self.config.ntraj = 1 if self.config.e_num == 0 or self.config.options.store_states: self.expect_out, self.psi_out = \ _evolve_no_collapse_psi_out(self.config) else: self.expect_out = _evolve_no_collapse_expect_out(self.config) else: # set arguments for input to monte carlo map_kwargs = {'progress_bar': self.config.progress_bar, 'num_cpus': self.config.options.num_cpus} map_kwargs.update(self.config.map_kwargs) task_args = (self.config, self.config.options, self.config.options.seeds) task_kwargs = {} results = config.map_func(_mc_alg_evolve, list(range(config.ntraj)), task_args, task_kwargs, **map_kwargs) for n, result in enumerate(results): state_out, expect_out, collapse_times, which_oper = result if self.config.e_num == 0 or self.config.options.store_states: self.psi_out[n] = state_out if self.config.e_num > 0: self.expect_out[n] = expect_out self.collapse_times_out[n] = collapse_times self.which_op_out[n] = which_oper self.psi_out = np.asarray(self.psi_out, dtype=object) # ----------------------------------------------------------------------------- # CODES FOR PYTHON FUNCTION BASED TIME-DEPENDENT RHS # ----------------------------------------------------------------------------- # RHS of ODE for time-dependent systems with no collapse operators def _tdRHS(t, psi, config): h_data = config.h_func(t, config.h_func_args).data return spmv(h_data, psi) # RHS of ODE for constant Hamiltonian and at least one function based # collapse operator def _cRHStd(t, psi, config): sys = cy_ode_rhs(t, psi, config.h_data, config.h_ind, config.h_ptr) col = np.array([np.abs(config.c_funcs[j](t, config.c_func_args)) ** 2 * spmv_csr(config.n_ops_data[j], config.n_ops_ind[j], config.n_ops_ptr[j], psi) for j in config.c_td_inds]) return sys - 0.5 * np.sum(col, 0) # RHS of ODE for list-function based Hamiltonian def _tdRHStd(t, psi, config): const_term = spmv_csr(config.h_data, config.h_ind, config.h_ptr, psi) h_func_term = np.array([config.h_funcs[j](t, config.h_func_args) * spmv_csr(config.h_td_data[j], config.h_td_ind[j], config.h_td_ptr[j], psi) for j in config.h_td_inds]) col_func_terms = np.array([np.abs( config.c_funcs[j](t, config.c_func_args)) ** 2 * spmv_csr(config.n_ops_data[j], config.n_ops_ind[j], config.n_ops_ptr[j], psi) for j in config.c_td_inds]) return (const_term + np.sum(h_func_term, 0) - 0.5 * np.sum(col_func_terms, 0)) def _tdRHStd_with_state(t, psi, config): const_term = spmv_csr(config.h_data, config.h_ind, config.h_ptr, psi) h_func_term = np.array([ config.h_funcs[j](t, psi, config.h_func_args) * spmv_csr(config.h_td_data[j], config.h_td_ind[j], config.h_td_ptr[j], psi) for j in config.h_td_inds]) col_func_terms = np.array([ np.abs(config.c_funcs[j](t, config.c_func_args)) ** 2 * spmv_csr(config.n_ops_data[j], config.n_ops_ind[j], config.n_ops_ptr[j], psi) for j in config.c_td_inds]) return (const_term + np.sum(h_func_term, 0) - 0.5 * np.sum(col_func_terms, 0)) # RHS of ODE for python function Hamiltonian def _pyRHSc(t, psi, config): h_func_data = - 1.0j * config.h_funcs(t, config.h_func_args) h_func_term = spmv(h_func_data, psi) const_col_term = 0 if len(config.c_const_inds) > 0: const_col_term = spmv_csr(config.h_data, config.h_ind, config.h_ptr, psi) return h_func_term + const_col_term def _pyRHSc_with_state(t, psi, config): h_func_data = - 1.0j * config.h_funcs(t, psi, config.h_func_args) h_func_term = spmv(h_func_data, psi) const_col_term = 0 if len(config.c_const_inds) > 0: const_col_term = spmv_csr(config.h_data, config.h_ind, config.h_ptr, psi) return h_func_term + const_col_term # ----------------------------------------------------------------------------- # evolution solver: return psi at requested times for no collapse operators # ----------------------------------------------------------------------------- def _evolve_no_collapse_psi_out(config): """ Calculates state vectors at times tlist if no collapse AND no expectation values are given. """ global _cy_rhs_func global _cy_col_spmv_func, _cy_col_expect_func global _cy_col_spmv_call_func, _cy_col_expect_call_func num_times = len(config.tlist) psi_out = np.array([None] * num_times) expect_out = [] for i in range(config.e_num): if config.e_ops_isherm[i]: # preallocate real array of zeros expect_out.append(np.zeros(num_times, dtype=float)) else: # preallocate complex array of zeros expect_out.append(np.zeros(num_times, dtype=complex)) expect_out[i][0] = \ cy_expect_psi_csr(config.e_ops_data[i], config.e_ops_ind[i], config.e_ops_ptr[i], config.psi0, config.e_ops_isherm[i]) if debug: print(inspect.stack()[0][3]) if not _cy_rhs_func: _mc_func_load(config) opt = config.options if config.tflag in [1, 10, 11]: ODE = ode(_cy_rhs_func) code = compile('ODE.set_f_params(' + config.string + ')', '<string>', 'exec') exec(code) elif config.tflag == 2: ODE = ode(_cRHStd) ODE.set_f_params(config) elif config.tflag in [20, 22]: if config.options.rhs_with_state: ODE = ode(_tdRHStd_with_state) else: ODE = ode(_tdRHStd) ODE.set_f_params(config) elif config.tflag == 3: if config.options.rhs_with_state: ODE = ode(_pyRHSc_with_state) else: ODE = ode(_pyRHSc) ODE.set_f_params(config) else: ODE = ode(cy_ode_rhs) ODE.set_f_params(config.h_data, config.h_ind, config.h_ptr) # initialize ODE solver for RHS ODE.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) # set initial conditions ODE.set_initial_value(config.psi0, config.tlist[0]) psi_out[0] = Qobj(config.psi0, config.psi0_dims, config.psi0_shape) for k in range(1, num_times): ODE.integrate(config.tlist[k], step=0) # integrate up to tlist[k] if ODE.successful(): state = ODE.y / dznrm2(ODE.y) psi_out[k] = Qobj(state, config.psi0_dims, config.psi0_shape) for jj in range(config.e_num): expect_out[jj][k] = cy_expect_psi_csr( config.e_ops_data[jj], config.e_ops_ind[jj], config.e_ops_ptr[jj], state, config.e_ops_isherm[jj]) else: raise ValueError('Error in ODE solver') return expect_out, psi_out # ----------------------------------------------------------------------------- # evolution solver: return expectation values at requested times for no # collapse oper # ----------------------------------------------------------------------------- def _evolve_no_collapse_expect_out(config): """ Calculates expect.values at times tlist if no collapse ops. given """ global _cy_rhs_func global _cy_col_spmv_func, _cy_col_expect_func global _cy_col_spmv_call_func, _cy_col_expect_call_func if debug: print(inspect.stack()[0][3]) num_times = len(config.tlist) expect_out = [] for i in range(config.e_num): if config.e_ops_isherm[i]: # preallocate real array of zeros expect_out.append(np.zeros(num_times, dtype=float)) else: # preallocate complex array of zeros expect_out.append(np.zeros(num_times, dtype=complex)) expect_out[i][0] = \ cy_expect_psi_csr(config.e_ops_data[i], config.e_ops_ind[i], config.e_ops_ptr[i], config.psi0, config.e_ops_isherm[i]) if not _cy_rhs_func: _mc_func_load(config) opt = config.options if config.tflag in [1, 10, 11]: ODE = ode(_cy_rhs_func) code = compile('ODE.set_f_params(' + config.string + ')', '<string>', 'exec') exec(code) elif config.tflag == 2: ODE = ode(_cRHStd) ODE.set_f_params(config) elif config.tflag in [20, 22]: if config.options.rhs_with_state: ODE = ode(_tdRHStd_with_state) else: ODE = ode(_tdRHStd) ODE.set_f_params(config) elif config.tflag == 3: if config.options.rhs_with_state: ODE = ode(_pyRHSc_with_state) else: ODE = ode(_pyRHSc) ODE.set_f_params(config) else: ODE = ode(cy_ode_rhs) ODE.set_f_params(config.h_data, config.h_ind, config.h_ptr) ODE.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) ODE.set_initial_value(config.psi0, config.tlist[0]) for jj in range(config.e_num): expect_out[jj][0] = cy_expect_psi_csr( config.e_ops_data[jj], config.e_ops_ind[jj], config.e_ops_ptr[jj], config.psi0, config.e_ops_isherm[jj]) for k in range(1, num_times): ODE.integrate(config.tlist[k], step=0) # integrate up to tlist[k] if ODE.successful(): state = ODE.y / dznrm2(ODE.y) for jj in range(config.e_num): expect_out[jj][k] = cy_expect_psi_csr( config.e_ops_data[jj], config.e_ops_ind[jj], config.e_ops_ptr[jj], state, config.e_ops_isherm[jj]) else: raise ValueError('Error in ODE solver') return expect_out # ----------------------------------------------------------------------------- # single-trajectory for monte carlo # ----------------------------------------------------------------------------- def _mc_alg_evolve(nt, config, opt, seeds): """ Monte Carlo algorithm returning state-vector or expectation values at times tlist for a single trajectory. """ global _cy_rhs_func global _cy_col_spmv_func, _cy_col_expect_func global _cy_col_spmv_call_func, _cy_col_expect_call_func tlist = config.tlist num_times = len(tlist) if not _cy_rhs_func: _mc_func_load(config) if config.options.steady_state_average: states_out = np.zeros((1), dtype=object) else: states_out = np.zeros((num_times), dtype=object) temp = sp.csr_matrix( np.reshape(config.psi0, (config.psi0.shape[0], 1)), dtype=complex) temp = csr2fast(temp) if (config.options.average_states and not config.options.steady_state_average): # output is averaged states, so use dm states_out[0] = Qobj(temp*temp.H, [config.psi0_dims[0], config.psi0_dims[0]], [config.psi0_shape[0], config.psi0_shape[0]], fast='mc-dm') elif (not config.options.average_states and not config.options.steady_state_average): # output is not averaged, so write state vectors states_out[0] = Qobj(temp, config.psi0_dims, config.psi0_shape, fast='mc') elif config.options.steady_state_average: states_out[0] = temp * temp.H # PRE-GENERATE LIST FOR EXPECTATION VALUES expect_out = [] for i in range(config.e_num): if config.e_ops_isherm[i]: # preallocate real array of zeros expect_out.append(np.zeros(num_times, dtype=float)) else: # preallocate complex array of zeros expect_out.append(np.zeros(num_times, dtype=complex)) expect_out[i][0] = \ cy_expect_psi_csr(config.e_ops_data[i], config.e_ops_ind[i], config.e_ops_ptr[i], config.psi0, config.e_ops_isherm[i]) collapse_times = [] which_oper = [] # SEED AND RNG AND GENERATE prng = RandomState(seeds[nt]) # first rand is collapse norm, second is which operator rand_vals = prng.rand(2) # CREATE ODE OBJECT CORRESPONDING TO DESIRED TIME-DEPENDENCE if config.tflag in [1, 10, 11]: ODE = ode(_cy_rhs_func) code = compile('ODE.set_f_params(' + config.string + ')', '<string>', 'exec') exec(code) elif config.tflag == 2: ODE = ode(_cRHStd) ODE.set_f_params(config) elif config.tflag in [20, 22]: if config.options.rhs_with_state: ODE = ode(_tdRHStd_with_state) else: ODE = ode(_tdRHStd) ODE.set_f_params(config) elif config.tflag == 3: if config.options.rhs_with_state: ODE = ode(_pyRHSc_with_state) else: ODE = ode(_pyRHSc) ODE.set_f_params(config) else: ODE = ode(_cy_rhs_func) ODE.set_f_params(config.h_data, config.h_ind, config.h_ptr) # initialize ODE solver for RHS ODE._integrator = qutip_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 not len(ODE._y): ODE.t = 0.0 ODE._y = np.array([0.0], complex) ODE._integrator.reset(len(ODE._y), ODE.jac is not None) # set initial conditions ODE.set_initial_value(config.psi0, tlist[0]) # make array for collapse operator inds cinds = np.arange(config.c_num) # RUN ODE UNTIL EACH TIME IN TLIST for k in range(1, num_times): # ODE WHILE LOOP FOR INTEGRATE UP TO TIME TLIST[k] while ODE.t < tlist[k]: t_prev = ODE.t y_prev = ODE.y norm2_prev = dznrm2(ODE._y) ** 2 # integrate up to tlist[k], one step at a time. ODE.integrate(tlist[k], step=1) if not ODE.successful(): raise Exception("ZVODE failed!") norm2_psi = dznrm2(ODE._y) ** 2 if norm2_psi <= rand_vals[0]: # collapse has occured: # find collapse time to within specified tolerance # ------------------------------------------------ ii = 0 t_final = ODE.t while ii < config.norm_steps: ii += 1 t_guess = t_prev + \ np.log(norm2_prev / rand_vals[0]) / \ np.log(norm2_prev / norm2_psi) * (t_final - t_prev) ODE._y = y_prev ODE.t = t_prev ODE._integrator.call_args[3] = 1 ODE.integrate(t_guess, step=0) if not ODE.successful(): raise Exception( "ZVODE failed after adjusting step size!") norm2_guess = dznrm2(ODE._y)**2 if (np.abs(rand_vals[0] - norm2_guess) < config.norm_tol * rand_vals[0]): break elif (norm2_guess < rand_vals[0]): # t_guess is still > t_jump t_final = t_guess norm2_psi = norm2_guess else: # t_guess < t_jump t_prev = t_guess y_prev = ODE.y norm2_prev = norm2_guess if ii > config.norm_steps: raise Exception("Norm tolerance not reached. " + "Increase accuracy of ODE solver or " + "Options.norm_steps.") collapse_times.append(ODE.t) # some string based collapse operators if config.tflag in [1, 11]: n_dp = [cy_expect_psi_csr(config.n_ops_data[i], config.n_ops_ind[i], config.n_ops_ptr[i], ODE._y, 1) for i in config.c_const_inds] _locals = locals() # calculates the expectation values for time-dependent # norm collapse operators exec(_cy_col_expect_call_func, globals(), _locals) n_dp = np.array(_locals['n_dp']) elif config.tflag in [2, 20, 22]: # some Python function based collapse operators n_dp = [cy_expect_psi_csr(config.n_ops_data[i], config.n_ops_ind[i], config.n_ops_ptr[i], ODE._y, 1) for i in config.c_const_inds] n_dp += [abs(config.c_funcs[i]( ODE.t, config.c_func_args)) ** 2 * cy_expect_psi_csr(config.n_ops_data[i], config.n_ops_ind[i], config.n_ops_ptr[i], ODE._y, 1) for i in config.c_td_inds] n_dp = np.array(n_dp) else: # all constant collapse operators. n_dp = np.array( [cy_expect_psi_csr(config.n_ops_data[i], config.n_ops_ind[i], config.n_ops_ptr[i], ODE._y, 1) for i in range(config.c_num)]) # determine which operator does collapse and store it kk = np.cumsum(n_dp / np.sum(n_dp)) j = cinds[kk >= rand_vals[1]][0] which_oper.append(j) if j in config.c_const_inds: state = spmv_csr(config.c_ops_data[j], config.c_ops_ind[j], config.c_ops_ptr[j], ODE._y) else: if config.tflag in [1, 11]: _locals = locals() # calculates the state vector for collapse by a # time-dependent collapse operator exec(_cy_col_spmv_call_func, globals(), _locals) state = _locals['state'] else: state = \ config.c_funcs[j](ODE.t, config.c_func_args) * \ spmv_csr(config.c_ops_data[j], config.c_ops_ind[j], config.c_ops_ptr[j], ODE._y) state = state / dznrm2(state) ODE._y = state ODE._integrator.call_args[3] = 1 rand_vals = prng.rand(2) # after while loop # ---------------- out_psi = ODE._y / dznrm2(ODE._y) if config.e_num == 0 or config.options.store_states: out_psi_csr = dense2D_to_fastcsr_cmode(np.reshape(out_psi, (out_psi.shape[0], 1)), out_psi.shape[0], 1) if (config.options.average_states and not config.options.steady_state_average): states_out[k] = Qobj( out_psi_csr * out_psi_csr.H, [config.psi0_dims[0], config.psi0_dims[0]], [config.psi0_shape[0], config.psi0_shape[0]], fast='mc-dm') elif config.options.steady_state_average: states_out[0] = ( states_out[0] + (out_psi_csr * out_psi_csr.H)) else: states_out[k] = Qobj(out_psi_csr, config.psi0_dims, config.psi0_shape, fast='mc') for jj in range(config.e_num): expect_out[jj][k] = cy_expect_psi_csr( config.e_ops_data[jj], config.e_ops_ind[jj], config.e_ops_ptr[jj], out_psi, config.e_ops_isherm[jj]) # Run at end of mc_alg function # ----------------------------- if config.options.steady_state_average: states_out = np.array([Qobj(states_out[0] / float(len(tlist)), [config.psi0_dims[0], config.psi0_dims[0]], [config.psi0_shape[0], config.psi0_shape[0]], fast='mc-dm')]) return (states_out, expect_out, np.array(collapse_times, dtype=float), np.array(which_oper, dtype=int)) def _mc_func_load(config): """Load cython functions""" global _cy_rhs_func global _cy_col_spmv_func, _cy_col_expect_func global _cy_col_spmv_call_func, _cy_col_expect_call_func if debug: print(inspect.stack()[0][3] + " in " + str(os.getpid())) if config.tflag in [1, 10, 11]: # compile time-depdendent RHS code if config.tflag in [1, 11]: code = compile('from ' + config.tdname + ' import cy_td_ode_rhs, col_spmv, col_expect', '<string>', 'exec') exec(code, globals()) _cy_rhs_func = cy_td_ode_rhs _cy_col_spmv_func = col_spmv _cy_col_expect_func = col_expect else: code = compile('from ' + config.tdname + ' import cy_td_ode_rhs', '<string>', 'exec') exec(code, globals()) _cy_rhs_func = cy_td_ode_rhs # compile wrapper functions for calling cython spmv and expect if config.col_spmv_code: _cy_col_spmv_call_func = compile( config.col_spmv_code, '<string>', 'exec') if config.col_expect_code: _cy_col_expect_call_func = compile( config.col_expect_code, '<string>', 'exec') elif config.tflag == 0: _cy_rhs_func = cy_ode_rhs def _mc_data_config(H, psi0, h_stuff, c_ops, c_stuff, args, e_ops, options, config): """Creates the appropriate data structures for the monte carlo solver based on the given time-dependent, or indepdendent, format. """ if debug: print(inspect.stack()[0][3]) config.soft_reset() # take care of expectation values, if any if any(e_ops): config.e_num = len(e_ops) for op in e_ops: if isinstance(op, list): op = op[0] config.e_ops_data.append(op.data.data) config.e_ops_ind.append(op.data.indices) config.e_ops_ptr.append(op.data.indptr) config.e_ops_isherm.append(op.isherm) config.e_ops_data = np.array(config.e_ops_data) config.e_ops_ind = np.array(config.e_ops_ind) config.e_ops_ptr = np.array(config.e_ops_ptr) config.e_ops_isherm = np.array(config.e_ops_isherm) # take care of collapse operators, if any if any(c_ops): config.c_num = len(c_ops) for c_op in c_ops: if isinstance(c_op, list): c_op = c_op[0] n_op = c_op.dag() * c_op config.c_ops_data.append(c_op.data.data) config.c_ops_ind.append(c_op.data.indices) config.c_ops_ptr.append(c_op.data.indptr) # norm ops config.n_ops_data.append(n_op.data.data) config.n_ops_ind.append(n_op.data.indices) config.n_ops_ptr.append(n_op.data.indptr) # to array config.c_ops_data = np.array(config.c_ops_data) config.c_ops_ind = np.array(config.c_ops_ind) config.c_ops_ptr = np.array(config.c_ops_ptr) config.n_ops_data = np.array(config.n_ops_data) config.n_ops_ind = np.array(config.n_ops_ind) config.n_ops_ptr = np.array(config.n_ops_ptr) if config.tflag == 0: # CONSTANT H & C_OPS CODE # ----------------------- if config.cflag: config.c_const_inds = np.arange(len(c_ops)) # combine Hamiltonian and constant collapse terms into one for c_op in c_ops: n_op = c_op.dag() * c_op H -= 0.5j * \ n_op # construct Hamiltonian data structures if options.tidy: H = H.tidyup(options.atol) config.h_data = -1.0j * H.data.data config.h_ind = H.data.indices config.h_ptr = H.data.indptr elif config.tflag in [1, 10, 11]: # STRING BASED TIME-DEPENDENCE # ---------------------------- # take care of arguments for collapse operators, if any if any(args): for item in args.items(): config.c_args.append(item[1]) # constant Hamiltonian / string-type collapse operators if config.tflag == 1: H_inds = np.arange(1) H_tdterms = 0 len_h = 1 C_inds = np.arange(config.c_num) # find inds of time-dependent terms C_td_inds = np.array(c_stuff[2]) # find inds of constant terms C_const_inds = np.setdiff1d(C_inds, C_td_inds) # extract time-dependent coefficients (strings) C_tdterms = [c_ops[k][1] for k in C_td_inds] # store indicies of constant collapse terms config.c_const_inds = C_const_inds # store indicies of time-dependent collapse terms config.c_td_inds = C_td_inds for k in config.c_const_inds: H -= 0.5j * (c_ops[k].dag() * c_ops[k]) if options.tidy: H = H.tidyup(options.atol) config.h_data = [H.data.data] config.h_ind = [H.data.indices] config.h_ptr = [H.data.indptr] for k in config.c_td_inds: op = c_ops[k][0].dag() * c_ops[k][0] config.h_data.append(-0.5j * op.data.data) config.h_ind.append(op.data.indices) config.h_ptr.append(op.data.indptr) config.h_data = -1.0j * np.array(config.h_data) config.h_ind = np.array(config.h_ind) config.h_ptr = np.array(config.h_ptr) else: # string-type Hamiltonian & at least one string-type # collapse operator # ----------------- H_inds = np.arange(len(H)) # find inds of time-dependent terms H_td_inds = np.array(h_stuff[2]) # find inds of constant terms H_const_inds = np.setdiff1d(H_inds, H_td_inds) # extract time-dependent coefficients (strings or functions) config.h_tdterms = [H[k][1] for k in H_td_inds] # combine time-INDEPENDENT terms into one. H = np.array([np.sum(H[k] for k in H_const_inds)] + [H[k][0] for k in H_td_inds], dtype=object) len_h = len(H) H_inds = np.arange(len_h) # store indicies of time-dependent Hamiltonian terms config.h_td_inds = np.arange(1, len_h) # if there are any collapse operators if config.c_num > 0: if config.tflag == 10: # constant collapse operators config.c_const_inds = np.arange(config.c_num) for k in config.c_const_inds: H[0] -= 0.5j * (c_ops[k].dag() * c_ops[k]) C_inds = np.arange(config.c_num) C_tdterms = np.array([]) else: # some time-dependent collapse terms C_inds = np.arange(config.c_num) # find inds of time-dependent terms C_td_inds = np.array(c_stuff[2]) # find inds of constant terms C_const_inds = np.setdiff1d(C_inds, C_td_inds) C_tdterms = [c_ops[k][1] for k in C_td_inds] # extract time-dependent coefficients (strings) # store indicies of constant collapse terms config.c_const_inds = C_const_inds # store indicies of time-dependent collapse terms config.c_td_inds = C_td_inds for k in config.c_const_inds: H[0] -= 0.5j * (c_ops[k].dag() * c_ops[k]) else: # set empty objects if no collapse operators C_const_inds = np.arange(config.c_num) config.c_const_inds = np.arange(config.c_num) config.c_td_inds = np.array([]) C_tdterms = np.array([]) C_inds = np.array([]) # tidyup if options.tidy: H = np.array([H[k].tidyup(options.atol) for k in range(len_h)], dtype=object) # construct data sets config.h_data = [H[k].data.data for k in range(len_h)] config.h_ind = [H[k].data.indices for k in range(len_h)] config.h_ptr = [H[k].data.indptr for k in range(len_h)] for k in config.c_td_inds: config.h_data.append(-0.5j * config.n_ops_data[k]) config.h_ind.append(config.n_ops_ind[k]) config.h_ptr.append(config.n_ops_ptr[k]) config.h_data = -1.0j * np.array(config.h_data) config.h_ind = np.array(config.h_ind) config.h_ptr = np.array(config.h_ptr) # set execuatble code for collapse expectation values and spmv col_spmv_code = ("state = _cy_col_spmv_func(j, ODE.t, " + "config.c_ops_data[j], config.c_ops_ind[j], " + "config.c_ops_ptr[j], ODE.y") col_expect_code = ("for i in config.c_td_inds: " + "n_dp.append(_cy_col_expect_func(i, ODE.t, " + "config.n_ops_data[i], " + "config.n_ops_ind[i], " + "config.n_ops_ptr[i], ODE.y") for kk in range(len(config.c_args)): col_spmv_code += ",config.c_args[" + str(kk) + "]" col_expect_code += ",config.c_args[" + str(kk) + "]" col_spmv_code += ")" col_expect_code += "))" config.col_spmv_code = col_spmv_code config.col_expect_code = col_expect_code # setup ode args string config.string = "" data_range = range(len(config.h_data)) for k in data_range: config.string += ("config.h_data[" + str(k) + "], config.h_ind[" + str(k) + "], config.h_ptr[" + str(k) + "]") if k != data_range[-1]: config.string += "," # Add objects to ode args string for k in range(len(config.h_tdterms)): if isinstance(config.h_tdterms[k], Cubic_Spline): config.string += ", config.h_tdterms["+str(k)+"].coeffs" # attach args to ode args string if len(config.c_args) > 0: for kk in range(len(config.c_args)): config.string += "," + "config.c_args[" + str(kk) + "]" name = "rhs" + str(os.getpid()) + str(config.cgen_num) config.tdname = name cgen = Codegen(H_inds, config.h_tdterms, config.h_td_inds, args, C_inds, C_tdterms, config.c_td_inds, type='mc', config=config) cgen.generate(name + ".pyx") elif config.tflag in [2, 20, 22]: # PYTHON LIST-FUNCTION BASED TIME-DEPENDENCE # ------------------------------------------ # take care of Hamiltonian if config.tflag == 2: # constant Hamiltonian, at least one function based collapse # operators H_inds = np.array([0]) H_tdterms = 0 len_h = 1 else: # function based Hamiltonian H_inds = np.arange(len(H)) H_td_inds = np.array(h_stuff[1]) H_const_inds = np.setdiff1d(H_inds, H_td_inds) config.h_funcs = np.array([H[k][1] for k in H_td_inds]) config.h_func_args = args Htd = np.array([H[k][0] for k in H_td_inds], dtype=object) config.h_td_inds = np.arange(len(Htd)) H = np.sum(H[k] for k in H_const_inds) # take care of collapse operators C_inds = np.arange(config.c_num) # find inds of time-dependent terms C_td_inds = np.array(c_stuff[1]) # find inds of constant terms C_const_inds = np.setdiff1d(C_inds, C_td_inds) # store indicies of constant collapse terms config.c_const_inds = C_const_inds # store indicies of time-dependent collapse terms config.c_td_inds = C_td_inds config.c_funcs = np.zeros(config.c_num, dtype=FunctionType) for k in config.c_td_inds: config.c_funcs[k] = c_ops[k][1] config.c_func_args = args # combine constant collapse terms with constant H and construct data for k in config.c_const_inds: H -= 0.5j * (c_ops[k].dag() * c_ops[k]) if options.tidy: H = H.tidyup(options.atol) Htd = np.array([Htd[j].tidyup(options.atol) for j in config.h_td_inds], dtype=object) # setup constant H terms data config.h_data = -1.0j * H.data.data config.h_ind = H.data.indices config.h_ptr = H.data.indptr # setup td H terms data config.h_td_data = np.array( [-1.0j * Htd[k].data.data for k in config.h_td_inds]) config.h_td_ind = np.array( [Htd[k].data.indices for k in config.h_td_inds]) config.h_td_ptr = np.array( [Htd[k].data.indptr for k in config.h_td_inds]) elif config.tflag == 3: # PYTHON FUNCTION BASED HAMILTONIAN # --------------------------------- # take care of Hamiltonian config.h_funcs = H config.h_func_args = args # take care of collapse operators config.c_const_inds = np.arange(config.c_num) config.c_td_inds = np.array([]) if len(config.c_const_inds) > 0: H = 0 for k in config.c_const_inds: H -= 0.5j * (c_ops[k].dag() * c_ops[k]) if options.tidy: H = H.tidyup(options.atol) config.h_data = -1.0j * H.data.data config.h_ind = H.data.indices config.h_ptr = H.data.indptr def _mc_dm_avg(psi_list): """ Private function that averages density matrices in parallel over all trajectories for a single time using parfor. """ ln = len(psi_list) dims = psi_list[0].dims shape = psi_list[0].shape out_data = np.sum([psi.data for psi in psi_list]) / ln return Qobj(out_data, dims=dims, shape=shape, fast='mc-dm')