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.2: Second order optimizers which greatly enhances the training of PINNs. We recommend testing them as default optimizers.
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:
-
First fork the library on Gitlab.
-
Then clone and install the library in development mode with
pip install -e .
- Install pre-commit and run it. Our pre-commit hooks consist in
ruff formatandruff check. You can installruffsimply bypip install ruff. We highly recommend you to check the code type hints withpyrighteven though we currently have no rule concerning type checking in the pipeline.
pip install pre-commit
pre-commit install
- 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 !
2 maintainers
Metadata
2 maintainers
Release files for jinns 1.10.2
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| jinns-1.10.2.tar.gz | 21.7 MB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| jinns-1.10.2-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 21.8 MB
Release files / jinns-1.10.2.tar.gz
| Download URL | jinns-1.10.2.tar.gz |
|---|---|
| Size | 21.7 MB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
7f08045c293d51be92be9dbd4b5b8b6728d67021fef7a930d076091f0b6a5e0b
|
|
BLAKE2b-256 checksum How to use checksums |
c5d44132a8352ce7a701e33d2b2e58cc61cd8d9d1cf7e4ab63c0006e69863682
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.9.25
|
Release files / jinns-1.10.2-py3-none-any.whl
| Download URL | jinns-1.10.2-py3-none-any.whl |
|---|---|
| Size | 164.1 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
aaf7bface07a13c1aaac3f27addb6fffa88cfd02144cf1958f4c129530f47822
|
|
BLAKE2b-256 checksum How to use checksums |
fa5e55c72237644d0d1edd0d06359d11815c9c7f1eb8b5bf2b0a24fd04fb9e87
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.9.25
|
