Spectral / Galerkin / PINNs experimentation toolkit built on top of JAX for fast differentiable ODE/PDE prototyping, variational forms, and mixed spectral bases.

The coverage badge will update once Codecov is fully configured.

• Orthogonal polynomial and Fourier bases (Chebyshev, Legendre, Jacobi, etc.)
• Tensor product and direct sum spaces with boundary conditions
• Assembly of bilinear / linear forms with symbolic (SymPy) coefficients
• A SymPy-based form-language for describing PDEs
• Curvilinear coordinates
• JAX-backed forward/backward transforms and differentiation
• Utilities for sparse conversion, preconditioning, and projection
• A friendly interface for experimenting with PINNs

Using uv (recommended):

pip install uv  # if not already installed
uv add jaxfun   # when published

From source:

git clone https://github.com/spectralDNS/jaxfun.git
cd jaxfun
uv sync

from jaxfun.galerkin import Chebyshev, TensorProduct, TestFunction, TrialFunction, Div, Grad
from jaxfun.galerkin.inner import inner

C = Chebyshev.Chebyshev(16)
T = TensorProduct((C, C))
v = TestFunction(T)
u = TrialFunction(T)
A = inner(Div(Grad(u)) * v)

See the examples for more patterns.

Run tests (excluding slow):

uv run pytest

Run full (including slow demos):

uv run pytest -m "slow or not slow"

Lint & format:

uv run pre-commit run --all-files

See CONTRIBUTING and the Code of Conduct.

BSD 2-Clause – see LICENSE.

• Mikael Mortensen: mikaem@math.uio.no
• August Femtehjell: august.femtehjell@uio.no