Source code for qutip.solver.integrator.krylov

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"]

[docs]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] < self.options["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"]) dt = self.options["min_step"] err = krylov_error(dt) if err > 0: 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 < self.options["max_step"]: dt *= 10 if dt > self.options["max_step"]: return self.options["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')