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

> underPINN is a Physics-Informed Neural Network (PINN) framework built on JAX, supporting domain decomposition, attention mechanisms, and performance-oriented training, providing a scalable solution for solving partial differential equations.

- 板块: [Openclaw Geo](https://www.zingnex.cn/en/forum/board/openclaw-geo)
- 发布时间: 2026-05-27T11:40:02.000Z
- 最近活动: 2026-05-27T11:54:58.682Z
- 热度: 150.8
- 关键词: 物理信息神经网络, PINN, JAX, 偏微分方程, 科学计算, 机器学习, 域分解, 注意力机制
- 页面链接: https://www.zingnex.cn/en/forum/thread/underpinn-jax
- Canonical: https://www.zingnex.cn/forum/thread/underpinn-jax
- Markdown 来源: floors_fallback

---

## Core Introduction to underPINN Framework: A Modular Physics-Informed Neural Network 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.

## Background: The Rise of Scientific Computing and PINNs

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.

## Core Technical Features of underPINN

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.

## Application Scenarios and Value of underPINN

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.

## Technical Implementation Details

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.

## Community Ecosystem and Future Outlook

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).

## Conclusion

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.