content: from ..integrator import IntegratorException, Integrator
  import numpy as np
  from qutip.core import data as _data
  from scipy.optimize import root_scalar
  from ..sesolve import SESolver
  
  
  __all__ = ["IntegratorKrylov"]
  
  
  class IntegratorKrylov(Integrator):
      """
      Evolve the state vector ("psi0") finding an approximation for the time
      evolution operator of Hamiltonian ("H") by obtaining the projection of
      the time evolution operator on a set of small dimensional Krylov
      subspaces (m << dim(H)).
      """
      integrator_options = {
          'atol': 1e-7,
          'nsteps': 100,
          'min_step': 1e-5,
          'max_step': 1e5,
          'krylov_dim': 0,
          'sub_system_tol': 1e-7,
          'always_compute_step': False,
      }
      support_time_dependant = False
      supports_blackbox = False
      method = 'krylov'
  
      def _prepare(self):
          if not self.system.isconstant:
              raise ValueError("krylov method only support constant system.")
          self._max_step = -np.inf
          krylov_dim = self.options["krylov_dim"]
          if krylov_dim < 0 or krylov_dim > self.system.shape[0]:
              raise ValueError("The options 'krylov_dim', must be a positive "
                               "integer smaller that the system size")
  
          if krylov_dim == 0:
              # TODO: krylov_dim, max_step and error (atol) are related by
              # err ~= exp(-krylov_dim / dt**(1~2))
              # We could ask for 2 and determine the third one.
              N = self.system.shape[0]
              krylov_dim = min(int((N + 100)**0.5), N-1)
              self.options["krylov_dim"]  = krylov_dim
  
          if not self.options["always_compute_step"]:
              from qutip import rand_ket
              N = self.system.shape[0]
              krylov_tridiag, krylov_basis = \
                  self._lanczos_algorithm(rand_ket(N).data)
              if (
                  krylov_tridiag.shape[0] < krylov_dim
                  or krylov_tridiag.shape[0] == N
              ):
                  self._max_step = np.inf
              else:
                  self._max_step = self._compute_max_step(krylov_tridiag,
                                                          krylov_basis)
  
      def _lanczos_algorithm(self, psi):
          """
          Computes a basis of the Krylov subspace for Hamiltonian 'H', a system
          state 'psi' and Krylov dimension 'krylov_dim'. The space is spanned
          by {psi, H psi, H^2 psi, ..., H^(krylov_dim) psi}.
  
          Parameters
          ------------
          psi: np.ndarray
              State used to calculate Krylov subspace.
          """
          krylov_dim = self.options['krylov_dim']
          H = (1j * self.system(0)).data
  
          v = []
          T_diag = np.zeros(krylov_dim + 1, dtype=complex)
          T_subdiag = np.zeros(krylov_dim + 1, dtype=complex)
  
          w_prime = _data.matmul(H, psi)
          T_diag[0] = _data.inner(w_prime, psi)
          v.append(psi)
          w = _data.add(w_prime, v[-1], -T_diag[0])
          T_subdiag[0] = _data.norm.l2(w)
          j = 0
  
          while j < krylov_dim and T_subdiag[j] > self.options['sub_system_tol']:
              j += 1
              v.append(_data.mul(w, 1 / T_subdiag[j-1]))
              w_prime = _data.matmul(H, v[-1])
              T_diag[j] = _data.inner(w_prime, v[-1])
              w = _data.add(w_prime, v[-1], -T_diag[j])
              w = _data.add(w, v[-2], -T_subdiag[j-1])
              T_subdiag[j] = _data.norm.l2(w)
  
          krylov_tridiag = _data.diag["dense"](
              [T_subdiag[:j], T_diag[:j+1], T_subdiag[:j]],
              [-1, 0, 1]
          )
          krylov_basis = _data.Dense(np.hstack([psi.to_array() for psi in v]))
  
          return krylov_tridiag, krylov_basis
  
      def _compute_krylov_set(self, krylov_tridiag, krylov_basis):
          """
          Compute the eigen energies, basis transformation operator (U) and e0.
          """
          eigenvalues, eigenvectors = _data.eigs(krylov_tridiag, True)
          N = eigenvalues.shape[0]
          U = _data.matmul(krylov_basis, eigenvectors)
          e0 = eigenvectors.adjoint() @ _data.one_element_dense((N, 1), (0, 0), 1.0)
          return eigenvalues, U, e0
  
      def _compute_psi(self, dt, eigenvalues, U, e0):
          """
          compute the state at time ``t``.
          """
          phases = _data.Dense(np.exp(-1j * dt * eigenvalues))
          aux = _data.multiply(phases, e0)
          return _data.matmul(U, aux)
  
      def _compute_max_step(self, krylov_tridiag, krylov_basis, krylov_state=None):
          """
          Compute the maximum step length to stay under the desired tolerance.
          """
          if not krylov_state:
              krylov_state = self._compute_krylov_set(krylov_tridiag, krylov_basis)
  
          small_tridiag = _data.Dense(krylov_tridiag.as_ndarray()[:-1, :-1])
          small_basis = _data.Dense(krylov_basis.as_ndarray()[:, :-1])
          reduced_state = self._compute_krylov_set(small_tridiag, small_basis)
  
          def krylov_error(t):
              # we divide by atol and take the log so that the error returned is 0
              # at atol, which is convenient for calling root_scalar with.
              return np.log(_data.norm.l2(
                  self._compute_psi(t, *krylov_state) -
                  self._compute_psi(t, *reduced_state)
              ) / self.options["atol"])
  
          # Under 0 will cause an infinite loop in the while loop bellow.
          dt = max(self.options["min_step"], 1e-14)
          max_step = max(self.options["max_step"], dt)
          err = krylov_error(dt)
          if err > 0:
              raise ValueError(
                  f"With the krylov dim of {self.options['krylov_dim']}, the "
                  f"error with the minimum step {dt} is {err}, higher than the "
                  f"desired tolerance of {self.options['atol']}."
              )
  
          while krylov_error(dt * 10) < 0 and dt < max_step:
              dt *= 10
  
          if dt > max_step:
              return max_step
  
          sol = root_scalar(f=krylov_error, bracket=[dt, dt * 10],
                            method="brentq", xtol=self.options['atol'])
          if sol.converged:
              return sol.root
          else:
              return dt
  
      def set_state(self, t, state0):
          self._t_0 = t
          krylov_tridiag, krylov_basis = self._lanczos_algorithm(state0)
          self._krylov_state = self._compute_krylov_set(krylov_tridiag, krylov_basis)
  
          if (
              krylov_tridiag.shape[0] <= self.options['krylov_dim']
              or krylov_tridiag.shape == self.system.shape
          ):
              # happy_breakdown
              self._max_step = np.inf
              return
  
          if (
              not np.isfinite(self._max_step)
              or self.options["always_compute_step"]
          ):
              self._max_step = self._compute_max_step(
                  krylov_tridiag, krylov_basis, self._krylov_state,
              )
  
      def get_state(self, copy=True):
          return self._t_0, self._compute_psi(0, *self._krylov_state)
  
      def integrate(self, t, copy=True):
          step = 0
          while t > self._t_0 + self._max_step:
              # The approximation in only valid in the range t_0, t_0 + max step
              # If outside, advance the range
              step += 1
              if step >= self.options["nsteps"]:
                  raise IntegratorException(f"Maximum number of integration steps ({self.options['nsteps']}) exceeded")
              new_psi = self._compute_psi(self._max_step, *self._krylov_state)
              self.set_state(self._t_0 + self._max_step, new_psi)
  
          delta_t = t - self._t_0
          out = self._compute_psi(delta_t, *self._krylov_state)
          return t, out
  
      @property
      def options(self):
          """
          Supported options by krylov method:
  
          atol : float, default: 1e-7
              Absolute tolerance.
  
          nsteps : int, default: 100
              Max. number of internal steps/call.
  
          min_step, max_step : float, default: (1e-5, 1e5)
              Minimum and maximum step size.
  
          krylov_dim: int, default: 0
              Dimension of Krylov approximation subspaces used for the time
              evolution approximation. If the defaut 0 is given, the dimension is calculated
              from the system size N, using `min(int((N + 100)**0.5), N-1)`.
  
          sub_system_tol: float, default: 1e-7
              Tolerance to detect a happy breakdown. A happy breakdown occurs
              when the initial ket is in a subspace of the Hamiltonian smaller
              than ``krylov_dim``.
  
          always_compute_step: bool, default: False
              If True, the step length is computed each time a new Krylov
              subspace is computed. Otherwise it is computed only once when
              creating the integrator.
          """
          return self._options
  
      @options.setter
      def options(self, new_options):
          Integrator.options.fset(self, new_options)
  
  
  SESolver.add_integrator(IntegratorKrylov, 'krylov')
  