add type hints and document trace classes

README.md

@@ -11,10 +11,12 @@
 
 ## About
 
-A majority of algorithms in randomized numerical linear algebra are sequential and additive in nature. However, many implementations do not effectively exploit this structure. Often, once an insufficiently accurate result is observed, the computation is often restarted from the beginning with a modified parameter set. This results in highly inefficient workflows.
+A majority of algorithms in randomized numerical linear algebra are sequential and additive in nature. However, many implementations do not effectively exploit this structure. Often, once an insufficiently accurate result is observed, the computation is restarted from the beginning with a modified parameter set. This results in highly inefficient workflows.
 
 The goal of roughly is to collect the most widespread algorithms of randomized numerical linear algebra and wrap them into an easy to use package where previous computations are stored in memory and available for the user to be reused.
 
+This project was based on the course [Advanced Scientific Programming in Python](https://github.com/JochenHinz/python_seminar) by Jochen Hinz.
+
 ## Example
 
 Computing a Krylov decomposition using the Arnoldi method:
@@ -52,34 +54,4 @@ import roughly as rly
 
 ## Features
 
-- Most implementations in roughly also work for linear operator only available as function handles instead of matrices
-- ...
-
-## Algorithms
-
-Currently, the following algorithms are implemented:
-
-### Krylov methods
-
-- Arnoldi method
-- Block Anroldi method
-- Lanczos method
-- Block Lanczos method
-
-### Low-rank approximations
-
-- Randomized SVD
-- Nyström approximation
-
-### Sketching
-
-- Randomized range finder
-- Subsampled randomized Hadamard transform (in progress)
-
-### Trace estimation
-
-- Hutchinson trace estimator
-- Subspace projection
-- Hutch++ trace estimator
-- Krylov-aware trace estimator (in progress)
-- XTrace (in progress)
+Most implementations in roughly also work for linear operator only available as function handles instead of matrices. Currently, roughly implements the Arnoldi, Lanczos, and blocked versions of them; the randomized SVD and Nyström approximation; the randomized range sketch; and the Girard-Hutchinson, subspace projection, and Hutch++ algorithms.

roughly/approximate/krylov.py

@@ -8,6 +8,7 @@
 import numpy as np
 import scipy as sp
 from abc import ABCMeta, abstractmethod
+from typing import Union
 
 class KrylovDecomposition(metaclass=ABCMeta):
     def __init__(self):
@@ -64,7 +65,7 @@ def _extend(self, k):
         """
         pass
 
-    def compute(self, A, X, k=10, dtype=None):
+    def compute(self, A : Union[np.ndarray, function], X : np.ndarray, k : int = 10, dtype : bool = None):
         """
         Compute Krylov decomposition of a linear operator
 
@@ -101,7 +102,7 @@ def compute(self, A, X, k=10, dtype=None):
 
         return self._result()
 
-    def refine(self, k=1):
+    def refine(self, k : int = 1):
         """
         Refine the existing Krylov decomposition of a linear operator,
 
@@ -168,7 +169,7 @@ class ArnoldiDecomposition(KrylovDecomposition):
         solution of the matrix eigenvalue problem". Quarterly of Applied
         Mathematics. 9 (1): 17–29. doi:10.1090/qam/42792.
     """
-    def __init__(self, reorth_tol=0.7):
+    def __init__(self, reorth_tol : float = 0.7):
         self.reorth_tol = reorth_tol
         super().__init__()
 
@@ -254,7 +255,7 @@ class LanczosDecomposition(KrylovDecomposition):
         Journal of Research of the National Bureau of Standards. 45 (4): 255–282.
         doi:10.6028/jres.045.026.
     """
-    def __init__(self, reorth_tol=0.7, return_matrix=True, extend_matrix=True):
+    def __init__(self, reorth_tol : float = 0.7, return_matrix : bool = True, extend_matrix : bool = True):
         self.return_matrix = return_matrix
         self.extend_matrix = extend_matrix
         self.reorth_tol = reorth_tol
@@ -424,7 +425,7 @@ class BlockLanczosDecomposition(LanczosDecomposition):
         Estimation". SIAM Journal on Matrix Analysis and ApplicationsVol.
         44 (3). doi:10.1137/22M1494257.
     """
-    def __init__(self, return_matrix=True, extend_matrix=True, reorth_steps=-1):
+    def __init__(self, return_matrix : bool = True, extend_matrix : bool = True, reorth_steps : int = -1):
         self.return_matrix = return_matrix
         self.extend_matrix = extend_matrix
         self.reorth_steps = reorth_steps

roughly/approximate/lowrank.py

@@ -7,10 +7,10 @@
 
 import numpy as np
 from abc import ABCMeta, abstractmethod
-from roughly.approximate.sketch import StandardSketch
+from roughly.approximate.sketch import RangeSketch, StandardSketch
 
 class LowRankApproximator(metaclass=ABCMeta):
-    def __init__(self, sketch=StandardSketch()):
+    def __init__(self, sketch : RangeSketch = StandardSketch()):
         self.sketch = sketch
 
     def _preprocess(self, A, k):
@@ -24,7 +24,7 @@ def _preprocess(self, A, k):
     def _factorize(self):
         pass
 
-    def compute(self, A, k=10):
+    def compute(self, A : np.ndarray, k : int = 10):
         """
         Compute the low-rank factorization.
 
@@ -49,7 +49,7 @@ def compute(self, A, k=10):
         U, S, Vh = self._factorize()
         return U, S, Vh
 
-    def refine(self, k=1):
+    def refine(self, k : int = 1):
         """
         Refine the low-rank factorization to a rank higher by k.
 

roughly/approximate/sketch.py

@@ -7,11 +7,12 @@
 
 import numpy as np
 from abc import ABCMeta, abstractmethod
+from typing import Union
 
 from roughly.core.random import gaussian, rademacher, spherical
 
 class RangeSketch(metaclass=ABCMeta):
-    def __init__(self, rng="gaussian", orthogonal=True):
+    def __init__(self, rng : Union[str, function] = "gaussian", orthogonal : bool = True):
         self.orthogonal = orthogonal
         if isinstance(rng, str):
             self.rng = eval(rng)
@@ -20,7 +21,7 @@ def __init__(self, rng="gaussian", orthogonal=True):
         else:
             raise ValueError("'{}' is not a valid random number generation.".format(rng))
 
-    def preprocess(self, A, k, n=None, dtype=None):
+    def _preprocess(self, A, k, n=None, dtype=None):
         self.k = k
         self.n = A.shape[0] if n is None else n
 
@@ -51,10 +52,10 @@ class StandardSketch(RangeSketch):
     orthogonal : bool
         Whether to orthogonalize the range sketch or not.
     """
-    def __init__(self, rng="gaussian", orthogonal=True):
+    def __init__(self, rng : Union[str, function] = "gaussian", orthogonal : bool = True):
         super().__init__(rng, orthogonal)
 
-    def compute(self, A, k=10, n=None, dtype=None, return_embedding=False):
+    def compute(self, A : Union[np.ndarray, function], k : int = 10, n : Union[int, None] = None, dtype  : Union[type, None] = None, return_embedding : bool = False):
         """
         Compute a randomized range sketch for the linear operator A.
 
@@ -81,7 +82,7 @@ def compute(self, A, k=10, n=None, dtype=None, return_embedding=False):
         self.Q : np.ndarray of shape (n, k)
             Approximated range of the linear operator A.
         """
-        self.preprocess(A, k, n, dtype)
+        self._preprocess(A, k, n, dtype)
 
         # Generate embedding and sketch linear operator
         self.S = self.rng(self.n, k)
@@ -94,7 +95,7 @@ def compute(self, A, k=10, n=None, dtype=None, return_embedding=False):
             return self.Q, self.S
         return self.Q
 
-    def refine(self, k=10, return_embedding=False):
+    def refine(self, k : int = 10, return_embedding : bool = False):
         """
         Increase the size of a randomized range sketch for a linear operator.