underPINN: A Modular Physics-Informed Neural Network Framework Based on JAX

underPINN is an open-source Physics-Informed Neural Network (PINN) framework developed by the Aeroscience Computations Analysis Lab, built on JAX. It supports domain decomposition, attention mechanisms, and performance-oriented training, aiming to provide a scalable solution for solving partial differential equations (PDEs). This article will introduce it from aspects such as background, technical features, and application scenarios.

In the field of engineering science, PDEs are core tools for describing physical phenomena. However, traditional numerical methods (such as FEM, FDM) face challenges of high computational cost in complex geometries, high-dimensional problems, or inverse problems. Physics-Informed Neural Networks (PINNs) combine deep learning with physical laws, embedding PDE residuals into the loss function, allowing training without large amounts of labeled data, thus providing new ideas for solving PDEs.

underPINN adopts a modular architecture, decoupling PINN components (networks, losses, optimizers, etc.) to support flexible expansion. It achieves high-performance computing based on JAX's automatic differentiation, JIT compilation, and vmap features; supports domain decomposition technology to handle large-scale problems; integrates attention mechanisms to capture key regions of solutions; and has built-in training strategies such as adaptive loss weights and learning rate scheduling to alleviate convergence difficulties.

1. Aerospace: Solve Euler equations/Navier-Stokes equations, simulate airflow fields, and handle complex geometries without grids; 2. Inverse problems: Infer material parameters, boundary conditions, etc., from observation data; 3. Data-driven simulation: Integrate physical priors with sparse/noisy data to achieve small-sample learning, suitable for fields such as nuclear engineering and astrophysics.

The network architecture supports MLP, ResNet, etc., with configurable hyperparameters; the loss function includes PDE residuals, initial/boundary condition losses, and optional data fitting losses, supporting adaptive weights; the training process includes adaptive sampling, forward propagation (calculating high-order derivatives), loss evaluation, JAX automatic differentiation backpropagation, and parameter updates.

underPINN seamlessly integrates JAX ecosystem tools (Optax, Flax, Orbax, etc.) and provides clear code structures and examples. Future evolution directions: mixed-precision training acceleration, integration of neural operators (such as FNO), multi-physics field coupling, and uncertainty quantification (Bayesian neural networks).

underPINN combines the flexibility of deep learning frameworks with the rigor of scientific computing, providing a powerful tool for solving PDEs. It is of great value to researchers and engineers in fields such as aerospace, energy, and materials, and is worth attention and trial.