Metadata-Version: 2.4
Name: jxpinn
Version: 0.1.0
Summary: A general-purpose JAX-based Physics-Informed Neural Network (PINN) framework.
Author-email: LBurny <siwuxiebuaa@gmail.com>
License: MIT
Project-URL: Homepage, https://github.com/LBurny/jxpinn
Project-URL: Repository, https://github.com/LBurny/jxpinn
Project-URL: Issues, https://github.com/LBurny/jxpinn/issues
Keywords: pinn,jax,physics-informed neural networks,deep learning,scientific machine learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4.20
Requires-Dist: jaxlib>=0.4.20
Requires-Dist: numpy>=1.24
Requires-Dist: matplotlib>=3.5
Requires-Dist: scipy>=1.10
Requires-Dist: shapely>=2.0
Requires-Dist: trimesh>=4.0
Requires-Dist: mapbox-earcut
Requires-Dist: tensorboardX>=2.6
Requires-Dist: IPython>=8.0
Requires-Dist: optax>=0.1.7
Provides-Extra: lbfgs
Requires-Dist: jaxopt>=0.8.0; extra == "lbfgs"
Provides-Extra: dev
Requires-Dist: jaxopt>=0.8.0; extra == "dev"
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: ruff>=0.1.0; extra == "dev"
Dynamic: license-file

# jxpinn

English | [中文](README.zh.md)

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](https://github.com/google/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_fn`s.
- 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:

```bash
python -m pip install jxpinn
```

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

```bash
python -m pip install "jxpinn[lbfgs]"
```

### From source

Alternatively, clone the repository and install in editable mode:

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

With the LBFGS optimizer:

```bash
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

```python
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:

| Example | Problem | Notes |
|---------|---------|-------|
| `examples/train_burgers.py` | 1D viscous Burgers equation | Hard IC/BC enforcement |
| `examples/train_poisson_square.py` | 2D Poisson on a square | Rectangular SDF support |
| `examples/train_poisson_star.py` | 2D Poisson on a star polygon | `PolygonDomain2D` |

Run an example from the repository root:

```bash
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:

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

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

```bash
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](LICENSE) for
details.

## Contributing

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

## Acknowledgments

This project was built with [JAX](https://github.com/google/jax),
[Optax](https://github.com/google-deepmind/optax), and
[Matplotlib](https://matplotlib.org/).
