# Source code for krotov.propagators

r"""Routines that can be passed as propagator to :func:.optimize_pulses

The numerical effort involved in the optimization is almost entirely within the
simulation of the system dynamics. In every iteration and for every objective,
the system must be "propagated" once forwards in time and once backwards in
time, see also :mod:krotov.parallelization.

The implementation of this time propagation must be inside the user-supplied
routine propagator that is passed to :func:.optimize_pulses and must
calculate the propagation over a single time step. In particular,
:func:qutip.mesolve.mesolve is not automatically used for simulating any
dynamics within the optimization.  The signature for any propagator
must be the same as the "reference" :func:expm propagator:

>>> str(inspect.signature(krotov.propagators.expm))
'(H, state, dt, c_ops=None, backwards=False, initialize=False)'

The arguments are as follows (cf. :class:Propagator):

* H is the system Hamiltonian or Liouvillian, in a nested-list format similar
to that used by :func:qutip.mesolve.mesolve, e.g., for a Hamiltonian
$\Op{H} = \Op{H}_0 + c \Op{H}_1$, where $c$ is the value of a control field
at a particular point in time, propagator would receive a list [H0, [H1,
c]] where H0 and H1 are :class:qutip.Qobj operators.
The nested-list for H used here, with scalar values for the controls, is
obtained internally from the format used by :func:~qutip.mesolve.mesolve,
with time-dependent controls over the entire time grid, via
:func:krotov.conversions.plug_in_pulse_values.
* state is the :class:qutip.Qobj state that should be propagated, either a
Hilbert space state, or a density matrix.
* dt is the time step (a float). It is always positive, even for
backwards=True.
* c_ops is None, or a list of collapse (Lindblad) operators, where each list
element is a :class:qutip.Qobj instance (or possibly a nested list, for
time-dependent Lindblad operators. Note that is generally preferred for H
to be a Liouvillian, for dissipative dynamics.
* backwards (:class:bool): If passed as True, the propagator should
propagate backwards in time. In Hilbert space, this means using -dt instead
of dt. In Liouville space, there is no difference between forward and
backward propagation. In the context of Krotov's method, the backward
propagation uses the conjugate Hamiltonian, respectively Liouvillian.
However, the propagator routine does not need to be aware of this fact: it
will receive the appropriate H and c_ops.
* initialize (:class:bool): A flag to indicate the beginning of a
propagation over a time grid. If False in subsequent calls, the propagator
may assume that the input state is the result of the previous call to
propagator.

The routines in this module are provided with no guarantee to be either
general or efficient. The :func:expm propagator is exact to machine
precision, but generally extremely slow.  For "production use", it is
recommended to supply a problem-specific propagator that is highly optimized
for speed. You might consider the use of Cython_. This is key to minimize the
runtime of the optimization.

The initialize flag enables "stateful" propagators that cache data between
calls. This can significantly improve numerical efficiency.
:class:DensityMatrixODEPropagator is an example for such a propagator. In
general, any stateful propagator should be an instance of
:class:Propagator.

.. _Cython: https://cython.org
"""
from abc import ABC, abstractmethod

import numpy as np
import qutip
import scipy
from qutip.cy.spconvert import dense2D_to_fastcsr_fmode
from qutip.cy.spmatfuncs import spmvpy_csr
from qutip.superoperator import mat2vec, vec2mat

__all__ = ['expm', 'Propagator', 'DensityMatrixODEPropagator']

[docs]def expm(H, state, dt, c_ops=None, backwards=False, initialize=False):
"""Propagate using matrix exponentiation

This supports H being a Hamiltonian (for a Hilbert space state) or a
Liouvillian (for state being a density matrix) in nested-list format.
Collapse operators c_ops are not supported. The propagator is not
stateful, thus initialize is ignored.
"""
if c_ops is None:
c_ops = []
if len(c_ops) > 0:
raise NotImplementedError("Liouville exponentiation not implemented")
assert isinstance(H, list) and len(H) > 0
eqm_factor = -1j  # factor in front of H on rhs of the equation of motion
if isinstance(H[0], list):
if H[0][1].type == 'super':
eqm_factor = 1
if backwards:
eqm_factor = eqm_factor.conjugate()
A = (eqm_factor * H[0][1]) * H[0][0]
else:
if H[0].type == 'super':
eqm_factor = 1
if backwards:
eqm_factor = eqm_factor.conjugate()
A = eqm_factor * H[0]
for part in H[1:]:
if isinstance(part, list):
A += (eqm_factor * part[1]) * part[0]
else:
A += eqm_factor * part
ok_types = (state.type == 'oper' and A.type == 'super') or (
state.type in ['ket', 'bra'] and A.type == 'oper'
)
if ok_types:
return ((A * dt).expm())(state)
else:
raise NotImplementedError(
"Cannot handle argument types A:%s, state:%s"
% (A.type, state.type)
)

[docs]class Propagator(ABC):
"""Abstract base class for stateful propagators"""

[docs]    @abstractmethod
def __call__(
self, H, state, dt, c_ops=None, backwards=False, initialize=False
):
"""Evaluation of a single propagation step

Args:
H (list): A Hamiltonian or Liouvillian in qutip's nested-list
format, with a scalar value in the place of a time-dependency.
For example, [H0, [H1, u]] for a drift Hamiltonian H0,
a control Hamiltonian H1, and a scalar value u that is
a time-dependent control evaluated for a particular point in
time.
state (qutip.Qobj): The state to propagate
dt (float): The time step over which to propagate
c_ops (list or None): A list of Lindblad operators. Using explicit
Lindblad operators should be avoided: it is usually more
efficient to convert them into a Lindbladian, passed as H
backwards (bool): Whether the propagation is forward in time or
backward in time
initialize (bool): Whether the propagator should (re-)initialize
for a new propagation, when the propagator is used to advance
on a time grid, initialize should be passed as True for the
initial time step (0 to dt in a forward propagation, or T to
T-dt for a backward propagation), and False otherwise.

Note:
A propagator may assume the propagation to be "sequential"
when initialize is False. That is, the state to propagate is the
result of the previous call to the propagator.
"""
pass

[docs]class DensityMatrixODEPropagator(Propagator):
"""Propagator for density matrix evolution under a Lindbladian

See :class:qutip.solver.Options for all arguments except reentrant.
Passing True for the reentrant re-initializes the propagator in every
time step.

Warning:
By default, the propagator is not "re-entrant". That is, you cannot use
more than one instance of :class:DensityMatrixODEPropagator in the
same process at the same time. This limitation is due to
:class:scipy.integrate.ode with the "zvode" integrator not being
re-entrant. Passing reentrant=True side-steps this problem by
re-initializating :class:scipy.integrate.ode in every time step. This
makes it possible to use :class:DensityMatrixODEPropagator in the
optimization of multiple objectives, but creates a significant
"""

def __init__(
self,
order=12,
atol=1e-8,
rtol=1e-6,
nsteps=1000,
first_step=0,
min_step=0,
max_step=0,
reentrant=False,
):
self.method = method
self.order = order
self.atol = atol
self.rtol = rtol
self.nsteps = nsteps
self.first_step = first_step
self.min_step = min_step
self.max_step = max_step
self._L_list = None  # cached Liouvillian data
self._control_indices = None  # which indices in L have a control val
self._r = None  # the integrator
self._t = 0.0  # time up to which we've integrated
self._y = None  # current vectorized state
self.reentrant = reentrant

[docs]    def __call__(
self, H, state, dt, c_ops=None, backwards=False, initialize=False
):
"""Evaluation of a single propagation step

Args:
H (list): A Liouvillian superoperator in qutip's nested-list
format, with a scalar value in the place of a time-dependency.
For example, [L0, [L1, u]] for a drift Liouvillian L0,
a control Liouvillian H1, and a scalar value u that is
a time-dependent control evaluated for a particular point in
time. If initialize is False, only the control values are
taken into account; any operators are assumed to be identical
to the internally cached values of H during initialization.
state (qutip.Qobj): The density matrix to propagate. The passed
value is ignored unless initialize is given as True.
Otherwise, it is assumed that state matches the (internally
stored) state that was the result from the previous propagation
step.
dt (float): The time step over which to propagate
c_ops (list or None): An empty list, or None. Since this propagator
assumes a full Liouvillian, it cannot be combined with Lindblad
operators.
backwards (bool): Whether the propagation is forward in time or
backward in time. Since the equation of motion for a
Liouvillian and conjugate Liouvillian is the same, this
parameter has no effect. Instead, for the backward propagation,
the conjugate Liouvillian must be passed for L.
initialize (bool): Whether to (re-)initialize for a new
propagation. This caches H (except for the control values)
and state internally.
"""
# H is really an L, but it's a very bad idea for a subclass not to
# follow the interface of the parent (including argument names).
# Internally, however, we'll use L instead of H
if initialize or self.reentrant:
self._initialize(H, state, dt, c_ops, backwards)
else:
if self.reentrant:
self._initialize_integrator(self._y)
# only update the control values
for i in self._control_indices:
self._L_list[i][1] = H[i][1]
self._t += dt
self._r.integrate(self._t)
self._y = self._r.y
return qutip.Qobj(
dense2D_to_fastcsr_fmode(
vec2mat(self._y), state.shape[0], state.shape[1]
),
dims=state.dims,
isherm=True,
)

@staticmethod
def _rhs(t, rho, L_list):
# _rhs being a staticmethod enables the propagator to be pickled (for
# parallelization)
out = np.zeros(rho.shape[0], dtype=complex)
L = L_list[0][0]
L_coeff = L_list[0][1]
spmvpy_csr(L.data, L.indices, L.indptr, rho, L_coeff, out)
for n in range(1, len(L_list)):
L = L_list[n][0]
L_coeff = L_list[n][1]
spmvpy_csr(L.data, L.indices, L.indptr, rho, L_coeff, out)
return out

def _initialize(self, L, rho, dt, c_ops, backwards):
self._initialize_data(L, rho, dt, c_ops, backwards)
self._initialize_integrator(self._y)

def _initialize_data(self, L, rho, dt, c_ops, backwards):
L_list = []
control_indices = []
if not (c_ops is None or len(c_ops) == 0):
# in principle, we could convert c_ops to a Lindbladian, here
raise NotImplementedError("c_ops not implemented")
for (i, spec) in enumerate(L):
if isinstance(spec, qutip.Qobj):
l_op = spec
l_coeff = 1
elif isinstance(spec, list) and isinstance(spec[0], qutip.Qobj):
l_op = spec[0]
l_coeff = spec[1]
control_indices.append(i)
else:
raise ValueError(
"Incorrect specification of time-dependent Liouvillian"
)
if l_op.type == 'super':
L_list.append([l_op.data, l_coeff, False])
else:
raise ValueError(
"Incorrect specification of time-dependent Liouvillian"
)
self._L_list = L_list
self._control_indices = control_indices
if rho.type == 'oper':
self._y = mat2vec(rho.full()).ravel('F')  # initial state
else:
raise ValueError("rho must be a density matrix")

def _initialize_integrator(self, initial_vector):
r = scipy.integrate.ode(self._rhs)
r.set_integrator(
'zvode',
method=self.method,
order=self.order,
atol=self.atol,
rtol=self.rtol,
nsteps=self.nsteps,
first_step=self.first_step,
min_step=self.min_step,
max_step=self.max_step,
)
r.set_initial_value(initial_vector)
r.set_f_params(self._L_list)
self._r = r
self._t = 0.0