"""Easy and efficient time evolutions.
Contains an evolution class, Evolution to easily and efficiently manage time
evolution of quantum states according to the Schrodinger equation,
and related functions.
"""
import functools
import numpy as np
from scipy.integrate import complex_ode
from scipy.sparse.linalg import LinearOperator
from .core import (qarray, isop, ldmul, rdmul, explt,
dot, issparse, qu, eye, dag, make_immutable)
from .linalg.base_linalg import eigh, norm, expm_multiply, Lazy
from .linalg.approx_spectral import norm_fro_approx
from .utils import continuous_progbar, progbar, ensure_dict
CALLABLE_TIME_INDEP_CLASSES = (LinearOperator, Lazy)
# --------------------------------------------------------------------------- #
# Quantum evolution equations #
# --------------------------------------------------------------------------- #
#
# This are mostly just to be used internally with the integrators
[docs]def schrodinger_eq_ket(ham):
"""Wavefunction schrodinger equation.
Parameters
----------
ham : operator
Time-independant Hamiltonian governing evolution.
Returns
-------
psi_dot(t, y) : callable
Function to calculate psi_dot(t) at psi(t).
"""
def psi_dot(_, y):
return -1.0j * dot(ham, y)
return psi_dot
[docs]def schrodinger_eq_ket_timedep(ham):
"""Wavefunction time dependent schrodinger equation.
Parameters
----------
ham : callable
Time-dependant Hamiltonian governing evolution, such that ``ham(t)``
returns an operator representation of the Hamiltonian at time ``t``.
Returns
-------
psi_dot(t, y) : callable
Function to calculate psi_dot(t) at psi(t).
"""
def psi_dot(t, y):
return -1.0j * dot(ham(t), y)
return psi_dot
[docs]def schrodinger_eq_dop(ham):
"""Density operator schrodinger equation, but with flattened input/output.
Note that this assumes both `ham` and `rho` are hermitian in order to speed
up the commutator, non-hermitian hamiltonians as used to model loss should
be treated explicilty or with `schrodinger_eq_dop_vectorized`.
Parameters
----------
ham : operator
Time-independant Hamiltonian governing evolution.
Returns
-------
rho_dot(t, y) : callable
Function to calculate rho_dot(t) at rho(t), input and
output both in ravelled (1D form).
"""
d = ham.shape[0]
def rho_dot(_, y):
hrho = dot(ham, y.reshape(d, d))
return -1.0j * (hrho - hrho.T.conj()).reshape(-1)
return rho_dot
[docs]def schrodinger_eq_dop_timedep(ham):
"""Time dependent density operator schrodinger equation, but with flattened
input/output.
Note that this assumes both `ham(t)` and `rho` are hermitian in order to
speed up the commutator, non-hermitian hamiltonians as used to model loss
should be treated explicilty or with `schrodinger_eq_dop_vectorized`.
Parameters
----------
ham : callable
Time-dependant Hamiltonian governing evolution, such that ``ham(t)``
returns an operator representation of the Hamiltonian at time ``t``.
Returns
-------
rho_dot(t, y) : callable
Function to calculate rho_dot(t) at rho(t), input and
output both in ravelled (1D form).
"""
d = ham(0).shape[0]
def rho_dot(t, y):
hrho = dot(ham(t), y.reshape(d, d))
return -1.0j * (hrho - hrho.T.conj()).reshape(-1)
return rho_dot
[docs]def schrodinger_eq_dop_vectorized(ham):
"""Density operator schrodinger equation, but with flattened input/output
and vectorised superoperator mode (no reshaping required).
Note that this is probably only more efficient for sparse Hamiltonians.
Parameters
----------
ham: time-independant hamiltonian governing evolution
Returns
-------
rho_dot(t, y) : callable
Function to calculate rho_dot(t) at rho(t), input and
output both in ravelled (1D form).
"""
d = ham.shape[0]
sparse = issparse(ham)
idt = eye(d, sparse=sparse)
evo_superop = -1.0j * ((ham & idt) - (idt & ham.T))
def rho_dot(_, y):
return dot(evo_superop, y)
return rho_dot
[docs]def lindblad_eq(ham, ls, gamma):
"""Lindblad equation, but with flattened input/output.
Parameters
----------
ham : operator
Time-independant hamiltonian governing evolution.
ls : sequence of matrices
Lindblad operators.
gamma : float
Dampening strength.
Returns
-------
rho_dot(t, y) : callable
Function to calculate rho_dot(t) at rho(t), input and
output both in ravelled (1D form).
"""
d = ham.shape[0]
lls = tuple(dot(dag(l), l) for l in ls)
def gen_l_terms(rho):
for l, ll in zip(ls, lls):
yield (dot(l, dot(rho, dag(l))) -
0.5 * (dot(rho, ll) + dot(ll, rho)))
def rho_dot(_, y):
rho = y.reshape(d, d)
rho_d = dot(ham, rho)
rho_d -= rho_d.T.conj()
rho_d *= -1.0j
rho_d += gamma * sum(gen_l_terms(rho))
return np.asarray(rho_d).reshape(-1)
return rho_dot
[docs]def lindblad_eq_vectorized(ham, ls, gamma, sparse=False):
"""Lindblad equation, but with flattened input/output and vectorised
superoperation mode (no reshaping required).
Parameters
----------
ham : operator
Time-independant hamiltonian governing evolution.
ls : sequence of matrices
Lindblad operators.
gamma : float
Dampening strength.
Returns
-------
rho_dot(t, y) : callable
Function to calculate rho_dot(t) at rho(t), input and
output both in ravelled (1D form).
"""
d = ham.shape[0]
ham_sparse = issparse(ham) or sparse
idt = eye(d, sparse=ham_sparse)
evo_superop = -1.0j * ((ham & idt) - (idt & ham.T))
def gen_lb_terms():
for l in ls:
lb_sparse = issparse(l) or sparse
idt = eye(d, sparse=lb_sparse)
yield ((l & l.conj()) - 0.5 * ((idt & dot(dag(l), l).T) +
(dot(dag(l), l) & idt)))
evo_superop += gamma * sum(gen_lb_terms())
def rho_dot(_, y):
return dot(evo_superop, y)
return rho_dot
def _calc_evo_eq(isdop, issparse, isopen=False, timedep=False):
"""Choose an appropirate dynamical equation to evolve with.
"""
eq_chooser = {
(0, 0, 0, 0): schrodinger_eq_ket,
(0, 1, 0, 0): schrodinger_eq_ket,
(1, 0, 0, 0): schrodinger_eq_dop,
(1, 1, 0, 0): schrodinger_eq_dop_vectorized,
(1, 0, 1, 0): lindblad_eq,
(1, 1, 1, 0): lindblad_eq_vectorized,
# time-dependent
(0, 0, 0, 1): schrodinger_eq_ket_timedep,
(0, 1, 0, 1): schrodinger_eq_ket_timedep,
(1, 0, 0, 1): schrodinger_eq_dop_timedep,
(1, 1, 0, 1): schrodinger_eq_dop_timedep,
}
return eq_chooser[(isdop, issparse, isopen, timedep)]
class Try2Then3Args:
def __init__(self, fn):
self.fn = fn
self.num_args = None
def first_call(self, t, p, H):
try:
res = self.fn(t, p)
self.num_args = 2
except TypeError as e:
if 'positional' in e.args[0]:
res = self.fn(t, p, H)
self.num_args = 3
else:
raise
return res
def __call__(self, t, p, H):
if self.num_args is None:
return self.first_call(t, p, H)
elif self.num_args == 2:
return self.fn(t, p)
elif self.num_args == 3:
return self.fn(t, p, H)
# --------------------------------------------------------------------------- #
# Quantum Evolution Class #
# --------------------------------------------------------------------------- #
[docs]class Evolution(object):
"""A class for evolving quantum systems according to Schrodinger equation.
The evolution can be performed in a number of ways:
- diagonalise the Hamiltonian (or use already diagonalised system).
- integrate the complex ODE, that is, the Schrodinger equation, using
scipy. Here either a mid- or high-order Dormand-Prince adaptive
time stepping scheme is used (see
:class:`scipy.integrate.complex_ode`).
Parameters
----------
p0 : quantum state
Inital state, either vector or operator. If vector, converted to ket.
ham : operator, tuple (1d array, operator), or callable
Governing Hamiltonian, if tuple then assumed to contain
``(eigvals, eigvecs)`` of presolved system. If callable (but not a
SciPy ``LinearOperator``), assume a time-dependent hamiltonian such
that ``ham(t)`` is the Hamiltonian at time ``t``. In this case, the
latest call to ``ham`` will be cached (and made immutable) in case it
is needed by callbacks passed to ``compute``.
t0 : float, optional
Initial time (i.e. time of state ``p0``), defaults to zero.
compute : callable, or dict of callable, optional
Function(s) to compute on the state at each time step. Function(s)
should take args (t, pt) or (t, pt, ham) if the Hamiltonian is
required. If ham is required, it will be passed in to the function
exactly as given to this ``Evolution`` instance, except if ``method``
is ``'solve'``, in which case it will be passed in as the solved system
``(eigvals, eigvecs)``. If supplied with:
- single callable : ``Evolution.results`` will contain the results
as a list,
- dict of callables : ``Evolution.results`` will contain the
results as a dict of lists with corresponding keys to those
given in ``compute``.
int_stop : callable, optional
A condition to terminate the integration early if ``method`` is
``'integrate'``. This callable is called at every successful
integration step and should take args (t, pt) or (t, pt, ham) similar
to the function(s) in the ``compute`` argument. It should return
``-1`` to stop the integration, otherwise it should return ``None``
or ``0``.
method : {'integrate', 'solve', 'expm'}
How to evolve the system:
- ``'integrate'``: use definite integration. Get system at each
time step, only need action of Hamiltonian on state. Generally
efficient.
- ``'solve'``: diagonalise dense hamiltonian. Best for small
systems and allows arbitrary time steps without loss of
precision.
- ``'expm'``: compute the evolved state using the action of the
operator exponential in a 'single shot' style. Only needs action
of Hamiltonian, for very large systems can use distributed MPI.
int_small_step : bool, optional
If ``method='integrate'``, whether to use a low or high order
integrator to give naturally small or large steps.
expm_backend : {'auto', 'scipy', 'slepc'}
How to perform the expm_multiply function if ``method='expm'``. Can
further specifiy ``'slepc-krylov'``, or ``'slepc-expokit'``.
expm_opts : dict
Supplied to :func:`~quimb.linalg.base_linalg.expm_multiply`
function if ``method='expm'``.
progbar : bool, optional
Whether to show a progress bar when calling ``at_times`` or integrating
with the ``update_to`` method.
"""
def __init__(self, p0, ham, t0=0,
compute=None,
int_stop=None,
method='integrate',
int_small_step=False,
expm_backend='AUTO',
expm_opts=None,
progbar=False):
self._p0 = qu(p0)
self._t = self.t0 = t0
self._isdop = isop(self._p0) # Density operator evolution?
self._d = p0.shape[0] # Hilbert space dimension
self._progbar = progbar
self._timedep = (
callable(ham) and
not isinstance(ham, CALLABLE_TIME_INDEP_CLASSES)
)
if self._timedep:
# cache the time-dependent Hamiltonian in case callbacks use it
noncacheing_ham = ham
@functools.lru_cache(1)
def ham(t):
Ht = noncacheing_ham(t)
if not isinstance(Ht, LinearOperator):
make_immutable(Ht)
return Ht
if (int_stop is not None) and (method != 'integrate'):
raise ValueError("You can't provide an integration stopping "
"condition (int_stop) if the method is not "
"'integrate'")
self._setup_callback(compute, int_stop)
self._method = method
if method == 'solve' or isinstance(ham, (tuple, list)):
if isinstance(ham, LinearOperator):
raise TypeError("You can't use the 'solve' method "
"with an abstract linear operator "
"Hamiltonian.")
elif self._timedep:
raise TypeError("You can't use the 'solve' method "
"with a time-dependent Hamiltonian.")
self._ham = ham
self._setup_solved_ham()
elif method == 'integrate':
self._start_integrator(ham, int_small_step)
self._ham = ham
elif method == 'expm':
if isinstance(ham, LinearOperator):
raise TypeError("You can't use the 'expm' method "
"with an abstract linear operator "
"Hamiltonian.")
elif self._timedep:
raise TypeError("You can't use the 'expm' method "
"with a time-dependent Hamiltonian.")
self._update_method = self._update_to_expm_ket
self._pt = self._p0
self.expm_backend = expm_backend
self.expm_opts = ensure_dict(expm_opts)
self._ham = ham
else:
raise ValueError(f"Did not understand evolution "
"method: '{method}'.")
def _setup_callback(self, fn, int_stop):
"""Setup callbacks in the correct place to compute into _results
"""
# if fn is None there is no callback
if fn is None:
step_callback = None
# else fn is a dict of callbacks or a single callback
else:
# dict of funcs input -> dict of funcs output
if isinstance(fn, dict):
fn_try2then3args = {k: Try2Then3Args(v) for k, v in fn.items()}
self._results = {k: [] for k in fn}
def step_callback(t, pt, H):
for k, v in fn_try2then3args.items():
fn_result = v(t, pt, H)
self._results[k].append(fn_result)
# else results -> single list of outputs of fn
else:
fn_try2then3args = Try2Then3Args(fn)
self._results = []
def step_callback(t, pt, H):
fn_result = fn_try2then3args(t, pt, H)
self._results.append(fn_result)
# For the integration callback, additionally need to convert
# back to 'quantum' (column vector) form
# if no compute callback, check if there is an int_stop callback
if step_callback is None:
# if there is only an int_stop callback, set this up
if int_stop is not None:
int_stop_try2then3args = Try2Then3Args(int_stop)
def int_step_callback(t, y, H):
pt = qarray(y.reshape(self._d, -1))
return int_stop_try2then3args(t, pt, H)
# else if there is neither kind of callback but a progbar is
# needed, set up a dummy callback so it gets updated
elif self._progbar:
def int_step_callback(t, y, H):
pass
# else there are no callbacks and no progbar
else:
int_step_callback = None
# else there is compute callback, but may need to add int_stop callback
else:
# if both kinds of callback, combine them
if int_stop is not None:
int_stop_try2then3args = Try2Then3Args(int_stop)
def int_step_callback(t, y, H):
# For the integration callback, additionally need to
# convert back to 'quantum' (column vector) form
pt = qarray(y.reshape(self._d, -1))
step_callback(t, pt, H)
return int_stop_try2then3args(t, pt, H)
# else no int_stop callback, so just set up compute callback
else:
def int_step_callback(t, y, H):
# For the integration callback, additionally need to
# convert back to 'quantum' (column vector) form
pt = qarray(y.reshape(self._d, -1))
step_callback(t, pt, H)
self._step_callback = step_callback
self._int_step_callback = int_step_callback
def _setup_solved_ham(self):
"""Solve the hamiltonian if needed and find the initial state
in the energy eigenbasis for quick evolution later.
"""
# See if already solved from tuple
try:
evals, evecs = self._ham
self._method = 'solve'
except ValueError:
evals, evecs = eigh(self._ham.A)
self._ham = (evals, evecs)
# Find initial state in energy eigenbasis at t0
if self._isdop:
self.pe0 = dot(dag(evecs), dot(self._p0, evecs))
self._update_method = self._update_to_solved_dop
else:
self.pe0 = dot(dag(evecs), self._p0)
self._update_method = self._update_to_solved_ket
# Current state (start with same as initial)
self._pt = self._p0
def _start_integrator(self, ham, small_step):
"""Initialize a stepping integrator.
"""
if self._timedep:
H0 = ham(0.0)
else:
H0 = ham
# set complex ode with governing equation
evo_eq = _calc_evo_eq(self._isdop, issparse(H0), False, self._timedep)
self._stepper = complex_ode(evo_eq(ham))
# 5th order stpper or 8th order stepper
int_mthd, step_fct = ('dopri5', 150) if small_step else ('dop853', 50)
if isinstance(H0, LinearOperator):
# approx norm doesn't need to be very accurate
nrm0 = norm_fro_approx(H0, tol=0.1)
else:
nrm0 = norm(H0, 'f')
first_step = nrm0 / step_fct
self._stepper.set_integrator(int_mthd, nsteps=0, first_step=first_step)
# Set step_callback to be evaluated with args (t, y) at each step
if self._int_step_callback is not None:
def solout(t, y):
res = self._int_step_callback(t, y, self._ham)
return res
self._stepper.set_solout(solout)
self._stepper.set_initial_value(self._p0.A.reshape(-1), self.t0)
# assign the correct update_to method
self._update_method = self._update_to_integrate
# Methods for updating the simulation ----------------------------------- #
def _update_to_expm_ket(self, t):
"""Update the simulation to time ``t``, without explicitly computing
the operator exponential itself.
"""
factor = -1j * (t - self.t)
self._pt = expm_multiply(factor * self._ham, self._pt,
backend=self.expm_backend, **self.expm_opts)
self._t = t
# compute any callbacks into -> self._results
if self._step_callback is not None:
self._step_callback(t, self._pt, self._ham)
def _update_to_solved_ket(self, t):
"""Update simulation consisting of a solved hamiltonian and a
wavefunction to time `t`.
"""
self._t = t
evals, evecs = self._ham
lt = explt(evals, t - self.t0)
self._pt = evecs @ ldmul(lt, self.pe0)
# compute any callbacks into -> self._results
if self._step_callback is not None:
self._step_callback(t, self._pt, self._ham)
def _update_to_solved_dop(self, t):
"""Update simulation consisting of a solved hamiltonian and a
density operator to time `t`.
"""
self._t = t
evals, evecs = self._ham
lt = explt(evals, t - self.t0)
lvpvl = rdmul(ldmul(lt, self.pe0), lt.conj())
self._pt = evecs @ (lvpvl @ dag(evecs))
# compute any callbacks into -> self._results
if self._step_callback is not None:
self._step_callback(t, self._pt, self._ham)
def _update_to_integrate(self, t):
"""Update simulation consisting of unsolved hamiltonian.
"""
self._stepper.integrate(t)
[docs] def update_to(self, t):
"""Update the simulation to time ``t`` using relevant method.
Parameters
----------
t : float
Time to update the evolution to.
"""
if self._progbar and hasattr(self, '_stepper'):
with continuous_progbar(self.t, t) as pbar:
if self._int_step_callback is not None:
def solout(t, y):
int_stop_res = self._int_step_callback(t, y, self._ham)
pbar.cupdate(t)
return int_stop_res
else:
def solout(t, _):
pbar.cupdate(t)
self._stepper.set_solout(solout)
self._update_method(t)
else:
self._update_method(t)
[docs] def at_times(self, ts):
"""Generator expression to yield state af list of times.
Parameters
----------
ts : sequence of floats
Times at which to evolve to, then yield the state.
Yields
------
pt : quantum state
Quantum state of evolution at next time in ``ts``.
Notes
-----
If integrating, currently any compute callbacks will be called at every
*integration* step, not just the times `ts` -- i.e. in general
len(Evolution.results) != len(ts) and if the adaptive step times are
needed they should be added as a callback, e.g.
``compute['t'] = lambda t, _: return t``.
"""
if self._progbar:
ts = progbar(ts)
for t in ts:
self._update_method(t)
yield self.pt
# Simulation properties ------------------------------------------------- #
@property
def t(self):
"""float : Current time of simulation.
"""
return self._stepper.t if self._method == 'integrate' else self._t
@property
def pt(self):
"""quantum state : State of the system at the current time (t).
"""
if self._method == 'integrate':
return qarray(self._stepper.y.reshape(self._d, -1))
else:
return self._pt
@property
def results(self):
"""list, or dict of lists, optional : Results of the compute
callback(s) for each time step.
"""
return self._results