Enable docstring-tests and write some missing docstrings (#216)

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

matfree/test_util.py

@@ -31,31 +31,41 @@ 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)
     return diag + offdiag1 + offdiag2
 
 
 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)))

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
 

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
 
 

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