__all__ = ['mcsolve']
import os
import numpy as np
from numpy.random import RandomState, randint
import scipy.sparse as sp
from scipy.integrate import ode
from scipy.integrate._ode import zvode
from types import FunctionType, BuiltinFunctionType
from functools import partial
from qutip.fastsparse import csr2fast
from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo
from qutip.parallel import parfor, parallel_map, serial_map
from qutip.cy.mcsolve import CyMcOde, CyMcOdeDiag
from qutip.cy.spconvert import dense1D_to_fastcsr_ket
from qutip.sesolve import sesolve
from qutip.solver import (Options, Result, ExpectOps,
solver_safe, SolverSystem)
from qutip.settings import debug
from qutip.ui.progressbar import TextProgressBar, BaseProgressBar
import qutip.settings
if debug:
import inspect
#
# Internal, global variables for storing references to dynamically loaded
# cython functions
# Todo: use real warning
def warn(text):
print(text)
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=0,
args={}, options=None, progress_bar=True,
map_func=parallel_map, map_kwargs={}, _safe_mode=True):
r"""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`, ``list``
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 : :class:`qutip.Qobj`, ``list``
single collapse operator or a ``list`` of collapse operators.
e_ops : :class:`qutip.Qobj`, ``list``
single operator as Qobj or ``list`` or equivalent of Qobj 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 isinstance(c_ops, (Qobj, QobjEvo)):
c_ops = [c_ops]
if options is None:
options = Options()
if options.rhs_reuse and not isinstance(H, SolverSystem):
# TODO: deprecate when going to class based solver.
if "mcsolve" in solver_safe:
# print(" ")
H = solver_safe["mcsolve"]
else:
pass
# raise Exception("Could not find the Hamiltonian to reuse.")
if not ntraj:
ntraj = options.ntraj
if len(c_ops) == 0 and not options.rhs_reuse:
warn("No c_ops, using sesolve")
return sesolve(H, psi0, tlist, e_ops=e_ops, args=args,
options=options, progress_bar=progress_bar,
_safe_mode=_safe_mode)
try:
num_traj = int(ntraj)
except TypeError:
num_traj = max(ntraj)
# set the physics
if not psi0.isket:
raise Exception("Initial state must be a state vector.")
# load monte carlo class
mc = _MC(options)
if isinstance(H, SolverSystem):
mc.ss = H
else:
mc.make_system(H, c_ops, tlist, args, options)
mc.reset(tlist[0], psi0)
mc.set_e_ops(e_ops)
if options.seeds is not None:
mc.seed(num_traj, options.seeds)
if _safe_mode:
mc.run_test()
# Run the simulation
mc.run(num_traj=num_traj, tlist=tlist,
progress_bar=progress_bar,
map_func=map_func, map_kwargs=map_kwargs)
return mc.get_result(ntraj)
# -----------------------------------------------------------------------------
# MONTE CARLO CLASS
# -----------------------------------------------------------------------------
class _MC():
"""
Private class for solving Monte Carlo evolution from mcsolve
"""
def __init__(self, options=None):
if options is None:
options = Options()
self.options = options
self.ss = None
self.tlist = None
self.e_ops = None
self.ran = False
self.psi0 = None
self.seeds = []
self.t = 0.
self.num_traj = 0
self.args_col = None
self._psi_out = []
self._expect_out = []
self._collapse = []
self._ss_out = []
def reset(self, t=0., psi0=None):
if psi0 is not None:
self.psi0 = psi0
if self.psi0 is not None:
self.initial_vector = self.psi0.full().ravel("F")
if self.ss is not None and self.ss.type == "Diagonal":
self.initial_vector = np.dot(self.ss.Ud, self.initial_vector)
self.t = t
self.ran = False
self._psi_out = []
self._expect_out = []
self._collapse = []
self._ss_out = []
def seed(self, ntraj, seeds=[]):
# setup seeds array
np.random.seed()
try:
seed = int(seeds)
np.random.seed(seed)
seeds = []
except TypeError:
pass
if len(seeds) < ntraj:
self.seeds = seeds + list(randint(0, 2**32,
size=ntraj-len(seeds),
dtype=np.uint32))
else:
self.seeds = seeds[:ntraj]
def make_system(self, H, c_ops, tlist=None, args={}, options=None):
if options is None:
options = self.options
else:
self.options = options
var = _collapse_args(args)
ss = SolverSystem()
ss.td_c_ops = []
ss.td_n_ops = []
ss.args = args
ss.col_args = var
for c in c_ops:
cevo = QobjEvo(c, args, tlist=tlist)
cdc = cevo._cdc()
cevo.compile()
cdc.compile()
ss.td_c_ops.append(cevo)
ss.td_n_ops.append(cdc)
try:
H_td = QobjEvo(H, args, tlist=tlist)
H_td *= -1j
for c in ss.td_n_ops:
H_td += -0.5 * c
if options.rhs_with_state:
H_td._check_old_with_state()
H_td.compile()
ss.H_td = H_td
ss.makefunc = _qobjevo_set
ss.set_args = _qobjevo_args
ss.type = "QobjEvo"
except Exception:
ss.h_func = H
ss.Hc_td = -0.5 * sum(ss.td_n_ops)
ss.Hc_td.compile()
ss.with_state = options.rhs_with_state
ss.makefunc = _func_set
ss.set_args = _func_args
ss.type = "callback"
solver_safe["mcsolve"] = ss
self.ss = ss
self.reset()
def set_e_ops(self, e_ops=[]):
if e_ops:
self.e_ops = ExpectOps(e_ops)
else:
self.e_ops = ExpectOps([])
self.e_ops.check_dims(self.ss.td_c_ops[0].cte.dims)
ss = self.ss
if ss is not None and ss.type == "Diagonal" and not self.e_ops.isfunc:
e_ops = [
Qobj(ss.Ud @ e.full() @ ss.U, dims=e.dims)
for e in self.e_ops.e_ops
]
self.e_ops = ExpectOps(e_ops)
if not self.e_ops:
self.options.store_states = True
def run_test(self):
try:
for c_op in self.ss.td_c_ops:
c_op.mul_vec(0, self.psi0)
except Exception as e:
raise Exception("c_ops are not consistant with psi0") from e
if self.ss.type == "QobjEvo":
try:
self.ss.H_td.mul_vec(0., self.psi0)
except Exception as e:
raise Exception("Error calculating H") from e
else:
try:
rhs, ode_args = self.ss.makefunc(self.ss)
rhs(0, self.psi0.full().ravel(), ode_args)
except Exception as e:
raise Exception("Error calculating H") from e
def run(self, num_traj=0, psi0=None, tlist=None,
args={}, e_ops=None, options=None,
progress_bar=True,
map_func=parallel_map, map_kwargs={}):
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# 4 situation for run:
# - first run
# - change parameters
# - add trajectories
# (self.add_traj) Not Implemented
# - continue from the last time and states
# (self.continue_runs) Not Implemented
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
options = options if options is not None else self.options
if self.ran and tlist[0] == self.t:
# psi0 is ignored since we restart from a
# different states for each trajectories
self.continue_runs(num_traj, tlist, args, e_ops, options,
progress_bar, map_func, map_kwargs)
return
if args and args != self.ss.args:
self.ss.set_args(self.ss, args)
self.reset()
if e_ops and e_ops != self.e_ops:
self.set_e_ops(e_ops)
self.reset()
if psi0 is not None and psi0 != self.psi0:
self.psi0 = psi0
self.reset()
tlist = np.array(tlist)
if tlist is not None and np.all(tlist != self.tlist):
self.tlist = tlist
self.reset()
if self.ran:
if options.store_states and self._psi_out[0].shape[0] == 1:
self.reset()
else:
# if not reset here, add trajectories
self.add_traj(num_traj, progress_bar, map_func, map_kwargs)
return
if not num_traj:
num_traj = options.ntraj
if options.num_cpus == 1 or num_traj == 1:
map_func = serial_map
if len(self.seeds) != num_traj:
self.seed(num_traj, self.seeds)
if not progress_bar:
progress_bar = BaseProgressBar()
elif progress_bar is True:
progress_bar = TextProgressBar()
# set arguments for input to monte carlo
map_kwargs_ = {'progress_bar': progress_bar,
'num_cpus': options.num_cpus}
map_kwargs_.update(map_kwargs)
map_kwargs = map_kwargs_
if self.e_ops is None:
self.set_e_ops()
if self.ss.type == "Diagonal":
results = map_func(self._single_traj_diag, list(range(num_traj)),
**map_kwargs)
else:
results = map_func(self._single_traj, list(range(num_traj)),
**map_kwargs)
self.t = self.tlist[-1]
self.num_traj = num_traj
self.ran = True
for result in results:
state_out, ss_out, expect, collapse = result
self._psi_out.append(state_out)
self._ss_out.append(ss_out)
self._expect_out.append(expect)
self._collapse.append(collapse)
self._psi_out = np.stack(self._psi_out)
self._ss_out = np.stack(self._ss_out)
def add_traj(self, num_traj,
progress_bar=True,
map_func=parallel_map, map_kwargs={}):
raise NotImplementedError
def continue_runs(self, num_traj, tlist, args={}, e_ops=[], options=None,
progress_bar=True,
map_func=parallel_map, map_kwargs={}):
raise NotImplementedError
# --------------------------------------------------------------------------
# results functions
# --------------------------------------------------------------------------
@property
def states(self):
dims = self.psi0.dims[0]
len_ = self._psi_out.shape[2]
if self._psi_out.shape[1] == 1:
dm_t = np.zeros((len_, len_), dtype=complex)
for i in range(self.num_traj):
vec = self._psi_out[i, 0]
dm_t += np.outer(vec, vec.conj())
return Qobj(dm_t/self.num_traj, dims=[dims, dims])
else:
states = np.empty((len(self.tlist)), dtype=object)
for j in range(len(self.tlist)):
dm_t = np.zeros((len_, len_), dtype=complex)
for i in range(self.num_traj):
vec = self._psi_out[i, j]
dm_t += np.outer(vec, vec.conj())
states[j] = Qobj(dm_t/self.num_traj, dims=[dims, dims])
return states
@property
def final_state(self):
dims = self.psi0.dims[0]
len_ = self._psi_out.shape[2]
dm_t = np.zeros((len_, len_), dtype=complex)
for i in range(self.num_traj):
vec = self._psi_out[i, -1]
dm_t += np.outer(vec, vec.conj())
return Qobj(dm_t/self.num_traj, dims=[dims, dims])
@property
def runs_final_states(self):
dims = self.psi0.dims[0]
psis = np.empty((self.num_traj), dtype=object)
for i in range(self.num_traj):
psis[i] = Qobj(dense1D_to_fastcsr_ket(self._psi_out[i, -1]),
dims=dims, fast='mc')
return psis
@property
def expect(self):
return self.expect_traj_avg()
@property
def runs_expect(self):
return [expt.finish() for expt in self._expect_out]
def expect_traj_avg(self, ntraj=0):
if not ntraj:
ntraj = len(self._expect_out)
expect = np.stack([expt.raw_out for expt in self._expect_out[:ntraj]])
expect = np.mean(expect, axis=0)
result = []
for ii in range(self.e_ops.e_num):
if self.e_ops.e_ops_isherm[ii]:
result.append(np.real(expect[ii, :]))
else:
result.append(expect[ii, :])
if self.e_ops.e_ops_dict:
result = {e: result[n]
for n, e in enumerate(self.e_ops.e_ops_dict.keys())}
return result
@property
def steady_state(self):
if self._ss_out is not None:
dims = self.psi0.dims[0]
len_ = self.psi0.shape[0]
return Qobj(np.mean(self._ss_out, axis=0),
dims=[dims, dims], shape=(len_, len_))
# TO-DO rebuild steady_state from _psi_out if needed
# elif self._psi_out is not None:
# return sum(self.state_average) / self.num_traj
else:
return None
@property
def runs_states(self):
dims = self.psi0.dims
psis = np.empty((self.num_traj, len(self.tlist)), dtype=object)
for i in range(self.num_traj):
for j in range(len(self.tlist)):
psis[i, j] = Qobj(dense1D_to_fastcsr_ket(self._psi_out[i, j]),
dims=dims, fast='mc')
return psis
@property
def collapse(self):
return self._collapse
@property
def collapse_times(self):
out = []
for col_ in self._collapse:
col = list(zip(*col_))
col = ([] if len(col) == 0 else col[0])
out.append(np.array(col))
return out
return [np.array(list(zip(*col_))[0]) for col_ in self._collapse]
@property
def collapse_which(self):
out = []
for col_ in self._collapse:
col = list(zip(*col_))
col = ([] if len(col) == 0 else col[1])
out.append(np.array(col))
return out
return [np.array(list(zip(*col_))[1]) for col_ in self._collapse]
def get_result(self, ntraj=[]):
# Store results in the Result object
if not ntraj:
ntraj = [self.num_traj]
elif not isinstance(ntraj, list):
ntraj = [ntraj]
output = Result()
output.solver = 'mcsolve'
output.seeds = self.seeds
options = self.options
output.options = options
if options.steady_state_average:
output.states = self.steady_state
elif options.average_states and options.store_states:
output.states = self.states
elif options.store_states:
output.states = self.runs_states
if options.store_final_state:
if options.average_states:
output.final_state = self.final_state
else:
output.final_state = self.runs_final_states
if options.average_expect:
output.expect = [self.expect_traj_avg(n) for n in ntraj]
if len(output.expect) == 1:
output.expect = output.expect[0]
else:
output.expect = self.runs_expect
# simulation parameters
output.times = self.tlist
output.num_expect = self.e_ops.e_num
output.num_collapse = len(self.ss.td_c_ops)
output.ntraj = self.num_traj
output.col_times = self.collapse_times
output.col_which = self.collapse_which
return output
# --------------------------------------------------------------------------
# single-trajectory for monte carlo
# --------------------------------------------------------------------------
def _single_traj(self, nt):
"""
Monte Carlo algorithm returning state-vector or expectation values
at times tlist for a single trajectory.
"""
# SEED AND RNG AND GENERATE
prng = RandomState(self.seeds[nt])
opt = self.options
# set initial conditions
ss = self.ss
tlist = self.tlist
e_ops = self.e_ops.copy()
opt = self.options
rhs, ode_args = self.ss.makefunc(ss)
ODE = self._build_integration_func(rhs, ode_args, opt)
ODE.set_initial_value(self.initial_vector, tlist[0])
e_ops.init(tlist)
cymc = CyMcOde(ss, opt)
states_out, ss_out, collapses = cymc.run_ode(ODE, tlist, e_ops, prng)
# Run at end of mc_alg function
# -----------------------------
if opt.steady_state_average:
ss_out /= float(len(tlist))
return (states_out, ss_out, e_ops, collapses)
def _build_integration_func(self, rhs, ode_args, opt):
"""
Create the integration function while fixing the parameters
"""
ODE = ode(rhs)
if ode_args:
ODE.set_f_params(ode_args)
# initialize ODE solver for RHS
ODE.set_integrator('zvode', method="adams")
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)
return ODE
# --------------------------------------------------------------------------
# In development diagonalize the Hamiltonian before solving
# Same seeds give same evolution
# 3~5 time faster
# constant system only.
# --------------------------------------------------------------------------
def make_diag_system(self, H, c_ops):
ss = SolverSystem()
ss.td_c_ops = []
ss.td_n_ops = []
H_ = H.copy()
H_ *= -1j
for c in c_ops:
H_ += -0.5 * c.dag() * c
w, v = np.linalg.eig(H_.full())
arg = np.argsort(np.abs(w))
eig = w[arg]
U = v.T[arg].T
Ud = U.T.conj()
for c in c_ops:
c_diag = Qobj(Ud @ c.full() @ U, dims=c.dims)
cevo = QobjEvo(c_diag)
cdc = cevo._cdc()
cevo.compile()
cdc.compile()
ss.td_c_ops.append(cevo)
ss.td_n_ops.append(cdc)
ss.H_diag = eig
ss.Ud = Ud
ss.U = U
ss.args = {}
ss.type = "Diagonal"
solver_safe["mcsolve"] = ss
if self.e_ops and not self.e_ops.isfunc:
e_ops = [
Qobj(Ud @ e.full() @ U, dims=e.dims)
for e in self.e_ops.e_ops
]
self.e_ops = ExpectOps(e_ops)
self.ss = ss
self.reset()
def _single_traj_diag(self, nt):
"""
Monte Carlo algorithm returning state-vector or expectation values
at times tlist for a single trajectory.
"""
# SEED AND RNG AND GENERATE
prng = RandomState(self.seeds[nt])
opt = self.options
ss = self.ss
tlist = self.tlist
e_ops = self.e_ops.copy()
opt = self.options
e_ops.init(tlist)
cymc = CyMcOdeDiag(ss, opt)
states_out, ss_out, collapses =\
cymc.run_ode(self.initial_vector, tlist, e_ops, prng)
if opt.steady_state_average:
ss_out = ss.U @ ss_out @ ss.Ud
states_out = np.inner(ss.U, states_out).T
if opt.steady_state_average:
ss_out /= float(len(tlist))
return (states_out, ss_out, e_ops, collapses)
# -----------------------------------------------------------------------------
# CODES FOR PYTHON FUNCTION BASED TIME-DEPENDENT RHS
# -----------------------------------------------------------------------------
def _qobjevo_set(ss, psi0=None, args={}, opt=None):
if args:
self.set_args(args)
rhs = ss.H_td.compiled_qobjevo.mul_vec
return rhs, ()
def _qobjevo_args(ss, args):
var = _collapse_args(args)
ss.col_args = var
ss.args = args
ss.H_td.solver_set_args(args, psi0, e_ops)
for c in ss.td_c_ops:
c.solver_set_args(args, psi0, e_ops)
for c in ss.td_n_ops:
c.solver_set_args(args, psi0, e_ops)
def _func_set(HS, psi0=None, args={}, opt=None):
if args:
self.set_args(args)
else:
args = ss.args
if ss.with_state:
rhs = _funcrhs
else:
rhs = _funcrhs_with_state
return rhs, (ss.h_func, ss.Hc_td, args)
def _func_args(ss, args):
var = _collapse_args(args)
ss.col_args = var
ss.args = args
for c in ss.td_c_ops:
c.solver_set_args(args, psi0, e_ops)
for c in ss.td_n_ops:
c.solver_set_args(args, psi0, e_ops)
return rhs, (ss.h_func, ss.Hc_td, args)
# RHS of ODE for python function Hamiltonian
def _funcrhs(t, psi, h_func, Hc_td, args):
h_func_data = -1.0j * h_func(t, args).data
h_func_term = h_func_data * psi
return h_func_term + Hc_td.mul_vec(t, psi)
def _funcrhs_with_state(t, psi, h_func, Hc_td, args):
h_func_data = - 1.0j * h_func(t, psi, args).data
h_func_term = h_func_data * psi
return h_func_term + Hc_td.mul_vec(t, psi)
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 = sum([psi.data for psi in psi_list]) / ln
return Qobj(out_data, dims=dims, shape=shape, fast='mc-dm')
def _collapse_args(args):
for key in args:
if key == "collapse":
if not isinstance(args[key], list):
args[key] = []
return key
return ""