From 2406b9f2206db954e18a0c4c76af5acabcb7e9d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicholas=20Kr=C3=A4mer?= <nicokraemer@gmx.net>
Date: Tue, 10 Sep 2024 08:04:22 +0200
Subject: [PATCH] Enable docstring-tests and write some missing docstrings
 (#216)

---
 matfree/funm.py                            |  6 +++---
 matfree/test_util.py                       | 10 ++++++++++
 pyproject.toml                             |  7 +++++++
 tutorials/4_control_variates.py            |  1 +
 tutorials/6_low_memory_trace_estimation.py |  1 +
 5 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/matfree/funm.py b/matfree/funm.py
index 59d59c7..1ab67fe 100644
--- a/matfree/funm.py
+++ b/matfree/funm.py
@@ -1,16 +1,16 @@
-"""Matrix-free implementations of functions of matrices.
+r"""Matrix-free implementations of functions of matrices.
 
 This includes matrix-function-vector products
 
 $$
-(f, A, v, p) \\mapsto f(A(p))v
+(f, A, v, p) \mapsto f(A(p))v
 $$
 
 as well as matrix-function extensions for stochastic trace estimation,
 which provide
 
 $$
-(f, A, v, p) \\mapsto v^\\top f(A(p))v.
+(f, A, v, p) \mapsto v^\top f(A(p))v.
 $$
 
 Plug these integrands into
diff --git a/matfree/test_util.py b/matfree/test_util.py
index f165d0e..dfe9557 100644
--- a/matfree/test_util.py
+++ b/matfree/test_util.py
@@ -31,12 +31,14 @@ def asymmetric_matrix_from_singular_values(vals, /, nrows, ncols):
 
 
 def to_dense_bidiag(d, e, /, offset=1):
+    """Materialize a bidiagonal matrix."""
     diag = linalg.diagonal_matrix(d)
     offdiag = linalg.diagonal_matrix(e, offset=offset)
     return diag + offdiag
 
 
 def to_dense_tridiag_sym(d, e, /):
+    """Materialize a symmetric tridiagonal matrix."""
     diag = linalg.diagonal_matrix(d)
     offdiag1 = linalg.diagonal_matrix(e, offset=1)
     offdiag2 = linalg.diagonal_matrix(e, offset=-1)
@@ -44,18 +46,26 @@ def to_dense_tridiag_sym(d, e, /):
 
 
 def tree_random_like(key, tree, *, generate_func=prng.normal):
+    """Fill a tree with random values."""
     flat, unflatten = tree_util.ravel_pytree(tree)
     flat_like = generate_func(key, shape=flat.shape, dtype=flat.dtype)
     return unflatten(flat_like)
 
 
 def assert_columns_orthonormal(Q, /):
+    """Assert that the columns in a matrix are orthonormal."""
     eye_like = Q.T @ Q
     ref = np.eye(len(eye_like))
     assert_allclose(eye_like, ref)
 
 
 def assert_allclose(a, b, /):
+    """Assert that two arrays are close.
+
+    This function uses a different default tolerance to
+    jax.numpy.allclose. Instead of fixing values, the tolerance
+    depends on the floating-point precision of the input variables.
+    """
     a = np.asarray(a)
     b = np.asarray(b)
     tol = np.sqrt(np.finfo_eps(np.dtype(b)))
diff --git a/pyproject.toml b/pyproject.toml
index b13108c..f2ed29e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -108,6 +108,8 @@ select = [
     "EM",
     # tryceratops:
     "TRY",
+    # Docstrings:
+    "D",
 ]
 ignore = [
     # warning: `one-blank-line-before-class` (D203) and `no-blank-line-before-class` (D211) are incompatible.
@@ -125,6 +127,11 @@ line-ending = "lf"
 quote-style = "double"
 skip-magic-trailing-comma = true
 
+[tool.ruff.per-file-ignores]
+# Ignore all directories named `tests`.
+"tests/**" = ["D"]
+"matfree/backend/**" = ["D"]
+
 [tool.ruff.lint.isort]
 split-on-trailing-comma = false
 
diff --git a/tutorials/4_control_variates.py b/tutorials/4_control_variates.py
index d11b6a0..04487ae 100644
--- a/tutorials/4_control_variates.py
+++ b/tutorials/4_control_variates.py
@@ -34,6 +34,7 @@
 
 
 def matvec_ctrl(v):
+    """Evaluate a matrix-vector product with a control variate."""
     return A @ v - diagonal_ctrl * v
 
 
diff --git a/tutorials/6_low_memory_trace_estimation.py b/tutorials/6_low_memory_trace_estimation.py
index cf3b661..25983a2 100644
--- a/tutorials/6_low_memory_trace_estimation.py
+++ b/tutorials/6_low_memory_trace_estimation.py
@@ -29,6 +29,7 @@
 
 
 def large_matvec(v):
+    """Evaluate a (dummy for a) large matrix-vector product."""
     return 1.2345 * v