Add (squared) Frobenius norm to hutch.py (#96)

Author: pnkraemer

matfree/hutch.py

@@ -21,6 +21,16 @@ def quadform(vec):
     return stochastic_estimate(quadform, **kwargs)
 
 
+def frobeniusnorm_squared(matvec_fun, /, **kwargs):
+    """Estimate the squared Frobenius norm of a matrix stochastically."""
+
+    def quadform(vec):
+        Av = matvec_fun(vec)
+        return np.vecdot(Av, Av)
+
+    return stochastic_estimate(quadform, **kwargs)
+
+
 def diagonal(matvec_fun, /, **kwargs):
     """Estimate the diagonal of a matrix stochastically."""
 

tests/test_hutch.py

@@ -20,6 +20,32 @@ def fixture_key():
     return prng.PRNGKey(seed=1)
 
 
+@testing.parametrize("num_batches", [1_000])
+@testing.parametrize("num_samples_per_batch", [1_000])
+@testing.parametrize("dim", [1, 10])
+@testing.parametrize("sample_fun", [montecarlo.normal, montecarlo.rademacher])
+def test_frobeniusnorm_squared(
+    fun, key, num_batches, num_samples_per_batch, dim, sample_fun
+):
+    """Assert that the Frobenius norm estimate is accurate."""
+    # Linearise function
+    x0 = prng.uniform(key, shape=(dim,))  # random lin. point
+    _, jvp = func.linearize(fun, x0)
+    J = func.jacfwd(fun)(x0)
+
+    # Estimate the trace
+    fun = sample_fun(shape=np.shape(x0), dtype=np.dtype(x0))
+    estimate = hutch.frobeniusnorm_squared(
+        jvp,
+        num_batches=num_batches,
+        key=key,
+        num_samples_per_batch=num_samples_per_batch,
+        sample_fun=fun,
+    )
+    truth = np.trace(J.T @ J)
+    assert np.allclose(estimate, truth, rtol=1e-2)
+
+
 @testing.parametrize("num_batches", [1_000])
 @testing.parametrize("num_samples_per_batch", [1_000])
 @testing.parametrize("dim", [1, 10])