content: # Required for Sphinx to follow autodoc_type_aliases
  from __future__ import annotations
  
  __all__ = ['Propagator', 'propagator', 'propagator_steadystate']
  
  import numbers
  import numpy as np
  
  from .. import Qobj, qeye, qeye_like, unstack_columns, QobjEvo, liouvillian
  from ..core import data as _data
  from ..typing import QobjEvoLike
  from .mesolve import mesolve, MESolver
  from .sesolve import sesolve, SESolver
  from .heom.bofin_solvers import HEOMSolver
  from .solver_base import Solver
  from .multitraj import MultiTrajSolver
  from numbers import Number
  from typing import Any
  from qutip.settings import settings
  
  
  def propagator_piecewise(
      H: QobjEvo,
      tlist: list[Number],
      piecewise_t: list[Number],
      c_ops: list[QobjEvo] = None,
      args: dict[str, Any] = None,
  ) -> Qobj | list[Qobj]:
      r"""
      Calculate propagators for piecewise constant Hamiltonians or
      Liouvillians.
  
      This function efficiently computes propagators when the Hamiltonian or
      Liouvillian is piecewise constant, i.e., it remains constant over
      certain time intervals and changes discontinuously at specified times.
      By exponentiating the generator directly on each constant interval,
      this approach is significantly faster than numerically integrating the
      differential equation.
  
      Parameters
      ----------
      H : :obj:`.QobjEvo`
          Time-dependent Hamiltonian or Liouvillian. Must be callable with
          signature ``H(t, args)`` that returns the constant operator valid at
          time ``t``.
  
      tlist : list of float
          List of times at which to evaluate the propagator.
  
      piecewise_t : list of float
          Times where the Hamiltonian or Liouvillian or collapse operators change
          discontinuously.  These times define the boundaries between constant
          intervals.
          Times outside the range ``(tlist[0], tlist[-1]]`` are ignored.
  
      c_ops : list of :obj:`.QobjEvo`, optional
          List of collapse operators for open quantum systems.
          Can be piecewise constant.
  
      args : dict, optional
          Dictionary of parameters to pass to the Hamiltonian and collapse
          operator functions.
  
      Returns
      -------
      propagators : list of :obj:`.Qobj`
          List of propagators :math:`U(t)` evaluated at each time in ``tlist``,
          such that :math:`\psi(t) = U(t)\psi(t_0)` for state vectors or
          :math:`\rho_{\mathrm{vec}}(t) = U(t)\rho_{\mathrm{vec}}(t_0)` for
          vectorized density matrices.
      """
  
      piecewise_times = {tt for tt in piecewise_t if tlist[0] < tt <= tlist[-1]}
      eval_times = set(tlist)
      times = sorted(eval_times | piecewise_times)
      atol = settings.core["atol"]
  
      out = []
      prev = times[0]
      prev_dt = None
      dU = None
  
      for nxt in times[1:]:
          dt = nxt - prev
  
          # Evaluate at midpoint to avoid discontinuities at boundaries
          t_eval = (prev + nxt) / 2
          H_step = H(t_eval, args)
          if c_ops:
              c_ops_q = [op(t_eval, args) for op in c_ops]
              gen = liouvillian(H_step, c_ops_q)
          else:
              gen = H_step if H_step.issuper else -1j * H_step
  
          if prev == times[0]:
              U = qeye_like(gen)
              if prev in eval_times:
                  out.append(U)
  
          cannot_reuse = (
              dU is None
              or prev_dt is None
              or abs(dt - prev_dt) > atol
              or prev in piecewise_times  # Moved past a switching point
          )
  
          if cannot_reuse:
              dU = (gen * dt).expm()
          U = dU @ U
          prev_dt = dt
  
          if nxt in eval_times:
              out.append(U)
          prev = nxt
      return out
  
  
  def propagator(
      H: QobjEvoLike,
      t: Number | list[Number],
      c_ops: QobjEvoLike | list[QobjEvoLike] = None,
      args: dict[str, Any] = None,
      options: dict[str, Any] = None,
      piecewise_t: list[Number] | None = None,
      **kwargs,
  ) -> Qobj | list[Qobj]:
      r"""
      Calculate the propagator U(t) for the density matrix or wave function such
      that :math:`\psi(t) = U(t)\psi(0)` or
      :math:`\rho_{\mathrm vec}(t) = U(t) \rho_{\mathrm vec}(0)`
      where :math:`\rho_{\mathrm vec}` is the vector representation of the
      density matrix.
  
      Parameters
      ----------
      H : :obj:`.Qobj`, :obj:`.QobjEvo`, :obj:`.QobjEvo` compatible format
          Possibly time-dependent system Liouvillian or Hamiltonian as a Qobj or
          QobjEvo. ``list`` of [:obj:`.Qobj`, :obj:`.Coefficient`] or callable
          that can be made into :obj:`.QobjEvo` are also accepted.
  
      t : float or array-like
          Time or list of times for which to evaluate the propagator. If a single
          time ``t`` is passed, the propagator from ``0`` to ``t`` is computed.
          When ``t`` is a list, the propagators from the first time in the list
          to each elements in ``t`` is returned. In that case, the first output
          will always be the identity matrix.
  
      c_ops : list, optional
          List of collapse operators as Qobj, QobjEvo or list that can be made
          into QobjEvo.
  
      args : dictionary, optional
          Parameters to callback functions for time-dependent Hamiltonians and
          collapse operators.
  
      options : dict, optional
          Options for the solver.
  
      piecewise_t : list, optional
          Times where the Hamiltonian or Liouvillian change values when
          they are piecewise constant. Providing these allows for a faster
          computation by exponentiating the Liouvillian or Hamiltonian directly
          on each interval.
  
      **kwargs :
          Extra parameters to use when creating the
          :obj:`.QobjEvo` from a list format ``H``. The most common are ``tlist``
          and ``order`` for array-based time dependance. See :obj:`.QobjEvo` for
          the full list.
  
      Returns
      -------
      U : :obj:`.Qobj`, list
          Instance representing the propagator(s) :math:`U(t)`. Return a single
          Qobj when ``t`` is a number or a list when ``t`` is a list.
  
      Notes
      -----
      Unlike :func:`.sesolve` or :func:`.mesolve`, the output times in ``t`` are
      not used for array time dependent system. ``tlist`` must be passed as a
      keyword argument in those case. ``tlist`` and ``t`` can have different
      length and values.
  
      """
      if isinstance(t, numbers.Real):
          tlist = [0, t]
          list_output = False
      else:
          tlist = t
          list_output = True
  
      if not isinstance(H, (Qobj, QobjEvo)):
          H = QobjEvo(H, args=args, **kwargs)
  
      if isinstance(c_ops, list):
          c_ops = [QobjEvo(op, args=args, **kwargs) for op in c_ops]
      elif c_ops is not None:
          c_ops = [QobjEvo(c_ops, args=args, **kwargs)]
  
      if piecewise_t is not None:
          out = propagator_piecewise(H, tlist, piecewise_t, c_ops, args)
          return out if list_output else out[-1]
  
      if c_ops:
          H = liouvillian(H, c_ops)
  
      U0 = qeye_like(H)
  
      if H.issuper:
          out = mesolve(H, U0, tlist, args=args, options=options).states
      else:
          out = sesolve(H, U0, tlist, args=args, options=options).states
  
      if list_output:
          return out
      else:
          return out[-1]
  
  
  def propagator_steadystate(U: Qobj) -> Qobj:
      r"""Find the steady state for successive applications of the propagator
      :math:`U`.
  
      Parameters
      ----------
      U : :obj:`.Qobj`
          Operator representing the propagator.
  
      Returns
      -------
      a : :obj:`.Qobj`
          Instance representing the steady-state density matrix.
      """
      evals, estates = U.eigenstates()
      shifted_vals = np.abs(evals - 1.0)
      ev_idx = np.argmin(shifted_vals)
      rho_data = unstack_columns(estates[ev_idx].data)
      rho_data = _data.mul(rho_data, 0.5 / _data.trace(rho_data))
      return Qobj(_data.add(rho_data, _data.adjoint(rho_data)),
                  dims=U._dims[0].oper,
                  isherm=True,
                  copy=False)
  
  
  class Propagator:
      """
      A generator of propagator for a system.
  
      Usage:
  
          U = Propagator(H, c_ops)
  
          psi_t = U(t) @ psi_0
  
      Save some previously computed propagator are stored to speed up subsequent
      computation. Changing ``args`` will erase these stored probagator.
  
      Parameters
      ----------
      system : :obj:`.Qobj`, :obj:`.QobjEvo`, :class:`.Solver`
          Possibly time-dependent system driving the evolution, either already
          packaged in a solver, such as :class:`.SESolver` or :class:`.BRSolver`,
          or the Liouvillian or Hamiltonian as a :obj:`.Qobj`,
          :obj:`.QobjEvo`. ``list`` of [:obj:`.Qobj`, :obj:`.Coefficient`]
          or callable that can be made into :obj:`.QobjEvo` are also accepted.
  
          Solvers that run non-deterministacilly, such as :class:`.MCSolver`, are
          not supported.
  
      c_ops : list, optional
          List of :obj:`.Qobj` or :obj:`.QobjEvo` collapse operators.
  
      args : dictionary, optional
          Parameters to callback functions for time-dependent Hamiltonians and
          collapse operators.
  
      options : dict, optional
          Options for the solver.
  
      memoize : int, default: 10
          Max number of propagator to save.
  
      tol : float, default: 1e-14
          Absolute tolerance for the time. If a previous propagator was computed
          at a time within tolerance, that propagator will be returned.
  
      Notes
      -----
      The :class:`Propagator` is not a :obj:`.QobjEvo` so
      it cannot be used for operations with :obj:`.Qobj` or
      :obj:`.QobjEvo`. It can be made into a
      :obj:`.QobjEvo` with ::
  
          U = QobjEvo(Propagator(H))
  
      """
      def __init__(
          self,
          system: Qobj | QobjEvo | Solver,
          *,
          c_ops: QobjEvoLike | list[QobjEvoLike] = None,
          args: dict[str, Any] = None,
          options: dict[str, Any] = None,
          memoize: int = 10,
          tol: float = 1e-14,
      ):
          if isinstance(system, MultiTrajSolver):
              raise TypeError("Non-deterministic solvers cannot be used "
                              "as a propagator system")
          elif isinstance(system, HEOMSolver):
              raise NotImplementedError(
                  "HEOM is not supported by Propagator. "
                  "Please, tell us on GitHub issues if you need it!"
              )
          elif isinstance(system, Solver):
              self.solver = system
          else:
              Hevo = QobjEvo(system, args=args)
              c_ops = c_ops if c_ops is not None else []
              c_ops = [QobjEvo(op, args=args) for op in c_ops]
              if Hevo.issuper or c_ops:
                  self.solver = MESolver(Hevo, c_ops=c_ops, options=options)
              else:
                  self.solver = SESolver(Hevo, options=options)
  
          self.times = [0]
          self.invs = [None]
          self.props = [qeye(self.solver._sys_dims)]
          self.solver.start(self.props[0], self.times[0])
          self.cte = self.solver.rhs.isconstant
          H_0 = self.solver.rhs(0)
          self.unitary = not H_0.issuper and H_0.isherm
          self.args = args
          self.memoize = max(3, int(memoize))
          self.tol = tol
  
      def _lookup_or_compute(self, t):
          """
          Get U(t) from cache or compute it.
          """
          idx = np.searchsorted(self.times, t)
          if idx < len(self.times) and abs(t-self.times[idx]) <= self.tol:
              U = self.props[idx]
          elif idx > 0 and abs(t-self.times[idx-1]) <= self.tol:
              U = self.props[idx-1]
          else:
              U = self._compute(t, idx)
              self._insert(t, U, idx)
          return U
  
      def __call__(self, t: float, t_start: float = 0, **args):
          """
          Get the propagator from ``t_start`` to ``t``.
  
          Parameters
          ----------
          t : float
              Time at which to compute the propagator.
          t_start: float [0]
              Time at which the propagator start such that:
                  ``psi[t] = U.prop(t, t_start) @ psi[t_start]``
          args : dict
              Argument to pass to a time dependent Hamiltonian.
              Updating ``args`` take effect since ``t=0`` and the new ``args``
              will be used in future call.
          """
          # We could improve it when the system is constant using U(2t) = U(t)**2
          if not self.cte and args and args != self.args:
              self.args = args
              self.solver._argument(args)
              self.times = [0]
              self.props = [qeye_like(self.props[0])]
              self.solver.start(self.props[0], self.times[0])
  
          if t_start:
              if t == t_start:
                  U = self._lookup_or_compute(0)
              if self.cte:
                  U = self._lookup_or_compute(t - t_start)
              else:
                  Uinv = self._inv(self._lookup_or_compute(t_start))
                  U = self._lookup_or_compute(t) @ Uinv
          else:
              U = self._lookup_or_compute(t)
          return U
  
      def inv(self, t: float, **args):
          """
          Get the inverse of the propagator at ``t``, such that
              ``psi_0 = U.inv(t) @ psi_t``
  
          Parameters
          ----------
          t : float
              Time at which to compute the propagator.
          args : dict
              Argument to pass to a time dependent Hamiltonian.
              Updating ``args`` take effect since ``t=0`` and the new ``args``
              will be used in future call.
          """
          return self._inv(self(t, **args))
  
      def _compute(self, t, idx):
          """
          Compute the propagator at ``t``, ``idx`` point to a pair of
          (time, propagator) close to the desired time.
          """
          t_last = self.solver._integrator.get_state(copy=False)[0]
          if self.times[idx-1] <= t_last <= t:
              U = self.solver.step(t)
          elif idx > 0:
              self.solver.start(self.props[idx-1], self.times[idx-1])
              U = self.solver.step(t)
          else:
              # Evolving backward in time is not supported by all integrator.
              self.solver.start(qeye_like(self.props[0]), t)
              Uinv = self.solver.step(self.times[idx])
              U = self._inv(Uinv)
          return U
  
      def _inv(self, U):
          return U.dag() if self.unitary else U.inv()
  
      def _insert(self, t, U, idx):
          """
          Insert a new pair of (time, propagator) to the memorized states.
          """
          while len(self.times) >= self.memoize:
              rm_idx = self.memoize // 2
              if self.times[rm_idx] < t:
                  idx -= 1
              del self.times[rm_idx]
              del self.props[rm_idx]
          self.times.insert(idx, t)
          self.props.insert(idx, U)
  