SPDEax

Composable spatial discretisations and operators for solving SPDEs with Diffrax

SPDEax is a library for constructing spatially discretised PDEs and SPDEs which can then be timestepped using Diffrax. This includes:

• a lightweight Field abstraction for storing spatial discretisations over grids
• a composable and extensible operator framework
• periodic, finite-difference differential operators Divergence, Gradient, Laplacian, etc.
• compatibility with jit, grad, vmap, etc.
• inter-operability with Lineax

Installation

pip install spdeax

Quickstart

Solving a one-dimensional heat equation is extremely easy!

import jax
import jax.numpy as jnp
import spdeax as spx
import diffrax as dfx

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

x0 = (0.0,)
x_final = (2 * jnp.pi,)
spatial_shape = (128,)
ic = spx.ScalarField.discretise_fn(
    x0,
    x_final,
    spatial_shape,
    lambda x: jnp.sin(x),
    coordinates="unpacked",
)

lplc = spx.Laplacian()

sol = dfx.diffeqsolve(
    terms=dfx.ODETerm(lambda t, y, args: lplc(y)),
    solver=dfx.Euler(),
    t0=0.0,
    t1=1.0,
    dt0=1e-3,
    y0=ic,
    saveat=dfx.SaveAt(t1=True),
    max_steps=10**4,
)

Further examples

For more examples showcasing the functionality of the library, see examples. These include

• non_linear_heat.ipynb
  • we demonstrate how a discretisation of $\partial_x((1+u^2)\partial_x u)$ is trivial to represent in SPDEax
• stochastic_heat.ipynb
  • we demonstrate how to solve SPDEs using SPDEax
  • in particular, we show how to wrap Lineax operators to output an AbstractField for use in diffrax.ControlTerm
• batching.ipynb
  • we demonstrate how batching is handled in SPDEax
• custom_advection_operator.ipynb
  • we show how to define your own custom AbstractOperator using a one-dimensional advection problem as an example
• two_dimensional_heat.ipynb
  • we demonstrate that SPDEax is dimension-agnostic!
• dean_kawasaki.ipynb
  • we show how an advanced SPDE can be solved with relative ease

Future development ideas

• add support for Dirichlet and Neumann boundary conditions
• add timesteppers for semi-linear (S)PDEs e.g. exponential Euler-type methods
• add documentation
• add TensorField