underPINN：基于JAX的模块化物理信息神经网络框架

underPINN是由Aeroscience Computations Analysis Lab开发的开源物理信息神经网络（PINN）框架，基于JAX构建，支持域分解、注意力机制和面向性能的训练，旨在为偏微分方程（PDE）求解提供可扩展解决方案。本文将从背景、技术特性、应用场景等方面展开介绍。

在工程科学领域，PDE是描述物理现象的核心工具，但传统数值方法（如FEM、FDM）在复杂几何、高维或逆问题中面临计算成本高的挑战。物理信息神经网络（PINNs）将深度学习与物理定律结合，通过在损失函数中嵌入PDE残差，无需大量标注数据即可训练，为PDE求解提供新思路。

underPINN采用模块化架构，解耦PINN组件（网络、损失、优化器等）以支持灵活扩展。基于JAX的自动微分、JIT编译和vmap特性实现高性能计算；支持域分解技术处理大规模问题；集成注意力机制捕捉解的关键区域；内置自适应损失权重、学习率调度等训练策略缓解收敛困难。

1.航空航天：求解欧拉方程/纳维-斯托克斯方程，模拟气流场，无需网格处理复杂几何；2.逆问题：从观测数据反推材料参数、边界条件等；3.数据驱动仿真：融合物理先验与稀疏/噪声数据，实现小样本学习，适用于核工程、天体物理等领域。

网络架构支持MLP、ResNet等，可配置超参数；损失函数包含PDE残差、初始/边界条件损失及可选数据拟合损失，支持自适应权重；训练流程包括自适应采样、前向传播（计算高阶导数）、损失评估、JAX自动微分反向传播及参数更新。

underPINN无缝集成JAX生态工具（Optax、Flax、Orbax等），提供清晰代码结构与示例。未来将演进方向：混合精度训练加速、神经算子集成（如FNO）、多物理场耦合、不确定性量化（贝叶斯神经网络）。

underPINN结合深度学习框架灵活性与科学计算严谨性，为PDE求解提供强大工具，对航空航天、能源、材料等领域的研究人员与工程师具有重要价值，值得关注与尝试。