Write a docstring for tridiag_sym (#215)

* Display the docstrings in a list because it looks tidier * Correct a docstring that used to show a non-existent argument * Write a docstring for tridiag_sym * Fix the wrong length of a numpy-doc header * Remove a placeholder (that wasn't supposed to be there)
pnkraemer · Sep 9, 2024 · 64d31c0 · 64d31c0
1 parent da7e004
commit 64d31c0
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 5 deletions.
diff --git a/matfree/decomp.py b/matfree/decomp.py
@@ -28,7 +28,7 @@ class _DecompResult(containers.NamedTuple):
 
 
 def tridiag_sym(
-    krylov_depth,
+    krylov_depth: int,
     /,
     *,
     materialize: bool = True,
@@ -63,8 +63,47 @@ def tridiag_sym(
         }
         ```
 
+    Parameters
+    ----------
+    krylov_depth
+        The depth of the Krylov space.
+        Read this as ``number of matrix-vector products''.
+        The deeper the Krylov space, the more accurate the factorisation tends to be.
+        However, the computational complexity increases linearly
+        with the depth of the Krylov space.
+    materialize
+        The value of this flag indicates whether the tridiagonal matrix
+        should be returned in a sparse format (which means, as a tuple of diagonas)
+        or as a dense matrix.
+        The dense matrix is helpful if different decompositions should be used
+        interchangeably. The sparse representation requires less memory.
+    reortho
+        The value of this parameter indicates whether to reorthogonalise the
+        basis vectors during the forward pass.
+        Reorthogonalisation makes the forward pass more expensive, but helps
+        (significantly) with numerical stability.
+    custom_vjp
+        The value of this flag indicates whether to use a custom vector-Jacobian
+        product as proposed by Krämer et al. (2024; bibtex above).
+        Generally, using a custom VJP tends to be a good idea.
+        However, due to JAX's mechanics, a custom VJP precludes the use of forward-mode
+        differentiation
+        ([see here](https://jax.readthedocs.io/en/latest/_autosummary/jax.custom_vjp.html)),
+        so don't use a custom VJP if you need forward-mode differentiation.
+
+    Returns
+    -------
+    decompose
+        A decomposition function that maps
+        ``(matvec, vector, *params)`` to the decomposition.
+        The decomposition is a tuple of (nested) arrays.
+        The first element is the Krylov basis,
+        the second element represents the tridiagonal matrix
+        (how it is represented depends on the value of ``materialize''),
+        the third element is
+        the residual, and the fourth element is
+        the (inverse of the) length of the initial vector.
     """
-
     if reortho == "full":
         return _tridiag_reortho_full(
             krylov_depth, custom_vjp=custom_vjp, materialize=materialize

diff --git a/matfree/eig.py b/matfree/eig.py
@@ -24,8 +24,6 @@ def svd_partial(v0: Array, depth: int, Av: Callable, vA: Callable):
         Matrix-vector product function.
     vA:
         Vector-matrix product function.
-    matrix_shape:
-        Shape of the matrix involved in matrix-vector and vector-matrix products.
     """
     # Factorise the matrix
     algorithm = decomp.bidiag(depth, materialize=True)

diff --git a/mkdocs.yml b/mkdocs.yml
@@ -37,13 +37,13 @@ plugins:
             show_object_full_path: true
             show_category_heading: false
             docstring_style: numpy
+            docstring_section_style: table
             show_if_no_docstring: true
             members_order: alphabetical
             annotations_path: brief
             show_signature: true
             show_signature_annotations: true
             separate_signature: false
-            docstring_section_style: list
             show_source: true
 watch: [matfree]
 extra: