This paper presents ‘PINNs-JAX’, an innovative implementation that utilizes the JAX framework to leverage the distinct capabilities of XLA compilers. This approach aims to improve computational efficiency and flexibility within PINN applications.
Each subplot corresponds to a problem, with its iteration count displayed at the top. The logarithmic x-axis shows the speed-up factor w.r.t the original code in TensorFlow v1, and the y-axis illustrates the mean relative error.
For more information, please refer to our paper:
Comparing PINNs Across Frameworks: JAX, TensorFlow, and PyTorch. Reza Akbarian Bafghi, and Maziar Raissi. AI4DiffEqtnsInSci, ICLR, 2024.
PINNs-JAX requires following dependencies to be installed:
Then, you can install PINNs-JAX itself via [pip]:
pip install pinnsjaxIf you intend to introduce new functionalities or make code modifications, we suggest duplicating the repository and setting up a local installation:
git clone https://github.com/rezaakb/pinns-jax
cd pinns-jax
# [OPTIONAL] create conda environment
conda create -n myenv python=3.9
conda activate myenv
# install package
pip install -e .Explore a variety of implemented examples within the examples folder. To run a specific code, such as the one for the Allen Cahn PDE, you can use:
python examples/ac/train.pyYou can train the model using a specified configuration, like the one found in examples/ac/configs/config.yaml. Parameters can be overridden directly from the command line. For instance:
python examples/ac/train.py trainer.max_epochs=20To utilize our package, there are two primary options:
- Implement your training structures using Hydra, as illustrated in our provided examples.
- Directly incorporate our package to solve your custom problem.
For a practical guide on directly using our package to solve the Burgers PDE in a continuous forward problem, refer to our tutorial here: tutorials/0-Burgers.ipynb.
The data located on the server and will be downloaded automatically upon running each example.
As this is the first version of our package, there might be scope for enhancements and bug fixes. We highly value community contributions. If you find any issues, missing features, or unusual behavior during your usage of this library, please feel free to open an issue or submit a pull request on GitHub. For any queries, suggestions, or feedback, please send them to Reza Akbarian Bafghi at reza.akbarianbafghi@colorado.edu.
Distributed under the terms of the [MIT] license, "pinnsjax" is free and open source software.
We employed this template to develop the package, drawing from its structure and design principles. For a deeper understanding, we recommend visiting their GitHub repository.
@inproceedings{
bafghi2024comparing,
title={Comparing {PINN}s Across Frameworks: {JAX}, TensorFlow, and PyTorch},
author={Reza Akbarian Bafghi and Maziar Raissi},
booktitle={ICLR 2024 Workshop on AI4DifferentialEquations In Science},
year={2024},
url={https://openreview.net/forum?id=BPFzolSSrI}
}