# 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
import numpy as np
import scipy.integrate
from scipy.linalg import norm as la_norm
from qutip.cy.stochastic import normalize_inplace
import qutip.settings as qset
from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo
from qutip.cy.spconvert import dense1D_to_fastcsr_ket, dense2D_to_fastcsr_fmode
from qutip.cy.spmatfuncs import (cy_expect_psi, cy_ode_psi_func_td,
cy_ode_psi_func_td_with_state)
from qutip.solver import Result, Options, config, solver_safe, SolverSystem
from qutip.superoperator import vec2mat
from qutip.ui.progressbar import (BaseProgressBar, TextProgressBar)
from qutip.cy.openmp.utilities import check_use_openmp, openmp_components
[docs]def sesolve(H, psi0, tlist, e_ops=None, args=None, options=None,
progress_bar=None, _safe_mode=True):
"""
Schrödinger equation evolution of a state vector or unitary matrix for a
given Hamiltonian.
Evolve the state vector (``psi0``) using a given Hamiltonian (``H``), by
integrating the set of ordinary differential equations that define the
system. Alternatively evolve a unitary matrix in solving the Schrodinger
operator equation.
The output is either the state vector or unitary matrix 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. ``e_ops`` cannot be used in conjunction
with solving the Schrodinger operator equation
Parameters
----------
H : :class:`~Qobj`, :class:`~QobjEvo`, list, or callable
System Hamiltonian as a :obj:`~Qobj , list of :obj:`Qobj` and
coefficient, :obj:`~QObjEvo`, or a callback function for time-dependent
Hamiltonians. List format and options can be found in QobjEvo's
description.
psi0 : :class:`~Qobj`
Initial state vector (ket) or initial unitary operator ``psi0 = U``.
tlist : array_like of float
List of times for :math:`t`.
e_ops : list of :class:`~Qobj` or callback function, optional
Single operator or list of operators for which to evaluate expectation
values. For operator evolution, the overlap is computed: ::
(e_ops[i].dag() * op(t)).tr()
args : dict, optional
Dictionary of scope parameters for time-dependent Hamiltonians.
options : :obj:`~solver.Options`, optional
Options for the ODE solver.
progress_bar : :obj:`~BaseProgressBar`, optional
Optional instance of :obj:`~BaseProgressBar`, or a subclass thereof,
for showing the progress of the simulation.
Returns
-------
output: :class:`~solver.Result`
An instance of the class :class:`~solver.Result`, which contains either
an array of expectation values for the times specified by ``tlist``, or
an array or state vectors 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 e_ops is None:
e_ops = []
if isinstance(e_ops, Qobj):
e_ops = [e_ops]
elif isinstance(e_ops, dict):
e_ops_dict = e_ops
e_ops = [e for e in e_ops.values()]
else:
e_ops_dict = None
if progress_bar is None:
progress_bar = BaseProgressBar()
if progress_bar is True:
progress_bar = TextProgressBar()
if not (psi0.isket or psi0.isunitary):
raise TypeError("The unitary solver requires psi0 to be"
" a ket as initial state"
" or a unitary as initial operator.")
if options is None:
options = Options()
if options.rhs_reuse and not isinstance(H, SolverSystem):
# TODO: deprecate when going to class based solver.
if "sesolve" in solver_safe:
H = solver_safe["sesolve"]
if args is None:
args = {}
check_use_openmp(options)
if isinstance(H, SolverSystem):
ss = H
elif isinstance(H, (list, Qobj, QobjEvo)):
ss = _sesolve_QobjEvo(H, tlist, args, options)
elif callable(H):
ss = _sesolve_func_td(H, args, options)
else:
raise Exception("Invalid H type")
func, ode_args = ss.makefunc(ss, psi0, args, e_ops, options)
if _safe_mode:
v = psi0.full().ravel('F')
func(0., v, *ode_args) + v
res = _generic_ode_solve(func, ode_args, psi0, tlist, e_ops, options,
progress_bar, dims=psi0.dims)
if e_ops_dict:
res.expect = {e: res.expect[n]
for n, e in enumerate(e_ops_dict.keys())}
res.SolverSystem = ss
return res
# -----------------------------------------------------------------------------
# A time-dependent unitary wavefunction equation on the list-function format
#
def _sesolve_QobjEvo(H, tlist, args, opt):
"""
Prepare the system for the solver, H can be an QobjEvo.
"""
H_td = -1.0j * QobjEvo(H, args, tlist=tlist)
if opt.rhs_with_state:
H_td._check_old_with_state()
nthread = opt.openmp_threads if opt.use_openmp else 0
H_td.compile(omp=nthread)
ss = SolverSystem()
ss.H = H_td
ss.makefunc = _qobjevo_set
solver_safe["sesolve"] = ss
return ss
def _qobjevo_set(HS, psi, args, e_ops, opt):
"""
From the system, get the ode function and args
"""
H_td = HS.H
H_td.solver_set_args(args, psi, e_ops)
if psi.isunitary:
func = H_td.compiled_qobjevo.ode_mul_mat_f_vec
elif psi.isket:
func = H_td.compiled_qobjevo.mul_vec
else:
raise TypeError("The unitary solver requires psi0 to be"
" a ket as initial state"
" or a unitary as initial operator.")
return func, ()
# -----------------------------------------------------------------------------
# Wave function evolution using a ODE solver (unitary quantum evolution), for
# time dependent hamiltonians.
#
def _sesolve_func_td(H_func, args, opt):
"""
Prepare the system for the solver, H is a function.
"""
ss = SolverSystem()
ss.H = H_func
ss.makefunc = _Hfunc_set
solver_safe["sesolve"] = ss
return ss
def _Hfunc_set(HS, psi, args, e_ops, opt):
"""
From the system, get the ode function and args
"""
H_func = HS.H
if psi.isunitary:
if not opt.rhs_with_state:
func = _ode_oper_func_td
else:
func = _ode_oper_func_td_with_state
else:
if not opt.rhs_with_state:
func = cy_ode_psi_func_td
else:
func = cy_ode_psi_func_td_with_state
return func, (H_func, args)
# -----------------------------------------------------------------------------
# evaluate dU(t)/dt according to the schrodinger equation
#
def _ode_oper_func_td(t, y, H_func, args):
H = H_func(t, args).data * -1j
ym = vec2mat(y)
return (H * ym).ravel("F")
def _ode_oper_func_td_with_state(t, y, H_func, args):
H = H_func(t, y, args).data * -1j
ym = vec2mat(y)
return (H * ym).ravel("F")
# -----------------------------------------------------------------------------
# Solve an ODE for func.
# Calculate the required expectation values or invoke callback
# function at each time step.
def _generic_ode_solve(func, ode_args, psi0, tlist, e_ops, opt,
progress_bar, dims=None):
"""
Internal function for solving ODEs.
"""
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# This function is made similar to mesolve's one for futur merging in a
# solver class
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# prepare output array
n_tsteps = len(tlist)
output = Result()
output.solver = "sesolve"
output.times = tlist
if psi0.isunitary:
initial_vector = psi0.full().ravel('F')
oper_evo = True
size = psi0.shape[0]
# oper_n = dims[0][0]
# norm_dim_factor = np.sqrt(oper_n)
elif psi0.isket:
initial_vector = psi0.full().ravel()
oper_evo = False
# norm_dim_factor = 1.0
r = scipy.integrate.ode(func)
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)
if ode_args:
r.set_f_params(*ode_args)
r.set_initial_value(initial_vector, tlist[0])
e_ops_data = []
output.expect = []
if callable(e_ops):
n_expt_op = 0
expt_callback = True
output.num_expect = 1
elif isinstance(e_ops, list):
n_expt_op = len(e_ops)
expt_callback = False
output.num_expect = n_expt_op
if n_expt_op == 0:
# fallback on storing states
opt.store_states = True
else:
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))
if oper_evo:
for e in e_ops:
if e.dims[1] != psi0.dims[0]:
raise TypeError(f"e_ops dims ({e.dims}) are not compatible"
f" with the state's ({psi0.dims})")
else:
e_ops_data.append(e.dag().data)
else:
for e in e_ops:
if e.dims[1] != psi0.dims[0]:
raise TypeError(f"e_ops dims ({e.dims}) are not compatible"
f" with the state's ({psi0.dims})")
else:
e_ops_data.append(e.data)
else:
raise TypeError("Expectation parameter must be a list or a function")
if opt.store_states:
output.states = []
if oper_evo:
def get_curr_state_data(r):
return vec2mat(r.y)
else:
def get_curr_state_data(r):
return r.y
#
# start evolution
#
dt = np.diff(tlist)
cdata = None
progress_bar.start(n_tsteps)
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.")
# get the current state / oper data if needed
if opt.store_states or opt.normalize_output or n_expt_op > 0 or expt_callback:
cdata = get_curr_state_data(r)
if opt.normalize_output:
# normalize per column
if oper_evo:
cdata /= la_norm(cdata, axis=0)
#cdata *= norm_dim_factor / la_norm(cdata)
r.set_initial_value(cdata.ravel('F'), r.t)
else:
#cdata /= la_norm(cdata)
norm = normalize_inplace(cdata)
if norm > 1e-12:
# only reset the solver if state changed
r.set_initial_value(cdata, r.t)
else:
r._y = cdata
if opt.store_states:
if oper_evo:
fdata = dense2D_to_fastcsr_fmode(cdata, size, size)
output.states.append(Qobj(fdata, dims=dims))
else:
fdata = dense1D_to_fastcsr_ket(cdata)
output.states.append(Qobj(fdata, dims=dims, fast='mc'))
if expt_callback:
# use callback method
output.expect.append(e_ops(t, Qobj(cdata, dims=dims)))
if oper_evo:
for m in range(n_expt_op):
output.expect[m][t_idx] = (e_ops_data[m] * cdata).trace()
else:
for m in range(n_expt_op):
output.expect[m][t_idx] = cy_expect_psi(e_ops_data[m], cdata,
e_ops[m].isherm)
if t_idx < n_tsteps - 1:
r.integrate(r.t + dt[t_idx])
progress_bar.finished()
if opt.store_final_state:
cdata = get_curr_state_data(r)
if opt.normalize_output:
cdata /= la_norm(cdata, axis=0)
# cdata *= norm_dim_factor / la_norm(cdata)
output.final_state = Qobj(cdata, dims=dims)
return output