jxpinn

English | 中文

A general-purpose JAX-based Physics-Informed Neural Network (PINN) framework.

jxpinn lets you define a PDE problem, a neural-network architecture, and a training domain through a single Config object, then train with either Optax optimizers or jaxopt.LBFGS. It supports hard (analytical) boundary/initial conditions, inverse problems, and polygonal 2D domains out of the box.

Features

• Built on JAX — JIT compilation, GPU support, and automatic differentiation.
• Unified Config dataclass that wires together domain, problem, network, and optimizer.
• Efficient derivative computation via chained JVPs instead of repeated jax.grad calls.
• Hard boundary/initial condition enforcement via problem-specific constraining_fns.
• Support for inverse problems (trainable problem parameters).
• 2D polygon domains (L-shapes, stars, rings, etc.) using pure-JAX helpers.
• Examples for harmonic oscillators, Burgers equation, wave equation, and Poisson problems.

Installation

Install from PyPI:

python -m pip install jxpinn

If you want to use the LBFGS optimizer, install the optional dependency:

python -m pip install "jxpinn[lbfgs]"

From source

Alternatively, clone the repository and install in editable mode:

git clone https://github.com/LBurny/jxpinn.git
cd jxpinn
python -m pip install -e .

With the LBFGS optimizer:

python -m pip install -e ".[lbfgs]"

Requirements

• Python >= 3.9
• JAX, NumPy, Matplotlib, SciPy, Optax
• jaxopt (optional, required for optimiser="lbfgs")

See pyproject.toml for the full dependency list.

Quick start

import numpy as np

from jxpinn.config import Config
from jxpinn.domains import RectangularDomainND
from jxpinn.networks import FCN
from jxpinn.problems import HarmonicOscillator1D
from jxpinn.trainers import PINNTrainer

config = Config(
    run="harmonic_oscillator",
    domain=RectangularDomainND,
    domain_init_kwargs={"xmin": np.array([0.0]), "xmax": np.array([1.0])},
    problem=HarmonicOscillator1D,
    problem_init_kwargs={"d": 2, "w0": 20},
    network=FCN,
    network_init_kwargs={"layer_sizes": [1, 32, 1]},
    n_steps=5000,
    ns=((60,),),
    n_test=(200,),
    optimiser="adam",
    optimiser_kwargs={"learning_rate": 1e-3},
    show_figures=False,
)

trainer = PINNTrainer(config)
all_params = trainer.train()

Examples

The examples/ directory contains standalone training scripts:

Run an example from the repository root:

python examples/train_poisson_star.py

Output figures and summaries are written to results/plots/ and results/summaries/ respectively (these directories are gitignored).

Running tests

Tests are plain Python scripts that can be run directly:

for f in test/test_*.py; do python "$f"; done

If you prefer pytest, install the dev extras and run:

python -m pip install -e ".[dev]"
pytest

Project structure

jxpinn/
├── jxpinn/           # Main package
│   ├── config.py      # Configuration dataclass
│   ├── domains.py     # Domain definitions (rectangular, polygonal)
│   ├── networks.py    # Neural network architectures
│   ├── problems.py    # PDE problems and constraints
│   ├── trainers.py    # Training loop and derivative machinery
│   └── util/          # Helpers (geometry, JAX utilities, I/O)
├── examples/          # Standalone training examples
├── test/              # Tests
├── research/          # Research notebooks and experiments
├── docs/              # Documentation and planning notes
└── results/           # Generated figures and summaries (gitignored)

License

This project is released under the MIT License. See LICENSE for details.

Contributing

Contributions are welcome! Please see CONTRIBUTING.md and open an issue or pull request on GitHub.

Acknowledgments

This project was built with JAX, Optax, and Matplotlib.