jinns: a JAX Library for Physics-Informed Neural Networks

TL;DR

jinns is an open-source JAX-based Python library designed for efficiently implementing physics-informed neural networks to solve forward, inverse, and meta-model learning problems in scientific computing.

Contribution

It introduces a versatile, extendable framework built on JAX, integrating popular libraries for model definition and optimization, with comprehensive examples and tutorials.

Findings

Provides baseline models for various physics problems

Enables efficient prototyping and extensions in scientific modeling

Leverages JAX ecosystem for performance and flexibility

Abstract

jinns is an open-source Python library for physics-informed neural networks, built to tackle both forward and inverse problems, as well as meta-model learning. Rooted in the JAX ecosystem, it provides a versatile framework for efficiently prototyping real-problems, while easily allowing extensions to specific needs. Furthermore, the implementation leverages existing popular JAX libraries such as equinox and optax for model definition and optimisation, bringing a sense of familiarity to the user. Many models are available as baselines, and the documentation provides reference implementations of different use-cases along with step-by-step tutorials for extensions to specific needs. The code is available on Gitlab https://gitlab.com/mia_jinns/jinns.