Logdet example in README (#98)

Author: pnkraemer

README.md

@@ -26,7 +26,7 @@ Imports:
 ```python
 >>> import jax
 >>> import jax.numpy as jnp
->>> from matfree import hutch, montecarlo
+>>> from matfree import hutch, montecarlo, slq
 
 >>> a = jnp.reshape(jnp.arange(12.), (6, 2))
 >>> key = jax.random.PRNGKey(1)
@@ -85,6 +85,24 @@ Here is how to use it:
 [220. 286.]
 
 
+```
+
+### Determinants
+
+
+Estimate log-determinants as such:
+```python
+>>> a = jnp.reshape(jnp.arange(36.), (6, 6)) / 36
+>>> sample_fun = montecarlo.normal(shape=(6,))
+>>> matvec = lambda x: a.T @ (a @ x) + x
+>>> order = 3
+>>> logdet, _ = slq.trace_of_matfun(jnp.log, matvec, order, key=key, sample_fun=sample_fun)
+>>> print(jnp.round(logdet))
+3.0
+>>> # for comparison:
+>>> print(jnp.round(jnp.linalg.slogdet(a.T @ a + jnp.eye(6))[1]))
+3.0
+
 ```
 
 ## Contributing

matfree/slq.py

@@ -1,7 +1,7 @@
 """Stochastic Lanczos quadrature."""
 
 from matfree import decomp, montecarlo
-from matfree.backend import func, linalg, np, prng
+from matfree.backend import func, linalg, np
 
 
 def trace_of_matfun(
@@ -11,22 +11,16 @@ def trace_of_matfun(
     /,
     *,
     key,
-    num_samples_per_batch,
-    num_batches,
-    tangents_shape,
-    tangents_dtype,
-    sample_fun=prng.normal,
+    sample_fun,
+    num_samples_per_batch=10,
+    num_batches=1,
 ):
     """Compute the trace of the function of a matrix.
 
     For example, logdet(M) = trace(log(M)) ~ trace(U log(D) Ut) = E[v U log(D) Ut vt].
     """
-
-    def sample(k, /):
-        return sample_fun(k, shape=tangents_shape, dtype=tangents_dtype)
-
     quadform = quadratic_form_slq(matfun, Av, order)
-    quadform_mc = montecarlo.montecarlo(quadform, sample_fun=sample)
+    quadform_mc = montecarlo.montecarlo(quadform, sample_fun=sample_fun)
 
     quadform_batch = montecarlo.mean_vmap(quadform_mc, num_samples_per_batch)
     quadform_batch = montecarlo.mean_map(quadform_batch, num_batches)

tests/test_autodiff.py

@@ -1,7 +1,7 @@
 """Tests for (selected) autodiff functionality."""
 
 
-from matfree import slq, test_util
+from matfree import montecarlo, slq, test_util
 from matfree.backend import np, prng, testing
 
 
@@ -32,14 +32,14 @@ def fun(s):
 
 def _logdet(A, order, key):
     n, _ = np.shape(A)
+    fun = montecarlo.normal(shape=(n,))
     received, num_nans = slq.trace_of_matfun(
         np.log,
         lambda v: A @ v,
         order,
         key=key,
         num_samples_per_batch=10,
         num_batches=1,
-        tangents_shape=(n,),
-        tangents_dtype=float,
+        sample_fun=fun,
     )
     return received

tests/test_slq.py

@@ -1,6 +1,6 @@
 """Tests for Lanczos functionality."""
 
-from matfree import slq, test_util
+from matfree import montecarlo, slq, test_util
 from matfree.backend import linalg, np, prng, testing
 
 
@@ -23,15 +23,15 @@ def test_logdet(A, order):
     """Assert that the log-determinant estimation matches the true log-determinant."""
     n, _ = np.shape(A)
     key = prng.PRNGKey(1)
+    fun = montecarlo.normal(shape=(n,))
     received, num_nans = slq.trace_of_matfun(
         np.log,
         lambda v: A @ v,
         order,
         key=key,
         num_samples_per_batch=10,
         num_batches=1,
-        tangents_shape=(n,),
-        tangents_dtype=np.dtype(A),
+        sample_fun=fun,
     )
     expected = linalg.slogdet(A)[1]
     print_if_assert_fails = ("error", np.abs(received - expected), "target:", expected)