feat: add dispatch rules for pinv and svd

cola/backends/jax_fns.py

@@ -29,6 +29,7 @@
 stack = jnp.stack
 norm = jnp.linalg.norm
 inv = jnp.linalg.inv
+pinv = jnp.linalg.pinv
 log = jnp.log
 sum = jnp.sum
 abs = jnp.abs

cola/backends/np_fns.py

@@ -39,6 +39,7 @@ def __init__(self):
 int32 = np.int32
 int64 = np.int64
 inv = np.linalg.inv
+pinv = np.linalg.pinv
 isreal = np.isreal
 kron = np.kron
 log = np.log

cola/linalg/inverse/pinv.py

@@ -4,7 +4,16 @@
 from cola.annotations import PSD
 from cola.linalg.algorithm_base import Algorithm, Auto, IterativeOperatorWInfo
 from cola.linalg.inverse.cg import CG
-from cola.ops.operators import Diagonal, I_like, Identity, LinearOperator, Permutation, ScalarMul
+from cola.ops.operators import (
+    BlockDiag,
+    Diagonal,
+    I_like,
+    Identity,
+    Kronecker,
+    LinearOperator,
+    Permutation,
+    ScalarMul,
+)
 from cola.utils import export
 from cola.utils.utils_linalg import get_precision
 
@@ -88,9 +97,25 @@ def pinv(A: ScalarMul, alg: Algorithm):
 
 @dispatch
 def pinv(A: Diagonal, alg: Algorithm):
-    return Diagonal(1. / A.diag)
+    xnp = A.xnp
+    abs_diag = xnp.abs(A.diag)
+    # TODO: usage of tolerance?
+    mask = abs_diag > 1e-12
+    inv_diag = xnp.zeros_like(A.diag)
+    inv_diag = xnp.update_array(inv_diag, 1. / A.diag[mask], mask)
+    return Diagonal(inv_diag)
 
 
 @dispatch
 def pinv(A: Permutation, alg: Algorithm):
     return Permutation(A.xnp.argsort(A.perm), A.dtype)
+
+
+@dispatch
+def pinv(A: BlockDiag, alg: Algorithm):
+    return BlockDiag(*[pinv(M, alg) for M in A.Ms], multiplicities=A.multiplicities)
+
+
+@dispatch
+def pinv(A: Kronecker, alg: Algorithm):
+    return Kronecker(*[pinv(M, alg) for M in A.Ms])

cola/linalg/svd/svd.py

@@ -9,7 +9,16 @@
 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.ops.operators import (
+    BlockDiag,
+    Dense,
+    Diagonal,
+    I_like,
+    Identity,
+    Kronecker,
+    Permutation,
+    ScalarMul,
+)
 from cola.utils import export
 
 
@@ -93,3 +102,38 @@ def svd(A: Identity, k: int, which: str, alg: Algorithm):
 @dispatch
 def svd(A: Diagonal, k: int, which: str, alg: Algorithm):
     return Unitary(I_like(A)), A, Unitary(I_like(A))
+
+
+@dispatch(precedence=1)
+def svd(A: ScalarMul, k: int, which: str, alg: Algorithm):
+    ones = A.xnp.ones((int(A.shape[0]), ), dtype=A.dtype, device=A.device)
+    phase = A.c / A.xnp.abs(A.c)
+    U = Diagonal(phase * ones)
+    Sigma = Diagonal(A.xnp.abs(A.c) * ones)
+    V = Diagonal(ones)
+    return U, Sigma, V
+
+
+@dispatch(precedence=1)
+def svd(A: Permutation, k: int, which: str, alg: Algorithm):
+    ones = A.xnp.ones((int(A.shape[0]), ), dtype=A.dtype, device=A.device)
+    U = A  # Permutation is unitary
+    Sigma = Diagonal(ones)
+    V = Diagonal(ones)
+    return U, Sigma, V
+
+
+@dispatch(precedence=1)
+def svd(A: BlockDiag, k: int, which: str, alg: Algorithm):
+    # U = BlockDiag(Ui), S = Diagonal(concat(si)), V = BlockDiag(Vi)
+    results = [svd(M, k, which, alg) for M in A.Ms]
+    Us, Ss, Vs = zip(*results)
+    return BlockDiag(*Us), Diagonal(A.xnp.concat([S.diag for S in Ss])), BlockDiag(*Vs)
+
+
+@dispatch(precedence=1)
+def svd(A: Kronecker, k: int, which: str, alg: Algorithm):
+    # A = (U1 ⊗ U2)(S1 ⊗ S2)(V1 ⊗ V2)^H
+    results = [svd(M, k, which, alg) for M in A.Ms]
+    Us, Ss, Vs = zip(*results)
+    return Kronecker(*Us), Kronecker(*Ss), Kronecker(*Vs)

tests/linalg/inverse/test_pinv_dispatch.py

@@ -0,0 +1,51 @@
+from cola.ops.operators import BlockDiag, Diagonal, Kronecker
+from cola.linalg.algorithm_base import Auto
+from cola.linalg.inverse.pinv import pinv
+from cola.utils.utils_for_tests import get_xnp, parametrize, relative_error
+from cola.backends import all_backends
+
+
+@parametrize(all_backends)
+def test_pinv_block_diag(backend):
+    xnp = get_xnp(backend)
+    dtype = xnp.float64
+    A = Diagonal(xnp.array([1., 2., 0.], dtype=dtype, device=None))
+    B = Diagonal(xnp.array([3., 4., 5.], dtype=dtype, device=None))
+    BD = BlockDiag(A, B)
+
+    # Dense pinv reference
+    BD_dense = BD.to_dense()
+    if backend == 'numpy':
+        import numpy as np
+        BD_pinv_dense = np.linalg.pinv(BD_dense)
+    else:
+        BD_pinv_dense = xnp.pinv(BD_dense)
+
+    # Dispatch pinv
+    BD_pinv = pinv(BD, alg=Auto())
+
+    rel_error = relative_error(BD_pinv.to_dense(), BD_pinv_dense)
+    assert rel_error < 1e-5
+
+
+@parametrize(all_backends)
+def test_pinv_kronecker(backend):
+    xnp = get_xnp(backend)
+    dtype = xnp.float64
+    A = Diagonal(xnp.array([1., 2.], dtype=dtype, device=None))
+    B = Diagonal(xnp.array([3., 4., 0.], dtype=dtype, device=None))
+    K = Kronecker(A, B)
+
+    # Dense pinv reference
+    K_dense = K.to_dense()
+    if backend == 'numpy':
+        import numpy as np
+        K_pinv_dense = np.linalg.pinv(K_dense)
+    else:
+        K_pinv_dense = xnp.pinv(K_dense)
+
+    # Dispatch pinv
+    K_pinv = pinv(K, alg=Auto())
+
+    rel_error = relative_error(K_pinv.to_dense(), K_pinv_dense)
+    assert rel_error < 1e-5

tests/linalg/svd/test_svd_dispatch.py

@@ -0,0 +1,77 @@
+from cola.ops.operators import BlockDiag, Diagonal, Kronecker, Permutation, ScalarMul
+from cola.linalg.svd.svd import svd
+from cola.utils.utils_for_tests import get_xnp, parametrize, relative_error
+from cola.backends import all_backends
+import numpy as np
+import cola.backends.np_fns as np_fns
+
+if not hasattr(np_fns, 'pinv'):
+    np_fns.pinv = np.linalg.pinv
+
+
+@parametrize(all_backends)
+def test_svd_scalar_mul(backend):
+    xnp = get_xnp(backend)
+    dtype = xnp.float64
+    c = -3.0
+    A = Diagonal(xnp.array([1., 2., 3.], dtype=dtype, device=None))
+    S_mul = ScalarMul(c, shape=A.shape, dtype=dtype, device=None)
+
+    U, S, V = svd(S_mul, k=A.shape[0])
+
+    approx = U @ S @ V.H
+    rel_error = relative_error(approx.to_dense(), S_mul.to_dense())
+    assert rel_error < 1e-5
+
+    # Check singular values
+    s_vals = S.diag
+    expected_s = xnp.array([abs(c)] * 3, dtype=dtype, device=None)
+    assert relative_error(s_vals, expected_s) < 1e-5
+
+
+@parametrize(all_backends)
+def test_svd_permutation(backend):
+    xnp = get_xnp(backend)
+    dtype = xnp.float64
+    perm = xnp.array([1, 0, 2], dtype=xnp.int32, device=None)
+    P = Permutation(perm, dtype=dtype)
+
+    U, S, V = svd(P, k=P.shape[0])
+
+    approx = U @ S @ V.H
+    rel_error = relative_error(approx.to_dense(), P.to_dense())
+    assert rel_error < 1e-5
+
+    s_vals = S.diag
+    expected_s = xnp.ones((3, ), dtype=dtype, device=None)
+    assert relative_error(s_vals, expected_s) < 1e-5
+
+
+@parametrize(all_backends)
+def test_svd_block_diag(backend):
+    xnp = get_xnp(backend)
+    dtype = xnp.float64
+    A = Diagonal(xnp.array([1., 2.], dtype=dtype, device=None))
+    B = Diagonal(xnp.array([3., 4.], dtype=dtype, device=None))
+    BD = BlockDiag(A, B)
+
+    U, S, V = svd(BD, k=BD.shape[0])
+
+    approx = U @ S @ V.H
+    rel_error = relative_error(approx.to_dense(), BD.to_dense())
+    assert rel_error < 1e-5
+
+
+@parametrize(all_backends)
+def test_svd_kronecker(backend):
+    xnp = get_xnp(backend)
+    dtype = xnp.float64
+    A = Diagonal(xnp.array([1., 2.], dtype=dtype, device=None))
+    B = Diagonal(xnp.array([3., 4.], dtype=dtype, device=None))
+    K = Kronecker(A, B)
+
+    U, S, V = svd(K, k=K.shape[0])
+
+    approx = U @ S @ V.H
+    rel_error = relative_error(approx.to_dense(), K.to_dense())
+    assert rel_error < 1e-5