Incorporated SVD (#102)

Incorporated SVD functionality into CoLA.

Author: AndPotap

cola/backends/jax_fns.py

@@ -10,7 +10,7 @@
 from jax.lax import dynamic_slice, expand_dims
 from jax.lax import fori_loop as _for_loop
 from jax.lax import while_loop as _while_loop
-from jax.lax.linalg import cholesky, qr, svd
+from jax.lax.linalg import cholesky, qr
 from jax.random import PRNGKey, normal
 from jax.scipy.linalg import block_diag
 from jax.scipy.linalg import lu as lu_lax
@@ -58,7 +58,6 @@
 nan_to_num = jnp.nan_to_num
 dynamic_slice = dynamic_slice
 zeros_like = jnp.zeros_like
-svd = svd
 cholesky = cholesky
 solvetri = solvetri
 qr = qr
@@ -94,6 +93,11 @@
 finfo = jnp.finfo
 
 
+def svd(A, full_matrices):
+    U, S, VH = jnp.linalg.svd(A, full_matrices=full_matrices)
+    return U, S, VH.T.conj()
+
+
 def to_np(array):
     return jax.device_get(array)
 

cola/backends/np_fns.py

@@ -65,7 +65,6 @@ def __init__(self):
 sqrt = np.sqrt
 stack = np.stack
 sum = np.sum
-svd = np.linalg.svd
 where = np.where
 fft = np.fft.fft
 ifft = np.fft.ifft
@@ -86,6 +85,11 @@ def lstsq(A, b):
     return soln
 
 
+def svd(A, full_matrices):
+    U, S, VH = np.linalg.svd(A, full_matrices=full_matrices)
+    return U, S, VH.T.conj()
+
+
 def eig(array):
     eigvals, eigvecs = np.linalg.eig(array)
     eigvals, eigvecs = eigvals.astype(complex), eigvecs.astype(complex)

cola/backends/torch_fns.py

@@ -43,7 +43,6 @@
 eigh = torch.linalg.eigh
 solve = torch.linalg.solve
 copy = torch.clone
-svd = torch.linalg.svd
 diag = torch.diag
 zeros_like = torch.zeros_like
 cholesky = torch.linalg.cholesky
@@ -68,6 +67,11 @@
 iscomplexobj = torch.is_complex
 
 
+def svd(A, full_matrices):
+    U, S, VH = torch.linalg.svd(A, full_matrices=full_matrices)
+    return U, S, VH.H
+
+
 def to_np(array):
     return array.detach().cpu().numpy()
 

cola/linalg/decompositions/decompositions.py

@@ -1,16 +1,26 @@
 """ Decompositions of linear operators, e.g. LU, Cholesky, Lanczos, Arnoldi, SVD"""
-from plum import dispatch
 from dataclasses import dataclass
-from cola.ops import LinearOperator, Array
-from cola.ops import Triangular, Permutation, Diagonal
-from cola.ops import Identity, ScalarMul, Kronecker, BlockDiag
-from cola.utils import export
-from cola.linalg.algorithm_base import Algorithm
-from cola.linalg.decompositions.lanczos import lanczos
-from cola.linalg.decompositions.arnoldi import arnoldi
-from cola.linalg.algorithm_base import Auto
+from typing import Any, Optional
+
+from plum import dispatch
+
 import cola.linalg
-from typing import Optional, Any
+from cola.annotations import Unitary
+from cola.linalg.algorithm_base import Algorithm, Auto
+from cola.linalg.decompositions.arnoldi import arnoldi
+from cola.linalg.decompositions.lanczos import lanczos, lanczos_eigs
+from cola.ops import (
+    Array,
+    BlockDiag,
+    Diagonal,
+    Identity,
+    Kronecker,
+    LinearOperator,
+    Permutation,
+    ScalarMul,
+    Triangular,
+)
+from cola.utils import export
 
 PRNGKey = Any
 
@@ -78,6 +88,31 @@ def __call__(self, A: LinearOperator):
         return arnoldi(A, **self.__dict__)
 
 
+@dataclass
+class LanczosSVD(Algorithm):
+    """
+    Does the SVD using the Lanczos decomposition,
+    Args:
+        start_vector (Array, optional): (n,) or (n, b) vector to start the algorithm.
+        max_iters (int, optional): The maximum number of iterations to run.
+        tol (float, optional): Relative error tolerance.
+        pbar (bool, optional): Whether to show progress bar.
+        key (PRNGKey, optional): Random key to use for the algorithm.
+            PRNGKey for jax, long integer for numpy or pytorch.
+    """
+    start_vector: Array = None
+    max_iters: int = 1_000
+    tol: float = 1e-6
+    pbar: bool = False
+    key: Optional[PRNGKey] = None
+
+    def __call__(self, A: LinearOperator):
+        Sigma0_2, V, _ = lanczos_eigs(A.H @ A, **self.__dict__)
+        Sigma1_2, U, _ = lanczos_eigs(A @ A.H, **self.__dict__)
+        Sigma = A.xnp.sqrt(Sigma0_2)
+        return Unitary(U), Diagonal(Sigma), Unitary(V)
+
+
 @export
 @dataclass
 class Lanczos(Algorithm):
@@ -174,3 +209,16 @@ def plu(A: BlockDiag):
     P, L, U = zip(*[plu(Ai) for Ai in A.Ms])
     BD = lambda *args: BlockDiag(*args, multiplicities=A.multiplicities)  # noqa
     return BD(*P), BD(*L), BD(*U)
+
+
+def get_slice(num, which):
+    if num == -1:
+        raise ValueError(f"Number of eigenvalues {num} must be explicitly specified")
+    if which == "SM":
+        eig_slice = slice(0, num, None)
+    elif which == "LM":
+        id = -1 if num is None else -num
+        eig_slice = slice(id, None, None)
+    else:
+        raise NotImplementedError(f"which={which} is not implemented")
+    return eig_slice

cola/linalg/eig/eigs.py

@@ -1,27 +1,24 @@
 import numpy as np
 from plum import dispatch
-from cola.annotations import Unitary, Stiefel
+
+from cola.annotations import SelfAdjoint, Stiefel, Unitary
 from cola.fns import lazify
-from cola.annotations import SelfAdjoint
-from cola.linalg.trace.diag_trace import diag
-from cola.ops.operator_base import LinearOperator
-from cola.ops.operators import Diagonal
-from cola.ops.operators import I_like
-from cola.ops.operators import Identity
-from cola.ops.operators import Triangular
-from cola.linalg.decompositions.lanczos import lanczos_eigs
-from cola.linalg.decompositions.arnoldi import arnoldi_eigs
 from cola.linalg.algorithm_base import Algorithm, Auto
-from cola.linalg.decompositions.decompositions import Arnoldi, Lanczos
+from cola.linalg.decompositions.arnoldi import arnoldi_eigs
+from cola.linalg.decompositions.decompositions import Arnoldi, Lanczos, get_slice
+from cola.linalg.decompositions.lanczos import lanczos_eigs
 from cola.linalg.eig.lobpcg import LOBPCG, lobpcg
 from cola.linalg.eig.power_iteration import PowerIteration
+from cola.linalg.trace.diag_trace import diag
 from cola.linalg.unary.unary import Eig, Eigh
+from cola.ops.operator_base import LinearOperator
+from cola.ops.operators import Diagonal, I_like, Identity, Triangular
 from cola.utils import export
 
 
 @export
 @dispatch.abstract
-def eig(A: LinearOperator, k: int = -1, which: str = 'LM', alg: Algorithm = Auto()):
+def eig(A: LinearOperator, k: int, which: str = "LM", alg: Algorithm = Auto()):
     """
     Computes eigenvalues and eigenvectors of a linear operator.
 
@@ -183,16 +180,3 @@ def eig(A: Diagonal, k: int, which: str, alg: Algorithm):
     eig_vals = A.diag[sorted_ind]
     eig_vecs = I_like(A).to_dense()[:, sorted_ind]
     return eig_vals[eig_slice], Unitary(lazify(eig_vecs[:, eig_slice]))
-
-
-def get_slice(num, which):
-    if num == -1:
-        raise ValueError(f"Number of eigenvalues {num} must be explicitly specified")
-    if which == "SM":
-        eig_slice = slice(0, num, None)
-    elif which == "LM":
-        id = -1 if num is None else -num
-        eig_slice = slice(id, None, None)
-    else:
-        raise NotImplementedError(f"which={which} is not implemented")
-    return eig_slice

cola/linalg/svd/svd.py

@@ -0,0 +1,95 @@
+import numpy as np
+from plum import dispatch
+
+from cola.annotations import Unitary
+from cola.fns import lazify
+from cola.linalg.algorithm_base import Algorithm, Auto
+from cola.linalg.decompositions.decompositions import Lanczos, get_slice
+from cola.linalg.decompositions.lanczos import lanczos_eigs
+from cola.linalg.eig.lobpcg import LOBPCG, lobpcg
+from cola.linalg.inverse.inv import inv
+from cola.ops.operator_base import LinearOperator
+from cola.ops.operators import Dense, Diagonal, I_like, Identity
+from cola.utils import export
+
+
+@export
+class DenseSVD(Algorithm):
+    """
+    Performs SVD on A.
+    """
+
+
+@export
+@dispatch.abstract
+def svd(A: LinearOperator, k: int, which: str = "LM", alg: Algorithm = Auto()):
+    """
+    Computes the SVD of the linear operator A.
+
+    Args:
+        A (LinearOperator): The linear operator to decompose.
+        alg (Algorithm): (Auto, SVD, LanczosSVD)
+
+    Returns:
+        Tuple[LinearOperator]: A tuple U, D, V, such that U D V^{*} = A.
+    """
+
+
+@dispatch(precedence=-1)
+def svd(A: LinearOperator, k: int, which: str, alg: Auto):
+    """ Auto:
+        - if A is small, use dense SVD
+        - if A is large, use Lanczos
+    """
+    match bool(np.prod(A.shape) <= 1e6):
+        case True:
+            alg = DenseSVD()
+        case False:
+            alg = Lanczos(**alg.__dict__)
+    return svd(A, k, which, alg)
+
+
+@dispatch
+def svd(A: LinearOperator, k: int, which: str, alg: DenseSVD):
+    U, Sigma, V = A.xnp.svd(A.to_dense(), full_matrices=True)
+    idx = A.xnp.argsort(Sigma, axis=-1)
+    return Unitary(Dense(U[:, idx])), Diagonal(Sigma[..., idx]), Unitary(Dense(V[:, idx]))
+
+
+@dispatch
+def svd(A: LinearOperator, k: int, which: str, alg: Lanczos):
+    xnp = A.xnp
+    eig_slice = get_slice(k, which)
+    if A.shape[1] <= A.shape[0]:
+        eig_vals, V, _ = lanczos_eigs(A.H @ A, **alg.__dict__)
+        V = Unitary(V[:, eig_slice])
+        Sigma = Diagonal(xnp.sqrt(eig_vals[eig_slice]))
+        U = Unitary(lazify((A @ V @ inv(Sigma)).to_dense()))
+    else:
+        eig_vals, U, _ = lanczos_eigs(A @ A.H, **alg.__dict__)
+        U = Unitary(U[:, eig_slice])
+        Sigma = Diagonal(xnp.sqrt(eig_vals[eig_slice]))
+        V = Unitary(lazify((inv(Sigma) @ U.H @ A).to_dense().conj().T))
+    return U, Sigma, V
+
+
+@dispatch
+def svd(A: LinearOperator, k: int, which: str, alg: LOBPCG):
+    xnp = A.xnp
+    eig_slice = get_slice(k, which)
+    eig_vals, V = lobpcg(A.H @ A, **alg.__dict__)
+    V = Unitary(V[:, eig_slice])
+    Sigma = Diagonal(xnp.sqrt(eig_vals[eig_slice]))
+    U = Unitary(lazify((A @ V @ inv(Sigma)).to_dense()))
+    return U, Sigma, V
+
+
+@dispatch
+def svd(A: Identity, k: int, which: str, alg: Algorithm):
+    ones = A.xnp.ones(A.shape[0], device=A.device, dtype=A.dtype)
+    return Unitary(I_like(A)), Diagonal(ones), Unitary(I_like(A))
+
+
+@dispatch
+def svd(A: Diagonal, k: int, which: str, alg: Algorithm):
+    return Unitary(I_like(A)), A, Unitary(I_like(A))