rqutils.ground_locg module

Single-vector LOBPCG

Overview

This module defines a single function ground_locg, which implements a single-vector version of the Locally Optimal Block Preconditioned Conjugate Gradient (LOBPCG) solver[1]. The structure heavily borrows from jax.experimental.sparse.linalg.lobpcg_standard, with optimizations for single-vector (ground eigenpair) calculation.

LOBPCG is a matrix-free method for finding the extremal eigenpairs of a generalized eigenvalue problem

\[A x = \lambda B x\]

with \((A, B)\) Hermitian. Our implementation only solves the non-generalized (\(B = I\)) problem and finds the minimum eigenvalue \(\lambda_0\) and the corresponding eigenvector \(v_0\).

The basic arguments to the function are

  • Matrix \(A\), either as a JAX Array or a function that takes the vector \(x\) (as a JAX Array) as input and returns \(Ax\).

  • Initial vector, which must have a non-vanishing overlap with \(v_0\).

Algorithm

See reference [1] for details. The goal is to minimize the Rayleigh quotient

\[\{\lambda_0, v_0\} = \{\min_{x}, \mathrm{argmin}_{x}\} \left(\rho (x) := \frac{x^{\dagger} A x}{|x|^2} \right) .\]

Conceptually, the algorithm consists of gradient descent iterations

\[x_{i+1} = x_{i} + \alpha_{i} r_{i},\]

where \(r_{i} := A x_{i} - \rho (x_{i}) x_{i}\) can be proven to be proportional to the gradient of \(\rho (x_{i})\). Note that \(r_{i}\) is orthogonal to \(x_{i}\):

\[\begin{split}x_{i}^{\dagger} r_{i} & = x_{i}^{\dagger} A x_{i} - \frac{x_{i}^{\dagger} A x_{i}}{|x_{i}|^2} x_{i}^{\dagger} x_{i} \\ & = 0.\end{split}\]

In practice, instead of finding the optimal step size \(\alpha_{i}\), we can directly minimize \(\rho\) in the space spanned by \(\{x_{i}, r_{i}\}\) via the Rayleigh-Ritz method and identify the minimizing vector as \(x_{i+1}\). Furthermore, it is known that convergence of the algorithm is drastically improved if we search \(x_{i+1}\) in the extended space spanned by \(\{x_{i}, x_{i-1}, r_{i}\}\). Thus, one iteration of gradient descent is given by the following steps, with orthogonal \(\{x_{i}, y_{i}, r_{i}\}\) (\(x_{i}, y_{i}\) normal) as the carryover from the previous iteration and \(R_A\) indicating the Rayleigh-Ritz routine over matrix \(A\):

\[\begin{split}p & \leftarrow \frac{r_{i}}{|r_{i}|} \\ \theta, \kappa & \leftarrow R_A[x_{i}, y_{i}, p] \\ s & \leftarrow \kappa_1 y_{i} + \kappa_2 p \\ t & \leftarrow \frac{\kappa_0}{|s|} s - |s| x_{i} \\ u & \leftarrow \kappa_0 x_{i} + s \\ x_{i+1} & \leftarrow \frac{u}{|u|} \\ y_{i+1} & \leftarrow \frac{t}{|t|} \\ r_{i+1} & \leftarrow A x_{i+1} - \theta x_{i+1}.\end{split}\]

The normal vector \(y_{i}\) is orthogonal to \(x_{i}\) and \(r_{i}\) and lies in the space spanned by \(\{x_{i}, x_{i-1}, r_{i}\}\).

Single-vector optimization

The B of LOBPCG refers to the algorithm’s ability to determine multiple eigenvectors simultaneously as a block. We have however chosen to compute just the ground state vector in this implementation, eyeing running on extremely large vectors (memory requirement of LOBPCG scales with the number of eigenvectors to compute). This choice opens up further memory-footprint optimizations in the Rayleigh-Ritz subroutine.

In the Rayleigh-Ritz subroutine, we form the matrix

\[R_{jk} = w_{j}^{\dagger} A w_{k},\]

where \(w = {x, y, p}\), and diagonalize it. With an undetermined number of simultaneous vectors in \(w\), we’d have to concatenate \(x\), \(y\), and \(p\) (thus creating their copies) and then numerically invert \(R\). Since we know that there are only three vectors, we can construct \(R_{jk}\) “by hand” and analytically invert the 3x3 matrix.

Distributed arrays

This function works transparently over distributed (sharded) input \(v_0\) if the callable passed as the mat argument preserves the sharding in the output.

References

[1]: https://en.wikipedia.org/wiki/LOBPCG

Single-vector LOBPCG API

rqutils.ground_locg.ground_locg(mat: Callable[[Array], Array] | Array, xinit: Array | int, args: tuple = (), maxiter: int = 1000, tol: float | None = None, vspace: tuple[int, DTypeLike] | None = None, log_level: int = 30) → tuple[float, NDArray, int]

Single-vector LOBPCG.

Parameters:

  • mat – Matrix \(A\), either as an Array or a function \(x \mapsto Ax\).

  • xinit – Initial vector. If given as an integer (requires vspace if mat is callable), a one-hot vector is created internally.

  • args – Additional arguments to callable mat.

  • maxiter – Maximum number of gradient descent iterations.

  • tol – Convergence condition.

  • vspace – Specification (dimension, dtype) of the vector space. Required only when mat is a callable and vinit is an integer.

  • log_level – Verbosity level.

Returns:

The smallest eigenvalue, its eigenvector, and the number of gradient descent iterations required to achieve the solution.