diff-soup: differentiable, sparse-oriented unified PDE suite

diff-soup is a software package written in JAX for modeling, simulation, and inference of (stiff) dynamical systems governed by ODEs, DAEs, and PDEs.

For modeling:

• a Recipe API that states equations in math-like notation on structured and unstructured domains;
• finite difference, finite volume and finite element discretizations;
• sparse Jacobians assembled by automatic differentiation with graph coloring.

For simulation:

• three variable-step, variable-order stiff integrators written in JAX, JNDF, JBDF and JIDA, for ODEs and DAEs up to index one;
• sparse direct solves with NVIDIA cuDSS on the GPU, and KLU or dense solves on the CPU;
• forward and reverse mode differentiation of every solve, to second order;
• batched solves, with data and model parallelism across GPUs and cluster nodes.

For inference:

• parameter estimation using least squares estimation, maximum likelihood estimation (MLE), and maximum a posteriori (MAP) estimation;
• physically-constrained deep learning using differentiable simulation, with neural networks embedded inside the physical model.

The modeling, simulation and inference modules of diff-soup, and the JAX and GPU features they are built on.

Installation

pip install "diff-soup[cu13]"   # CUDA 13, driver 580 or newer
pip install "diff-soup[cu12]"   # CUDA 12, driver 525 or newer
pip install diff-soup           # CPU only

Requires Linux x86_64 and Python 3.12 or 3.13. The GPU backend needs an NVIDIA GPU of compute capability 8.0 or newer.

Turn on JAX's 64-bit mode with jax.config.update("jax_enable_x64", True) before building a problem.

Building from source is covered on the installation page.

Documentation

Available at https://stmorgenstern.github.io/diff-soup/.

Quick example

The two-dimensional Brusselator, a stiff reaction-diffusion system, on a 32 by 32 finite-volume grid with its walls held at the reaction steady state:

import diff_soup as ds
import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

n = 32
A, B, alpha = 1.0, 3.0, 1.0 / 50.0

spec = ds.problem("brusselator")
spec.domain("square", ds.domains.cartesian_grid_2d(n, n, 1.0, 1.0, method="fvm"))
spec.field("u", mass="differential",
           init=lambda x, y: A + jnp.sin(2 * jnp.pi * x) * jnp.sin(2 * jnp.pi * y))
spec.field("v", mass="differential", init=lambda x, y: jnp.full_like(x, B / A))

@spec.equations(writes=("u", "v"), reads=("u", "v"))
def brusselator(q):
    uuv = q.u**2 * q.v
    du = A + uuv - (B + 1.0) * q.u + alpha * q.ops.laplacian(q.u)
    dv = B * q.u - uuv + alpha * q.ops.laplacian(q.v)
    return {"u": du, "v": dv}

i, j = jnp.indices((n, n))
wall = ((i == 0) | (i == n - 1) | (j == 0) | (j == n - 1)).ravel()
spec.bc.dirichlet("u", on=wall, value=A)
spec.bc.dirichlet("v", on=wall, value=B / A)

problem = spec.build()
y0 = spec.assemble_state({})

sol = ds.pdesolve(problem, ds.JIDA(), (0.0, 10.0), y0)
u_end = sol.ys[-1][: n * n]
print(float(u_end.mean()))   # the average of u over the square at t = 10

u and v of the same model on a 128 by 128 grid, from t = 0 to t = 10.

Acknowledgments

diff-soup builds on the following projects.

Time stepping
SUNDIALS: IDA and CVODE, whose step-size, error-weight and solver-reuse rules the integrators follow (Hindmarsh et al. 2005).
OrdinaryDiffEq.jl: the QNDF method the NDF integrator follows (Rackauckas and Nie 2017).
SciPy: BDF dense output and the bounded least-squares trust region.

Sparse direct solvers
SuiteSparse KLU: the CPU sparse LU (Davis and Palamadai Natarajan 2010).
SuperLU_MT: the threaded CPU sparse LU for source builds (Li 2005).
NVIDIA cuDSS: the GPU sparse direct solver.

JAX ecosystem
JAX: the array library and compiler diff-soup is written in.
Equinox: modules and PyTree tools.
Diffrax: the interfaces for saving output and events, which diff-soup mirrors.
jaxtyping: shape and dtype annotations.
Optax: first-order optimisers for the embedded neural networks.
BlackJAX: posterior sampling.

diff-soup is released under the BSD 3-Clause License, see LICENSE. Third-party licences and notices are listed in THIRD_PARTY_NOTICES.md.

Citation

If you use diff-soup, please cite the paper, arXiv:2610.37333.

@article{degnanmorgenstern2026diffsoup,
    title={diff-soup: {GPU}-native Differentiable Simulation and Inference for Large-scale Stiff Dynamical Systems},
    author={Samuel Degnan-Morgenstern and Martin Z. Bazant},
    journal={arXiv preprint arXiv:2610.37333},
    year={2026},
}