# 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.
###############################################################################
"""
This module provides solvers for the unitary Schrodinger equation.
"""
__all__ = ['sesolve']
import os
import types
from functools import partial
import numpy as np
import scipy.integrate
from scipy.linalg import norm
import qutip.settings as qset
from qutip.qobj import Qobj, isket
from qutip.rhs_generate import rhs_generate
from qutip.solver import Result, Options, 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.cy.spmatfuncs import (cy_expect_psi, cy_ode_rhs,
cy_ode_psi_func_td,
cy_ode_psi_func_td_with_state,
spmvpy_csr)
from qutip.cy.codegen import Codegen
from qutip.cy.utilities import _cython_build_cleanup
from qutip.ui.progressbar import BaseProgressBar
from qutip.cy.openmp.utilities import check_use_openmp, openmp_components
if qset.has_openmp:
from qutip.cy.openmp.parfuncs import cy_ode_rhs_openmp
if debug:
import inspect
[docs]def sesolve(H, rho0, tlist, e_ops=[], args={}, options=None,
progress_bar=BaseProgressBar(),
_safe_mode=True):
"""
Schrodinger equation evolution of a state vector for a given Hamiltonian.
Evolve the state vector or density matrix (`rho0`) using a given
Hamiltonian (`H`), by integrating the set of ordinary differential
equations that define the system.
The output is either the state vector at arbitrary points in time
(`tlist`), or the expectation values of the supplied operators
(`e_ops`). If e_ops is a callback function, it is invoked for each
time in `tlist` with time and the state as arguments, and the function
does not use any return values.
Parameters
----------
H : :class:`qutip.qobj`
system Hamiltonian, or a callback function for time-dependent
Hamiltonians.
rho0 : :class:`qutip.qobj`
initial density matrix or state vector (ket).
tlist : *list* / *array*
list of times for :math:`t`.
e_ops : list of :class:`qutip.qobj` / callback function single
single operator or list of operators for which to evaluate
expectation values.
args : *dictionary*
dictionary of parameters for time-dependent Hamiltonians and
collapse operators.
options : :class:`qutip.Qdeoptions`
with options for the ODE solver.
Returns
-------
output: :class:`qutip.solver`
An instance of the class :class:`qutip.solver`, which contains either
an *array* of expectation values for the times specified by `tlist`, or
an *array* or state vectors or density matrices corresponding to the
times in `tlist` [if `e_ops` is an empty list], or
nothing if a callback function was given inplace of operators for
which to calculate the expectation values.
"""
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, rho0, c_ops=[], e_ops=e_ops, args=args)
# convert array based time-dependence to string format
H, _, args = _td_wrap_array_str(H, [], args, tlist)
# check for type (if any) of time-dependent inputs
n_const, n_func, n_str = _td_format_check(H, [])
if options is None:
options = Options()
if (not options.rhs_reuse) or (not config.tdfunc):
# reset config time-dependence flags to default values
config.reset()
#check if should use OPENMP
check_use_openmp(options)
if n_func > 0:
res = _sesolve_list_func_td(H, rho0, tlist, e_ops, args, options,
progress_bar)
elif n_str > 0:
res = _sesolve_list_str_td(H, rho0, tlist, e_ops, args, options,
progress_bar)
elif isinstance(H, (types.FunctionType,
types.BuiltinFunctionType,
partial)):
res = _sesolve_func_td(H, rho0, tlist, e_ops, args, options,
progress_bar)
else:
res = _sesolve_const(H, rho0, tlist, e_ops, args, options,
progress_bar)
if e_ops_dict:
res.expect = {e: res.expect[n]
for n, e in enumerate(e_ops_dict.keys())}
return res
# -----------------------------------------------------------------------------
# A time-dependent unitary wavefunction equation on the list-function format
#
def _sesolve_list_func_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
#
if not isket(psi0):
raise TypeError("The unitary solver requires a ket as initial state")
#
# construct liouvillian in list-function format
#
L_list = []
if not opt.rhs_with_state:
constant_func = lambda x, y: 1.0
else:
constant_func = lambda x, y, z: 1.0
# add all hamitonian terms to the lagrangian list
for h_spec in H_list:
if isinstance(h_spec, Qobj):
h = h_spec
h_coeff = constant_func
elif isinstance(h_spec, list):
h = h_spec[0]
h_coeff = h_spec[1]
else:
raise TypeError("Incorrect specification of time-dependent " +
"Hamiltonian (expected callback function)")
L_list.append([-1j * h.data, h_coeff])
L_list_and_args = [L_list, args]
#
# setup integrator
#
initial_vector = psi0.full().ravel()
if not opt.rhs_with_state:
r = scipy.integrate.ode(psi_list_td)
else:
r = scipy.integrate.ode(psi_list_td_with_state)
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])
r.set_f_params(L_list_and_args)
#
# call generic ODE code
#
return _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar,
dims=psi0.dims)
#
# evaluate dpsi(t)/dt according to the master equation using the
# [Qobj, function] style time dependence API
#
def psi_list_td(t, psi, H_list_and_args):
H_list = H_list_and_args[0]
args = H_list_and_args[1]
H = H_list[0][0]
H_td = H_list[0][1]
out = np.zeros(psi.shape[0],dtype=complex)
spmvpy_csr(H.data, H.indices, H.indptr, psi, H_td(t, args), out)
for n in range(1, len(H_list)):
#
# args[n][0] = the sparse data for a Qobj in operator form
# args[n][1] = function callback giving the coefficient
#
H = H_list[n][0]
H_td = H_list[n][1]
spmvpy_csr(H.data, H.indices, H.indptr, psi, H_td(t, args), out)
return out
def psi_list_td_with_state(t, psi, H_list_and_args):
H_list = H_list_and_args[0]
args = H_list_and_args[1]
H = H_list[0][0]
H_td = H_list[0][1]
out = np.zeros(psi.shape[0],dtype=complex)
spmvpy_csr(H.data, H.indices, H.indptr,
psi, H_td(t, args), out)
for n in range(1, len(H_list)):
#
# args[n][0] = the sparse data for a Qobj in operator form
# args[n][1] = function callback giving the coefficient
#
H = H_list[n][0]
H_td = H_list[n][1]
spmvpy_csr(H.data, H.indices, H.indptr, psi, H_td(t, args), out)
return out
# -----------------------------------------------------------------------------
# Wave function evolution using a ODE solver (unitary quantum evolution) using
# a constant Hamiltonian.
#
def _sesolve_const(H, psi0, tlist, e_ops, args, opt, progress_bar):
"""!
Evolve the wave function using an ODE solver
"""
if debug:
print(inspect.stack()[0][3])
if not isket(psi0):
raise TypeError("psi0 must be a ket")
#
# setup integrator.
#
initial_vector = psi0.full().ravel()
L = -1.0j * H
if opt.use_openmp and L.data.nnz >= qset.openmp_thresh:
r = scipy.integrate.ode(cy_ode_rhs_openmp)
r.set_f_params(L.data.data, L.data.indices, L.data.indptr,
opt.openmp_threads)
else:
r = scipy.integrate.ode(cy_ode_rhs)
r.set_f_params(L.data.data, L.data.indices, L.data.indptr)
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])
#
# call generic ODE code
#
return _generic_ode_solve(r, psi0, tlist, e_ops, opt,
progress_bar, dims=psi0.dims)
#
# evaluate dpsi(t)/dt [not used. using cython function is being used instead]
#
def _ode_psi_func(t, psi, H):
return H * psi
# -----------------------------------------------------------------------------
# A time-dependent disipative master equation on the list-string format for
# cython compilation
#
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)
# -----------------------------------------------------------------------------
# Wave function evolution using a ODE solver (unitary quantum evolution), for
# time dependent hamiltonians
#
def _sesolve_list_td(H_func, psi0, tlist, e_ops, args, opt, progress_bar):
"""!
Evolve the wave function using an ODE solver with time-dependent
Hamiltonian.
"""
if debug:
print(inspect.stack()[0][3])
if not isket(psi0):
raise TypeError("psi0 must be a ket")
#
# configure time-dependent terms and setup ODE solver
#
if len(H_func) != 2:
raise TypeError('Time-dependent Hamiltonian list must have two terms.')
if (not isinstance(H_func[0], (list, np.ndarray))) or \
(len(H_func[0]) <= 1):
raise TypeError('Time-dependent Hamiltonians must be a list with two '
+ 'or more terms')
if (not isinstance(H_func[1], (list, np.ndarray))) or \
(len(H_func[1]) != (len(H_func[0]) - 1)):
raise TypeError('Time-dependent coefficients must be list with ' +
'length N-1 where N is the number of ' +
'Hamiltonian terms.')
tflag = 1
if opt.rhs_reuse and config.tdfunc is None:
print("No previous time-dependent RHS found.")
print("Generating one for you...")
rhs_generate(H_func, args)
lenh = len(H_func[0])
if opt.tidy:
H_func[0] = [(H_func[0][k]).tidyup() for k in range(lenh)]
# create data arrays for time-dependent RHS function
Hdata = [-1.0j * H_func[0][k].data.data for k in range(lenh)]
Hinds = [H_func[0][k].data.indices for k in range(lenh)]
Hptrs = [H_func[0][k].data.indptr for k in range(lenh)]
# setup ode args string
string = ""
for k in range(lenh):
string += ("Hdata[" + str(k) + "], Hinds[" + str(k) +
"], Hptrs[" + str(k) + "],")
if args:
td_consts = args.items()
for elem in td_consts:
string += str(elem[1])
if elem != td_consts[-1]:
string += (",")
# run code generator
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)
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(' + string + ')', '<string>', 'exec')
exec(code)
#
# call generic ODE code
#
return _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar
, dims=psi0.dims)
# -----------------------------------------------------------------------------
# Wave function evolution using a ODE solver (unitary quantum evolution), for
# time dependent hamiltonians
#
def _sesolve_func_td(H_func, psi0, tlist, e_ops, args, opt, progress_bar):
"""!
Evolve the wave function using an ODE solver with time-dependent
Hamiltonian.
"""
if debug:
print(inspect.stack()[0][3])
if not isket(psi0):
raise TypeError("psi0 must be a ket")
#
# setup integrator
#
new_args = None
if type(args) is dict:
new_args = {}
for key in args:
if isinstance(args[key], Qobj):
new_args[key] = args[key].data
else:
new_args[key] = args[key]
elif type(args) is list or type(args) is tuple:
new_args = []
for arg in args:
if isinstance(arg, Qobj):
new_args.append(arg.data)
else:
new_args.append(arg)
if type(args) is tuple:
new_args = tuple(new_args)
else:
if isinstance(args, Qobj):
new_args = args.data
else:
new_args = args
initial_vector = psi0.full().ravel()
if not opt.rhs_with_state:
r = scipy.integrate.ode(cy_ode_psi_func_td)
else:
r = scipy.integrate.ode(cy_ode_psi_func_td_with_state)
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])
r.set_f_params(H_func, new_args)
#
# call generic ODE code
#
return _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar,
dims=psi0.dims)
#
# evaluate dpsi(t)/dt for time-dependent hamiltonian
#
def _ode_psi_func_td(t, psi, H_func, args):
H = H_func(t, args)
return -1j * (H * psi)
def _ode_psi_func_td_with_state(t, psi, H_func, args):
H = H_func(t, psi, args)
return -1j * (H * psi)
# -----------------------------------------------------------------------------
# Solve an ODE which solver parameters already setup (r). Calculate the
# required expectation values or invoke callback function at each time step.
#
def _generic_ode_solve(r, psi0, tlist, e_ops, opt, progress_bar, dims=None):
"""
Internal function for solving ODEs.
"""
if opt.normalize_output:
state_norm_func = norm
else:
state_norm_func = None
#
# prepare output array
#
n_tsteps = len(tlist)
output = Result()
output.solver = "sesolve"
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:
# fallback on storing states
output.states = []
opt.store_states = True
else:
output.expect = []
output.num_expect = n_expt_op
for op in e_ops:
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)
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 state_norm_func:
data = r.y / state_norm_func(r.y)
r.set_initial_value(data, r.t)
if opt.store_states:
output.states.append(Qobj(r.y, dims=dims))
if expt_callback:
# use callback method
e_ops(t, Qobj(r.y, dims=psi0.dims))
for m in range(n_expt_op):
output.expect[m][t_idx] = cy_expect_psi(e_ops[m].data,
r.y, e_ops[m].isherm)
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:
try:
os.remove(config.tdname + ".pyx")
except:
pass
if opt.store_final_state:
output.final_state = Qobj(r.y, dims=dims)
return output