picardax 0.0.2

Time-Implicit Differentiable Finite Difference Solvers in JAX.

Time-Implicit Differentiable PDE Solvers in JAX.

Installation • Quickstart • Design • Features • Documentation • Background

🚨 This is a pre-release for me to cite in my dissertation. It is based on code from the PRDP paper and Emulator Superiority Paper slightly rewritten and enhanced. Please check back end of May 2026 for a full release with extended features, documentation, and examples.

Picardax is a sister project to Exponax (which provides explicit and exponential time steppers). Where Exponax sidesteps nonlinear solves entirely, Picardax embraces them: it solves the implicit system arising from theta-method (and similar) time discretizations using iterative nonlinear solvers (Picard iteration, Newton's method) that themselves rely on iterative linear solvers (CG, GMRES, BiCGStab). Built on JAX and Equinox, every solver is automatically differentiable, JIT-compilable, and GPU/TPU-ready.

Installation

pip install picardax

Requires Python 3.10+ and JAX 0.4.13+. See the JAX install guide.

Quickstart

Use a convenience stepper for common PDEs:

import picardax
from picardax.periodic.stepper import Diffusion

stepper = Diffusion(
    num_spatial_dims=1,
    domain_extent=1.0,
    num_points=64,
    dt=0.01,
    diffusivity=0.01,
)

u_next = stepper(u)

Or assemble components manually for full control:

from picardax import CG, Picard, Stepper, ThetaResiduum
from picardax.periodic import CollocatedDerivatives
from picardax.periodic.rhs import Diffusion

derivatives = CollocatedDerivatives(
    num_spatial_dims=1, num_points=64, domain_extent=1.0
)
rhs = Diffusion(derivatives=derivatives, diffusivity=0.01)
residuum = ThetaResiduum(rhs=rhs, theta=0.5, dt=0.01)
solver = Picard(linear_solver=CG(), maxiter=10, tol=1e-6)
stepper = Stepper(residuum=residuum, solver=solver)

u_next = stepper(u)

# Diagnose convergence
n_iters = stepper.diagnose(u)
errors = stepper.get_error_iterates(u)

Design

TODO

Features

TODO

Documentation

Documentation is available at fkoehler.site/picardax.

Background

TODO

Related

TODO

License

MIT, see here

---

fkoehler.site  ·  GitHub @ceyron  ·  X @felix_m_koehler  ·  LinkedIn Felix Koehler