Merge pull request #6 from pyatsysh/docs/auto-differentiable-jax

Docs: auto-differentiable (JAX) namespace — API reference + tutorial

Author: psesh

source-docs/index.txt

@@ -41,5 +41,6 @@ Code documentation for :code:`equadratures` can be found by clicking on the link
 
    Core classes <core/index>
    Secondary classes <secondary/index>
+   Auto-differentiable (JAX) <jax/index>
    Theory <theory/index>
    Tutorials <tutorials/index>

source-docs/jax/basis.txt

@@ -0,0 +1,5 @@
+Basis and design matrix
+==================
+
+.. automodule:: equadratures.jax.basis
+    :members:

source-docs/jax/index.txt

@@ -0,0 +1,39 @@
+:orphan:
+
+.. _jax:
+
+============================
+Auto-differentiable (JAX)
+============================
+
+The :code:`equadratures.jax` namespace is a JAX-native, differentiable backend for
+``equadratures``. Quadrature, orthogonal-polynomial approximation, uncertainty
+quantification and learnable polynomial kernels are all ``jax.grad``-able,
+``jax.jit``-able and ``jax.vmap``-able, so models and kernels can be trained by
+gradient descent. It is a *parallel namespace*: the classic NumPy API is
+unchanged.
+
+Install the optional dependency with ``pip install equadratures[jax]`` (add
+``equadratures[jax-learn]`` for the Optax-based kernel training). See the
+:doc:`tutorial <../tutorials/tutorials/Auto_Differentiable_Equadratures>` for a
+worked introduction.
+
+Code documentation for the individual modules is linked below.
+
+
+.. toctree::
+   :maxdepth: 1
+
+   parameter
+   recurrence
+   quadrature
+   polynomials
+   basis
+   poly
+   kernels
+
+
+.. toctree:
+   :hidden:
+
+   ../documentation

source-docs/jax/kernels.txt

@@ -0,0 +1,5 @@
+Learnable polynomial kernels
+==================
+
+.. automodule:: equadratures.jax.kernel
+    :members:

source-docs/jax/parameter.txt

@@ -0,0 +1,5 @@
+Parameter
+==================
+
+.. automodule:: equadratures.jax.parameter
+    :members:

source-docs/jax/poly.txt

@@ -0,0 +1,5 @@
+Polynomial model
+==================
+
+.. automodule:: equadratures.jax.poly
+    :members:

source-docs/jax/polynomials.txt

@@ -0,0 +1,5 @@
+Orthonormal polynomials
+==================
+
+.. automodule:: equadratures.jax.polynomial
+    :members:

source-docs/jax/quadrature.txt

@@ -0,0 +1,5 @@
+Gauss quadrature
+==================
+
+.. automodule:: equadratures.jax.quadrature
+    :members:

source-docs/jax/recurrence.txt

@@ -0,0 +1,5 @@
+Recurrence coefficients
+==================
+
+.. automodule:: equadratures.jax.recurrence
+    :members:

source-docs/requirements.txt

@@ -10,3 +10,7 @@ sphinx-gallery
 sphinxcontrib-applehelp==1.0.2
 sphinxcontrib-bibtex
 sphinxcontrib-devhelp==1.0.2
+
+# for the equadratures.jax (auto-differentiable) docs + tutorial
+jax
+optax