From e26f77fc42d5a46a7a098ad7762f3a50700c157a Mon Sep 17 00:00:00 2001 From: FMatti Date: Mon, 10 Jun 2024 17:47:46 +0200 Subject: [PATCH] add type hints and document trace classes --- README.md | 36 +---- roughly/approximate/krylov.py | 11 +- roughly/approximate/lowrank.py | 8 +- roughly/approximate/sketch.py | 13 +- roughly/approximate/trace.py | 246 +++++++++++++++++++++++++++++++-- 5 files changed, 253 insertions(+), 61 deletions(-) diff --git a/README.md b/README.md index ad8adb8..f6ed43f 100644 --- a/README.md +++ b/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. diff --git a/roughly/approximate/krylov.py b/roughly/approximate/krylov.py index fbdd9bc..a130583 100644 --- a/roughly/approximate/krylov.py +++ b/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 diff --git a/roughly/approximate/lowrank.py b/roughly/approximate/lowrank.py index 850605f..3addaee 100644 --- a/roughly/approximate/lowrank.py +++ b/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. diff --git a/roughly/approximate/sketch.py b/roughly/approximate/sketch.py index f6ef894..f94e319 100644 --- a/roughly/approximate/sketch.py +++ b/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. diff --git a/roughly/approximate/trace.py b/roughly/approximate/trace.py index d0540f5..3e6d2fb 100644 --- a/roughly/approximate/trace.py +++ b/roughly/approximate/trace.py @@ -10,12 +10,13 @@ import numpy as np from abc import ABCMeta, abstractmethod +from typing import Union from roughly.core.random import gaussian, rademacher, spherical from roughly.approximate.sketch import StandardSketch class TraceEstimator(metaclass=ABCMeta): - def __init__(self, rng="gaussian"): + def __init__(self, rng : Union[str, function] = "gaussian"): if isinstance(rng, str): self.rng = eval(rng) elif isinstance(rng, function): @@ -23,7 +24,10 @@ def __init__(self, rng="gaussian"): 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): + """ + Preprocess the input matrix/operator and determine problem dimensions. + """ self.k = k self.n = A.shape[0] if n is None else n @@ -39,52 +43,266 @@ def refine(self): pass class HutchinsonTraceEstimator(TraceEstimator): - def __init__(self, rng="gaussian"): + """ + Implements the Girard-Hutchinson trace estimator for matrix or linear + operator A [1]. + + Parameters + ---------- + rng : function or str + Distribution to generator the randomized embedding with. Either a + function whose arguments are the sizes of the axes of the randomized + embedding or one of the following strings: + - "gaussian": Standard Gaussian random numbers + - "rademacher": Uniformly sample from {-1, +1} + - "spherical": Spherically normalized standard Gaussians + + Attributes + ---------- + .compute(A, X, k, ...) + Compute the Girard-Hutchinson trace estimate. + .refine(k, ...) + Refine the Girard-Hutchinson trace estimate. + + Example + ------- + >>> import numpy as np + >>> from roughly.approximate.trace import HutchinsonTraceEstimator + >>> estimator = HutchinsonTraceEstimator() + >>> A = np.diag(np.ones(100)) + >>> t_A = np.trace(A) + >>> t = estimator.compute(A, 100) + >>> assert(abs(t_A - t) / abs(t_A) < 1) + >>> t = estimator.refine(100) + >>> assert(abs(t_A - t) / abs(t_A) < 1e-1) + + [1] Girard, A. (1989). "A fast ‘Monte-Carlo cross-validation’ procedure for + large least squares problems with noisy data". Numerische Mathematik. + 56 (1): 1-23. doi:10.1007/BF01395775. + """ + def __init__(self, rng : Union[str, function] = "gaussian"): super().__init__(rng) - def compute(self, A, k=10, n=None, dtype=None): - self.preprocess(A, k, n, dtype=dtype) + def compute(self, A : Union[np.ndarray, function], k : int = 10, n : Union[int, None] = None, dtype : Union[type, None] = None): + """ + Compute the trace estimate. + + Parameters + ---------- + A : np.ndarray of shape (n, n) or function + The matrix or linear operator (given as function handle) for which a + basis of the trace is computed. + k : int >= 1 + The number of iterations in the Arnoldi method + n : int >= 1 + The length of each component in the sketch (only needed if A is + given as a function handle). + dtype : [np.float, np.complex, ...] + The expected data type of the Krylov basis. If None, dtype is + inferred from A. + + Returns + ------- + self.est : float + Trace estimate. + """ + self._preprocess(A, k, n, dtype=dtype) V = self.rng(self.n, self.k) self.est = np.trace(V.T @ self.matvec(V)) / self.k return self.est - def refine(self, k=1): + def refine(self, k : int = 1): + """ + Refine the trace estimate. + + Parameters + ---------- + k : int >= 1 + The number of iterations in the Arnoldi method + + Returns + ------- + self.est : float + Trace estimate. + """ V = self.rng(self.n, k) self.est = (self.est * self.k + np.trace(V.T @ self.matvec(V))) / (k + self.k) self.k += k return self.est class SubspaceProjectionEstimator(TraceEstimator): - def __init__(self, rng="gaussian"): + """ + Implements the subspace projection for trace estimation [2]. + + Parameters + ---------- + rng : function or str + Distribution to generator the randomized embedding with. Either a + function whose arguments are the sizes of the axes of the randomized + embedding or one of the following strings: + - "gaussian": Standard Gaussian random numbers + - "rademacher": Uniformly sample from {-1, +1} + - "spherical": Spherically normalized standard Gaussians + + Attributes + ---------- + .compute(A, k, ...) + Compute the trace estimate with subspace target rank k. + .refine(k, ...) + Refine the trace estimate by increasing the subspace target rank by k. + + Example + ------- + >>> import numpy as np + >>> from roughly.approximate.trace import SubspaceProjectionEstimator + >>> estimator = SubspaceProjectionEstimator() + >>> A = np.diag(np.ones(100)) + >>> t_A = np.trace(A) + >>> t = estimator.compute(A, 50) + >>> assert(abs(t_A - t) / abs(t_A) < 1) + >>> t = estimator.refine(50) + >>> assert(abs(t_A - t) / abs(t_A) < 1e-1) + + [2] Meyer, R. A., Musco, C., Musco, C., and Woodruff, D. P. (2021). + "Hutch++: Optimal Stochastic Trace Estimation". Symposium on Simplicity + in Algorithms (SOSA). doi:10.1137/1.9781611976496.16. + """ + def __init__(self, rng : Union[str, function] = "gaussian"): self.sketch = StandardSketch(rng) super().__init__(rng) - def compute(self, A, k=10, n=None, dtype=None): - self.preprocess(A, k, n, dtype=dtype) + def compute(self, A : Union[np.ndarray, function], k : int = 10, n : Union[int, None] = None, dtype : Union[type, None] = None): + """ + Compute the trace estimate. + + Parameters + ---------- + A : np.ndarray of shape (n, n) or function + The matrix or linear operator (given as function handle) for which a + basis of the trace is computed. + k : int >= 1 + The number of iterations in the Arnoldi method + n : int >= 1 + The length of each component in the sketch (only needed if A is + given as a function handle). + dtype : [np.float, np.complex, ...] + The expected data type of the Krylov basis. If None, dtype is + inferred from A. + + Returns + ------- + self.est : float + Trace estimate. + """ + self._preprocess(A, k, n, dtype=dtype) self.Q = self.sketch.compute(A, k=k, n=self.n, dtype=self.dtype) self.est = np.trace(self.Q.T @ self.matvec(self.Q)) return self.est - def refine(self, k=1): + def refine(self, k : int = 1): + """ + Refine the trace estimate. + + Parameters + ---------- + k : int >= 1 + The number of iterations in the Arnoldi method + + Returns + ------- + self.est : float + Trace estimate. + """ self.Q = self.sketch.refine(k=k) self.est = np.trace(self.Q.T @ self.matvec(self.Q)) self.k += k return self.est class DeflatedTraceEstimator(TraceEstimator): - def __init__(self, rng="gaussian"): + """ + Implements the Hutch++ algorithm for trace estimation [2]. + + Parameters + ---------- + rng : function or str + Distribution to generator the randomized embedding with. Either a + function whose arguments are the sizes of the axes of the randomized + embedding or one of the following strings: + - "gaussian": Standard Gaussian random numbers + - "rademacher": Uniformly sample from {-1, +1} + - "spherical": Spherically normalized standard Gaussians + + Attributes + ---------- + .compute(A, k, ...) + Compute the Hutch++ trace estimate with k matrix-vector products. + .refine(k, ...) + Refine the Hutch++ trace estimate by adding k matrix-vector products. + + Example + ------- + >>> import numpy as np + >>> from roughly.approximate.trace import DeflatedTraceEstimator + >>> estimator = DeflatedTraceEstimator() + >>> A = np.diag(np.ones(100)) + >>> t_A = np.trace(A) + >>> t = estimator.compute(A, 50) + >>> assert(abs(t_A - t) / abs(t_A) < 1) + >>> t = estimator.refine(50) + >>> assert(abs(t_A - t) / abs(t_A) < 1e-1) + + [2] Meyer, R. A., Musco, C., Musco, C., and Woodruff, D. P. (2021). + "Hutch++: Optimal Stochastic Trace Estimation". Symposium on Simplicity + in Algorithms (SOSA). doi:10.1137/1.9781611976496.16. + """ + def __init__(self, rng : Union[str, function] = "gaussian"): self.sketch = StandardSketch(rng) super().__init__(rng) - def compute(self, A, k=10, n=None, dtype=None): - self.preprocess(A, k, n, dtype=dtype) + def compute(self, A : Union[np.ndarray, function], k : int = 10, n : Union[int, None] = None, dtype : Union[type, None] = None): + """ + Compute the trace estimate. + + Parameters + ---------- + A : np.ndarray of shape (n, n) or function + The matrix or linear operator (given as function handle) for which a + basis of the trace is computed. + k : int >= 1 + The number of iterations in the Arnoldi method + n : int >= 1 + The length of each component in the sketch (only needed if A is + given as a function handle). + dtype : [np.float, np.complex, ...] + The expected data type of the Krylov basis. If None, dtype is + inferred from A. + + Returns + ------- + self.est : float + Trace estimate. + """ + self._preprocess(A, k, n, dtype=dtype) self.Q = self.sketch.compute(A, k=k // 3, n=self.n, dtype=self.dtype) G = self.rng(self.n, k // 3) G = G - self.Q @ (self.Q.T @ G) self.est = np.trace(self.Q.T @ self.matvec(self.Q)) + 1 / G.shape[-1] * np.trace(G.T @ self.matvec(G)) return self.est - def refine(self, k=1): + def refine(self, k : int = 1): + """ + Refine the trace estimate. + + Parameters + ---------- + k : int >= 1 + The number of iterations in the Arnoldi method + + Returns + ------- + self.est : float + Trace estimate. + """ self.Q = self.sketch.refine(k=k) G = self.rng(self.n, k) G = G - self.Q @ (self.Q.T @ G)