torch_vn_algebra.krylov¶
Source: torch_vn_algebra/krylov.py. Signatures and docstrings as in the source.
Module documentation¶
Krylov-subspace methods for large sectors, where dense matrices no longer fit.
A Hamiltonian is only needed through its action on vectors, `matvec(V)` with V of shape
(batch, k) -> (batch, k); every batch element may be a different operator (e.g. one disorder
realisation each). `SparseSectorHamiltonian` provides such an operator for spin chains: a
batch of diagonals plus one shared sparse off-diagonal part.
E0, psi, res = krylov.ground_state(H) # batched Lanczos with restarts
psi_t = krylov.evolve(H, psi0, times) # exp(-i t H) psi0, Krylov steps
Lanczos uses full re-orthogonalisation (memory batch * m * k).
SparseSectorHamiltonian¶
class SparseSectorHamiltonian(diag: torch.Tensor, offdiag: Optional[torch.Tensor], dtype: torch.dtype)
H_b = diag(d_b) + A for b = 1..batch, with d of shape (batch, k) (real) and one sparse k x k matrix A
shared by the batch (hopping terms). Memory O(batch * k + nnz).
SparseSectorHamiltonian.batch_size¶
@property
def batch_size(self) -> int
SparseSectorHamiltonian.dim¶
@property
def dim(self) -> int
SparseSectorHamiltonian.device¶
@property
def device(self)
SparseSectorHamiltonian.matvec¶
def matvec(self, V: torch.Tensor) -> torch.Tensor
(batch, k) or (batch, k, m) -> same shape.
SparseSectorHamiltonian.to_dense¶
def to_dense(self) -> torch.Tensor
lanczos¶
def lanczos(matvec: Callable, v0: torch.Tensor, m: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
m-step Lanczos with full re-orthogonalisation, batched. v0: (batch, k), normalised inside.
Returns alpha (batch, m), beta (batch, m-1) (real) and the Krylov basis Q (batch, m, k).
Breakdown (an invariant subspace) is handled by zero beta and zero further vectors.
ground_state¶
def ground_state(H, m: int=80, tol: Optional[float]=None, max_restarts: int=20, v0: Optional[torch.Tensor]=None, generator: Optional[torch.Generator]=None)
Lowest eigenvalue and eigenvector of every Hamiltonian in the batch by restarted Lanczos.
H: object with matvec, batch_size, dim, dtype, device (e.g. SparseSectorHamiltonian).
Stops when the residual ||H psi - E psi|| < tol * max(1, |E|) for all batch elements
(default tol: 1e-10 in double, 1e-5 in single precision).
Returns E0 (batch,), psi (batch, k), residual (batch,).
krylov_error_bound¶
def krylov_error_bound(beta_m: torch.Tensor, ritz: torch.Tensor, dt: float, m: int) -> torch.Tensor
Bound (Theorem 9.18(b) of the lecture notes) for || exp(-i dt H) v - ||v|| Q_m exp(-i dt T_m) e_1 || / ||v|| from
the data of one m-step Lanczos run: the residual norm beta_m = ||H q_m - Q_m T_m e_m|| (batch,) and the Ritz
values (batch, m) (eigenvalues of T_m). With w_T the half-width of the Ritz values and x = |dt| w_T / 2 it is
(4 beta_m / w_T) S_m(x) = 2 beta_m |dt| S_m(x) / x, S_m(x) = sum_{j >= m} x^j / j!. It holds in exact arithmetic
and needs no estimate of the spectrum of H; it is invariant under H -> H + c.
expm_multiply_step¶
def expm_multiply_step(H, psi: torch.Tensor, dt: float, m: int=30) -> Tuple[torch.Tensor, torch.Tensor]
exp(-i dt H) psi in an m-dimensional Krylov space, w = ||psi|| Q_m exp(-i dt T_m) e_1. Returns the new state
and a bound for its error per batch element, ||psi|| (4 beta_m / w_T) S_m(|dt| w_T / 2) (krylov_error_bound;
Theorem 9.18 of the lecture notes): it holds in exact arithmetic, uses only the Lanczos data (beta_m the norm of
the Lanczos residual, w_T the half-width of the Ritz values) and needs no estimate of the spectrum of H.
(Earlier versions returned beta_m |e_m^T exp(-i dt T_m) e_1|, the norm of the defect of the Schroedinger
equation at time dt: an error rate, not an error.) Rounding errors are not included.
evolve¶
def evolve(H, psi0: torch.Tensor, times, m: int=30, dt: Optional[float]=None, tol: float=1e-08, progress: bool=False) -> torch.Tensor
psi(t) = exp(-i t H) psi0 at every time of the increasing grid `times` (starting at times[0]), in Krylov steps.
Default (dt=None): every step from a state psi is the largest h (up to the next grid time) whose error bound
(krylov_error_bound) is <= tol * h / (times[-1] - times[0]) * ||psi0||. As every step is an isometry, the local
errors add up (Corollary 9.19 of the lecture notes): the Krylov error of the state at time t is at most
tol * ||psi0|| * (t - times[0]) / (times[-1] - times[0]), in exact arithmetic (rounding errors, which grow with the
number of steps, are not included). The step depends on H only through the Lanczos data, so a constant shift
H -> H + c does not change it. With a fixed dt the steps are dt (shortened to fit the grid), and a CostWarning
is issued if the sum of the error bounds exceeds tol * ||psi0||. Returns (T, batch, k), complex.