jinns

Physics Informed Neural Networks with JAX. jinns is a Python package for physics-informed neural networks (PINNs) in the JAX ecosystem. It provides an intuitive and flexible interface for

• forward problem: learning a PDE solution.
• inverse problem: learning the parameters of a PDE. See jinns.solve_alternate() for fine-grained and efficient inverse problems.
• meta-modeling: learning a family of PDE indexed by its parameters.
• New in jinns v1.9.0: natural gradient optimizers which greatly enhances the training of PINNs

jinns specific points:

• jinns uses JAX - It is directed to JAX users: forward and backward autodiff, vmapping, jitting and more! No reinventing the wheel: it relies on the JAX ecosystem whenever possible, such as equinox for neural networks or optax for optimization.

• jinns is highly modular - It gives users maximum control for defining their problems, and extending the package. The maths and computations are visible and not hidden behind layers of code!

• jinns is efficient - It compares favorably to other existing Python package for PINNs on the PINNacle benchmarks, as demonstrated in the table below. For more details on the benchmarks, checkout the PINN multi-library benchmark

• Implemented PINN architectures

  • Vanilla Multi-Layer Perceptron popular accross the PINNs litterature.

  • Separable PINNs: allows to leverage forward-mode autodiff for computational speed.

  • Hyper PINNs: useful for meta-modeling

• Other

  • Adaptative Loss Weights are now implemented. Some SoftAdapt, LRAnnealing and ReLoBRaLo are available and users can implement their own strategy. See the tutorial

• Get started: check out our various notebooks on the documentation.

	jinns	DeepXDE - JAX	DeepXDE - Pytorch	PINA	Nvidia Modulus
Burgers1D	445	723	671	1977	646
NS2d-C	265	278	441	1600	275
PInv	149	218	CC	1509	135
Diffusion-Reaction-Inv	284	NI	3424	4061	2541
Navier-Stokes-Inv	175	NI	1511	1403	498

Training time in seconds on an Nvidia T600 GPU. NI means problem cannot be implemented in the backend, CC means the code crashed.

Installation

Install the latest version with pip

pip install jinns

Documentation

The project's documentation is hosted on Gitlab page and available at https://mia_jinns.gitlab.io/jinns/index.html.

Found a bug / want a feature ?

Open an issue on the Gitlab repo.

Contributing

Here are the contributors guidelines:

1. First fork the library on Gitlab.

2. Then clone and install the library in development mode with

pip install -e .

3. Install pre-commit and run it. Our pre-commit hooks consist in ruff format and ruff check. You can install ruff simply by pip install ruff. We highly recommend you to check the code type hints with pyright even though we currently have no rule concerning type checking in the pipeline.

pip install pre-commit
pre-commit install

4. Open a merge request once you are done with your changes, the review will be done via Gitlab.

Contributors

Don't hesitate to contribute and get your name on the list here !

List of contributors: Hugo Gangloff, Nicolas Jouvin, Lucia Clarotto, Inass Soukarieh, Mohamed Badi

Cite us

Please consider citing our work if you found it useful to yours, using this ArXiV preprint

@article{gangloff_jouvin2024jinns,
  title={jinns: a JAX Library for Physics-Informed Neural Networks},
  author={Gangloff, Hugo and Jouvin, Nicolas},
  journal={arXiv preprint arXiv:2412.14132},
  year={2024}
}

Acknowledgement

The library benefited from discussions with other researchers and devs. Special thanks to the Scimba library team, go check out their work if you are interested in sciML and hybrid numerical methods !