[ADD] Squared Frobenius norm estimator (#80)

* [ADD] Squared Frobenius norm estimator

* [FIX] Typo

* [FIX] Math typo

Author: f-dangel

curvlinops/__init__.py

@@ -12,6 +12,7 @@
 )
 from curvlinops.jacobian import JacobianLinearOperator, TransposedJacobianLinearOperator
 from curvlinops.kfac import KFACLinearOperator
+from curvlinops.norm.hutchinson import HutchinsonSquaredFrobeniusNormEstimator
 from curvlinops.papyan2020traces.spectrum import (
     LanczosApproximateLogSpectrumCached,
     LanczosApproximateSpectrumCached,
@@ -23,22 +24,30 @@
 from curvlinops.trace.meyer2020hutch import HutchPPTraceEstimator
 
 __all__ = [
+    # linear operators
     "HessianLinearOperator",
     "GGNLinearOperator",
     "EFLinearOperator",
     "FisherMCLinearOperator",
     "KFACLinearOperator",
     "JacobianLinearOperator",
     "TransposedJacobianLinearOperator",
+    # inversion
     "CGInverseLinearOperator",
     "NeumannInverseLinearOperator",
     "KFACInverseLinearOperator",
+    # slicing
     "SubmatrixLinearOperator",
+    # spectral properties
     "lanczos_approximate_spectrum",
     "lanczos_approximate_log_spectrum",
     "LanczosApproximateSpectrumCached",
     "LanczosApproximateLogSpectrumCached",
+    # trace estimation
     "HutchinsonTraceEstimator",
     "HutchPPTraceEstimator",
+    # diagonal estimation
     "HutchinsonDiagonalEstimator",
+    # norm estimation
+    "HutchinsonSquaredFrobeniusNormEstimator",
 ]

curvlinops/norm/__init__.py

@@ -0,0 +1 @@
+"""Matrix norm estimation methods."""

curvlinops/norm/hutchinson.py

@@ -0,0 +1,62 @@
+"""Hutchinson-style matrix norm estimation."""
+
+from scipy.sparse.linalg import LinearOperator
+
+from curvlinops.trace.hutchinson import HutchinsonTraceEstimator
+
+
+class HutchinsonSquaredFrobeniusNormEstimator:
+    r"""Estimate the squared Frobenius norm of a matrix using Hutchinson's method.
+
+    Let :math:`\mathbf{A} \in \mathbb{R}^{M \times N}` be some matrix. It's Frobenius
+    norm :math:`\lVert\mathbf{A}\rVert_\text{F}` is defined via:
+
+    .. math::
+        \lVert\mathbf{A}\rVert_\text{F}^2
+        =
+        \sum_{m=1}^M \sum_{n=1}^N \mathbf{A}_{n,m}^2
+        =
+        \text{Tr}(\mathbf{A}^\top \mathbf{A}).
+
+    Due to the last equality, we can use Hutchinson-style trace estimation to estimate
+    the squared Frobenius norm.
+
+    Example:
+        >>> from numpy import mean, round
+        >>> from numpy.linalg import norm
+        >>> from numpy.random import rand, seed
+        >>> seed(0) # make deterministic
+        >>> A = rand(5, 5)
+        >>> fro2_A = norm(A, ord='fro')**2 # exact squared Frobenius norm as reference
+        >>> estimator = HutchinsonSquaredFrobeniusNormEstimator(A)
+        >>> # one- and multi-sample approximations
+        >>> fro2_A_low_prec = estimator.sample()
+        >>> fro2_A_high_prec = mean([estimator.sample() for _ in range(1_000)])
+        >>> assert abs(fro2_A - fro2_A_low_prec) > abs(fro2_A - fro2_A_high_prec)
+        >>> round(fro2_A, 4), round(fro2_A_low_prec, 4), round(fro2_A_high_prec, 4)
+        (10.7192, 8.3257, 10.6406)
+    """
+
+    def __init__(self, A: LinearOperator):
+        """Store the linear operator whose squared Frobenius norm will be estimated.
+
+        Args:
+            A: Linear operator whose squared Frobenius norm will be estimated.
+        """
+        self._trace_estimator = HutchinsonTraceEstimator(A.T @ A)
+
+    def sample(self, distribution: str = "rademacher") -> float:
+        """Draw a sample from the squared Frobenius norm estimator.
+
+        Multiple samples can be combined into a more accurate squared Frobenius norm
+        estimation via averaging.
+
+        Args:
+            distribution: Distribution of the vector along which the linear operator
+                will be evaluated. Either ``'rademacher'`` or ``'normal'``.
+                Default is ``'rademacher'``.
+
+        Returns:
+            Sample from the squared Frobenius norm estimator.
+        """
+        return self._trace_estimator.sample(distribution=distribution)

docs/rtd/linops.rst

@@ -78,6 +78,12 @@ Diagonal approximation
 .. autoclass:: curvlinops.HutchinsonDiagonalEstimator
    :members: __init__, sample
 
+Frobenius norm approximation
+============================
+
+.. autoclass:: curvlinops.HutchinsonSquaredFrobeniusNormEstimator
+   :members: __init__, sample
+
 Experimental
 ============