Algorithm implementation and refactor (#73)

Co-authored-by: AndPotap <[email protected]>

README.md

@@ -61,7 +61,7 @@ print(F @ v)
 ```
 ```
 [0.2       0.2       2.2       2.2       4.2       4.2       6.2
- 6.2       8.2       8.2       7.8121004 2.062    ]
+ 6.2       8.2       8.2       7.8       2.1    ]
 ```
 
 2. **Performing Linear Algebra**. With these objects we can perform linear algebra operations even when they are very big.
@@ -70,8 +70,8 @@ print(cola.linalg.trace(F))
 Q = F.T @ F + 1e-3 * cola.ops.I_like(F)
 b = cola.linalg.inv(Q) @ v
 print(jnp.linalg.norm(Q @ b - v))
-print(cola.linalg.eig(F)[0][:5])
-print(cola.sqrt(A))
+print(cola.linalg.eig(F, k=F.shape[0])[0][:5])
+print(cola.linalg.sqrt(A))
 ```
 
 ```

cola/__init__.py

@@ -12,7 +12,6 @@
 import_from_all("annotations", globals(), __all__, __name__)
 import_from_all("linalg", globals(), __all__, __name__)
 import_from_all("utils", globals(), __all__, __name__)
-import_from_all("decompositions", globals(), __all__, __name__)
 
 __all__.append("LinearOperator")
 # import_from_all("ops", globals(), __all__,__name__)

cola/backends/jax_fns.py

@@ -123,11 +123,12 @@ def is_cuda_available():
 
 def eig(A):
     # if GPU, convert to CPU first since jax doesn't support it
-    device = A.device_buffer.device()
-    if str(device)[:3] != 'cpu':
-        A = jax.device_put(A, jax.devices("cpu")[0])
+    # device = A.device_buffer.device()
+    # if str(device)[:3] != 'cpu':
+    # A = jax.device_put(A, jax.devices("cpu")[0])
     w, v = jnp.linalg.eig(A)
-    return jax.device_put(w, device), jax.device_put(v, device)
+    return w, v
+    # return jax.device_put(w, device), jax.device_put(v, device)
 
 
 def eye(n, m, dtype, device):

cola/linalg/algorithm_base.py

@@ -0,0 +1,31 @@
+from plum import parametric
+from cola.ops import LinearOperator
+from cola.utils import export
+from types import SimpleNamespace
+# import pytreeclass as tc
+
+
+@export
+class Algorithm:
+    pass
+
+
+@parametric
+class IterativeOperatorWInfo(LinearOperator):
+    def __init__(self, A: LinearOperator, alg: Algorithm):
+        super().__init__(A.dtype, A.shape)
+        self.A = A
+        self.alg = alg
+        self.info = {}
+
+    def _matmat(self, X):
+        Y, self.info = self.alg(self.A, X)
+        return Y
+
+    def __str__(self):
+        return f"{self.alg}({str(self.A)})"
+
+
+@export
+class Auto(SimpleNamespace, Algorithm):
+    pass

cola/algorithms/__init__.py → cola/linalg/decompositions/__init__.py

@@ -1,4 +1,3 @@
-""" Low-level algorithms used in CoLA (no dispatch rules). """
 import pkgutil
 from cola.utils import import_from_all
 

cola/algorithms/arnoldi.py → cola/linalg/decompositions/arnoldi.py

@@ -2,9 +2,8 @@
 from cola.ops import LinearOperator
 from cola.ops import Array, Dense
 from cola.ops import Householder, Product
-from cola.utils import export
-import cola
-from cola import Stiefel, lazify
+# from cola.utils import export
+from cola import lazify
 
 # def arnoldi_eigs_bwd(res, grads, unflatten, *args, **kwargs):
 #     val_grads, eig_grads, _ = grads
@@ -34,9 +33,8 @@
 
 # @export
 # @iterative_autograd(arnoldi_eigs_bwd)
-@export
 def arnoldi_eigs(A: LinearOperator, start_vector: Array = None, max_iters: int = 100, tol: float = 1e-7,
-                 use_householder: bool = False, pbar: bool = False):
+                 use_householder: bool = False, pbar: bool = False, key=None):
     """
     Computes eigenvalues and eigenvectors using Arnoldi.
 
@@ -48,6 +46,7 @@ def arnoldi_eigs(A: LinearOperator, start_vector: Array = None, max_iters: int =
         tol (float, optional): Stopping criteria.
         use_householder (bool, optional): Use Householder Arnoldi variant.
         pbar (bool, optional): Show a progress bar.
+        key (PNRGKey, optional): PRNGKey for random number generation.
 
     Returns:
         tuple:
@@ -56,16 +55,15 @@ def arnoldi_eigs(A: LinearOperator, start_vector: Array = None, max_iters: int =
             - info (dict): General information about the iterative procedure.
     """
     Q, H, info = arnoldi(A=A, start_vector=start_vector, max_iters=max_iters, tol=tol, use_householder=use_householder,
-                         pbar=pbar)
+                         pbar=pbar, key=key)
     xnp = A.xnp
     eigvals, vs = xnp.eig(H.to_dense())
     eigvectors = Q @ lazify(vs)
     return eigvals, eigvectors, info
 
 
-@export
 def arnoldi(A: LinearOperator, start_vector=None, max_iters=100, tol: float = 1e-7, use_householder: bool = False,
-            pbar: bool = False):
+            pbar: bool = False, key=None):
     """
     Computes the Arnoldi decomposition of the linear operator A, A = QHQ^*.
 
@@ -77,6 +75,7 @@ def arnoldi(A: LinearOperator, start_vector=None, max_iters=100, tol: float = 1e
         tol (float, optional): Stopping criteria.
         use_householder (bool, optional): Use Householder Arnoldi iteration.
         pbar (bool, optional): Show a progress bar.
+        key (PNRGKey, optional): PRNGKey for random number generation.
 
     Returns:
         tuple:
@@ -86,7 +85,8 @@ def arnoldi(A: LinearOperator, start_vector=None, max_iters=100, tol: float = 1e
     """
     xnp = A.xnp
     if start_vector is None:
-        start_vector = xnp.randn(A.shape[-1], dtype=A.dtype, device=A.device)
+        key = xnp.PRNGKey(42) if key is None else key
+        start_vector = xnp.randn(A.shape[-1], dtype=A.dtype, device=A.device, key=key)
     if len(start_vector.shape) == 1:
         rhs = start_vector[:, None]
     else:
@@ -103,17 +103,6 @@ def arnoldi(A: LinearOperator, start_vector=None, max_iters=100, tol: float = 1e
         return Q, H, infodict
 
 
-def ArnoldiDecomposition(A: LinearOperator, start_vector=None, max_iters=100, tol=1e-7, use_householder=False,
-                         pbar=False):
-    """ Provides the Arnoldi decomposition of a matrix A = Q H Q^H. LinearOperator form of arnoldi,
-        see arnoldi for arguments."""
-    Q, H, info = arnoldi(A=A, start_vector=start_vector, max_iters=max_iters, tol=tol, use_householder=use_householder,
-                         pbar=pbar)
-    A_approx = cola.UnitaryDecomposition(Q, H)
-    A_approx.info = info
-    return A_approx
-
-
 def get_householder_vec_simple(x, idx, xnp):
     indices = xnp.arange(x.shape[0])
     vec = xnp.where(indices >= idx, x, 0.)