PINNx: Physics-Informed Neural Networks for Scientific Machine Learning in JAX

PINNx is a library for scientific machine learning and physics-informed learning in JAX. It is a rewrite of DeepXDE but is enhanced by our brain modeling ecosystem. For example, it leverages

• brainstate for just-in-time compilation,
• saiunit for dimensional analysis,
• braintools for checkpointing, loss functions, and other utilities.

Quickstart

Define a PINN with explicit variables and physical units.

import braintools
import saiunit as u
import pinnx

# geometry
geometry = pinnx.geometry.GeometryXTime(
    geometry=pinnx.geometry.Interval(-1, 1.),
    timedomain=pinnx.geometry.TimeDomain(0, 0.99)
).to_dict_point(x=u.meter, t=u.second)

uy = u.meter / u.second
v = 0.01 / u.math.pi * u.meter ** 2 / u.second

# boundary conditions
bc = pinnx.icbc.DirichletBC(lambda x: {'y': 0. * uy})
ic = pinnx.icbc.IC(lambda x: {'y': -u.math.sin(u.math.pi * x['x'] / u.meter) * uy})

# PDE equation
def pde(x, y):
    jacobian = approximator.jacobian(x)
    hessian = approximator.hessian(x)
    dy_x = jacobian['y']['x']
    dy_t = jacobian['y']['t']
    dy_xx = hessian['y']['x']['x']
    residual = dy_t + y['y'] * dy_x - v * dy_xx
    return residual

# neural network
approximator = pinnx.nn.Model(
    pinnx.nn.DictToArray(x=u.meter, t=u.second),
    pinnx.nn.FNN(
        [geometry.dim] + [20] * 3 + [1],
        "tanh",
        braintools.init.KaimingUniform()
    ),
    pinnx.nn.ArrayToDict(y=uy)
)

# problem
problem = pinnx.problem.TimePDE(
    geometry,
    pde,
    [bc, ic],
    approximator,
    num_domain=2540,
    num_boundary=80,
    num_initial=160,
)

# training
trainer = pinnx.Trainer(problem)
trainer.compile(braintools.optim.Adam(1e-3)).train(iterations=15000)
trainer.compile(braintools.optim.LBFGS(1e-3)).train(2000, display_every=500)
trainer.saveplot(issave=True, isplot=True)

Installation

• Install the stable version with pip:

pip install pinnx --upgrade

• Install pinnx on CPU or GPU with JAX following the instructions on

pip install pinnx[cpu]  # for CPU

pip install pinnx[cuda12]  # for NVIDIA GPUs with CUDA 12

pip install pinnx[cuda13]  # for NVIDIA GPUs with CUDA 13

pip install pinnx[tpu]  # for Google TPUs

Documentation

The official documentation is hosted on Read the Docs: https://pinnx.readthedocs.io/

See also the ecosystem

pinnx is one part of our brain modeling ecosystem: https://brainmodeling.readthedocs.io/