Trajectree 0.0.0__py3-none-any.whl → 0.0.1__py3-none-any.whl

trajectree/quimb/quimb/evo.py

@@ -0,0 +1,712 @@
+"""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 (
+    dag,
+    dot,
+    explt,
+    eye,
+    isop,
+    issparse,
+    ldmul,
+    make_immutable,
+    qarray,
+    qu,
+    rdmul,
+)
+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
+
+
+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
+
+
+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
+
+
+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
+
+
+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
+
+
+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
+
+
+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
+
+
+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                                                     #
+# --------------------------------------------------------------------------- #
+
+
+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.toarray())
+            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.toarray().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)
+
+    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)
+
+    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