fix: improve docs and logic in cubature example

Completes docstrings that were missed in a previous commit, simplifies
some logic to aid in code interpretability, and corrects some typehints.

Refs: #798

examples/cubature_recombination.py

@@ -16,9 +16,9 @@
 End-to-end example for using recombination solvers.
 
 In this example we demonstrate how the recombination solvers in Coreax can be used to
-find non-product Cubature formulae, given a product Cubature formulae over the n-cube.
+find non-product Cubature formulae, given product Cubature formulae over the `n`-cube.
 These formulae can be used to exactly (*up to finite-precision arithmetic*) integrate
-all multi-variate polynomials of, at most, some given degree `m`, over the `n-cube`.
+all `n`-variate polynomials of, at most, some given degree `m`, over the `n`-cube.
 """
 
 import functools as ft
@@ -36,17 +36,17 @@
 from coreax.solvers import TreeRecombination
 
 
-# First we define some functions for generating a baseline product Cubature and
-# the multi-variate monomial test-functions which this Cubature should integrates
-# exactly of the n-cube.
+# First we define some functions for generating a baseline product Cubature and the
+# `n`-variate monomial test-functions which the product Cubature and the recombined
+# cubature (when it is computed later on) should integrate exactly over the `n`-cube.
 def leggauss_product_cubature(dimension: int, degree: int) -> tuple[Array, Array]:
     """
-    Construct a Legendre-Gauss product cubature over the n-cube.
+    Construct a Legendre-Gauss product cubature over the `n`-cube.
 
-    :param dimension: Dimension `n` of the n-cube (product space/domain)
-    :param degree: The algebraic degree of all multi-variate polynomials that the
-        produce cubature exactly integrates, over the n-cube (product space/domain)
-    :return: nodes and weights for the product cubature formula.
+    :param dimension: Dimension `n` of the `n`-cube (product space/domain)
+    :param degree: The algebraic degree `m` of all multi-variate polynomials that the
+        product cubature exactly integrates, over the `n`-cube (product space/domain)
+    :return: Nodes and weights for the product cubature formula.
     """
     nodes, weights = np.polynomial.legendre.leggauss(degree)
     prod_nodes = np.fromiter(
@@ -64,17 +64,17 @@ def monomial_power_generator(
     dimension: int, max_degree: int, *, mode: Literal["all", "even", "odd"] = "all"
 ) -> Iterable[tuple[int, ...]]:
     """
-    Return a generator for all combinations of multi-variate monomial powers.
+    Return a generator for all combinations of `n`-variate monomial powers.
 
-    :param dimension: Number of unique variates; dimension of the domain over which
+    :param dimension: Number of unique variates `n`; dimension of the domain over which
         the monomials are defined
-    :param max_degree: Maximal degree of any monomial power; equal to the sum of
+    :param max_degree: Maximal degree of any monomial power `m`; equal to the sum of
         powers for each variate in a given monomial E.G :math:`xy**2` has degree three
         :math:`xy` has degree two, :math:`x` has degree one, etc...
-    :param mode: If to return 'all', 'even', or 'odd' multi-variate monomial powers; a
+    :param mode: If to return 'all', 'even', or 'odd' `n`-variate monomial powers; a
         monomial is even if and only if its powers sum to a non-zero even value. E.G.
         :math:`xy` is even, :math:`xy**2` is odd, and :math:`1 == x^0 y^0` is odd.
-    :return: generator for all multi-variate monomial powers of the specified dimension
+    :return: Generator for all multi-variate monomial powers of the specified dimension
         and maximum degree.
     """
     monomial_degree_generator = itertools.product(
@@ -103,39 +103,81 @@ def test_functions_from_monomial_powers(
     monomial_powers: Int[Array, "k n"],
 ) -> Callable[[Float[Array, " n"]], Float[Array, " k"]]:
     """
-    Construct test functions given a set of multi-variate monomial powers.
+    Construct test functions given a set of `n`-variate monomial powers.
 
-    :param monomial_powers:
-    :return:
+    :param monomial_powers: The integer powers for `k` monomials in `n` variates
+    :return: A callable that evaluates a passed `n`-vector for the `k` monomials
+        defined by `monomial_powers`; can be used as the `test_functions` for a
+        :class:`~coreax.solvers.RecombinationSolver`.
     """
     _, n = monomial_powers.shape
-    coefficients = jnp.zeros((n, n))
+    maximum_degree = jnp.max(jnp.sum(monomial_powers, axis=-1))
+    coefficients = jnp.zeros((maximum_degree + 1, n))
     column_index = jnp.arange(n)
 
     @jax.vmap
-    def reversed_coefficients(_power: Int[Array, " n"]) -> Int[Array, "n n"]:
-        """Create a polyval coefficient matrix given a multi-variate monomial power."""
-        return coefficients.at[_power, column_index].set(1)
+    def polyval_coefficients(_power: Int[Array, "k n"]) -> Int[Array, "k m+1 n"]:
+        """
+        Create a polyval coefficient matrix given an `n`-variate monomial power.
+
+        :param _power: Powers for `k` monomials in `n` variates
+        :return: `k` polyval coefficient arrays where, for a given array, the [i,j]
+            index gives the `j`-th variate raised to the `m+1-i`-th power, where `m+1`
+            is the maximal degree of any of the `k` monomials plus one (the plus one is
+            required to account for the zero-degree coefficients required by polyval).
+        """
+        return coefficients.at[maximum_degree - _power, column_index].set(1)
 
     @ft.partial(jax.vmap, in_axes=(0, None))
     def monomial_test_function(
-        _coefficients: Int[Array, "n n"], x: Float[Array, " n"]
-    ) -> Float[Array, " "]:
-        """Evaluate a multi-variate monomial."""
-        return jnp.prod(jnp.polyval(_coefficients[::-1], x))
+        _coefficients: Int[Array, "k m+1 n"], x: Float[Array, " n"]
+    ) -> Float[Array, " k"]:
+        """
+        Evaluate an `n`-variate monomial.
 
-    return jtu.Partial(monomial_test_function, reversed_coefficients(monomial_powers))
+        :param _coefficients: `k` polyval coefficient arrays of maximum degree `m`; the
+            plus one in the type annotation is required to account for the zero-degree
+            coefficients required by polyval
+        :param x: An `n`-variate array at which to evaluate the monomial
+        :return: The `k` monomials, defined by `_coefficients`, evaluated at `x`.
+        """
+        return jnp.prod(jnp.polyval(_coefficients, x))
+
+    return jtu.Partial(monomial_test_function, polyval_coefficients(monomial_powers))
 
 
 def main(dimension: int = 3, max_degree: int = 4) -> Coresubset:
-    """
+    r"""
     Run the 'recombination' example for finding non-product cubature formulae.
 
-    Generates the
+    Generates the non-product (recombined) cubature formulae over the `n`-cube, for all
+    `n`-variate polynomials of maximal degree `m`, given a product Legendre-Gauss
+    cubature formulae over the same domain with the same degree `m`.
+
+    .. note::
+        Cubature formulae are simply quadrature formulae for `n`-dimensional domains.
+        However, it should be noted that many authors do not make this distinction and
+        simply refer to all formulae as ``quadratures``.
+
+    An `N`-point cubature formulae of degree `m`, with respect to some positive Borel
+    measure :math:`\mu` on :math:`X \subseteq \mathbb{R}^n`, is any discrete measure
+    :math:`\mathbb{Q}^N := \sum_{i=1}^N \lambda_i \delta_{x_i}`, with positive weights
+    :math:`\lambda \in \mathbb{R}_+` and nodes :math:`x_i \in \text{supp}(X)`, for which
+    the following equation
+
+    .. math::
+        \int_X f(x) d\mu(x) = \int_X f(x) d\mathbb{Q}^N(x):=\sum_{i=1}^N\lambda_i f(x_i)
+
+    holds for all `n`-variate polynomials of degree `m`, :math:`f \in \mathcal{P}_m(X)`.
+
+    For the purposes of computational efficiency, the number of points `N` in any given
+    cubature formula should be minimal. Hence, recombination is used in this example to
+    find a non-product formula with :math:`N_\text{recombination} \le N` points, where
+    `N` is the number of points in the Legendre-Gauss product formula over the `n`-cube.
 
-    :param dimension: Number of unique variates; dimension of the domain over which
+    :param dimension: Number of unique variates `n`; dimension of the domain over which
         the monomials are defined
-    :param max_degree: Maximal degree of any monomial power; equal to the sum of
+    :param max_degree: Maximal degree of any monomial power `m`; equal to the sum of
         powers for each variate in a given monomial E.G :math:`xy**2` has degree three
         :math:`xy` has degree two, :math:`x` has degree one, etc...
     :return: The coresubset (non-product) cubature.