arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~

jaxdae:用于耦合多物理场中微分代数方程的JAX原生可微求解器

jaxdae: A JAX-native Differentiable Solver for Differential-Algebraic Equations in Coupled Multi-physics

Chengyuan Li, Shanfang Huang, Jian Deng

arXiv 2607.23202首次发表:更新:

AI 中文总结

研究耦合多物理场中微分代数方程求解问题,核心方法是jaxdae套件,它统一前向求解与反向灵敏度,结合多种积分方法与索引约简等,使DAE成为可微原语,用于推理、控制和设计。

AI 中文摘要

许多工程模型最初是偏微分方程,经空间离散化后变为常微分方程并耦合代数约束,其联合演化构成微分代数方程(DAE)。参数反演、不确定性量化等都需要该求解的梯度,但工业因果建模工具和JAX中的可微物理框架未能满足需求。本文展示了可在一个JAX原生套件中统一前向DAE求解及其反向模式灵敏度。jaxdae将自适应BDF、Radau和Rosenbrock积分与Pantelides索引约简和虚拟导数配对,通过冻结接受步长网格并在其上重新求解可变步长BDF - 2进行反向传播,使自适应BDF路径可微。完整管道在一个`this http URL`调用下进行微分,XLA将批处理参数扫描融合为一个程序,其运行时间从批处理1到1000几乎保持不变,DAE成为推理、控制和设计的可微原语。

英文摘要

Many engineered models begin as partial differential equations. Spatial discretization converts them into ordinary differential equations coupled to algebraic constraints---conservation closures, constitutive laws, network topology---whose joint evolution is a differential-algebraic equation (DAE). Parameter inversion, uncertainty quantification, Bayesian inference, and optimal control all require gradients of this solve. The two software traditions that should supply them have not met: industrial acausal modeling tools simulate DAEs forward but stop at reverse-mode differentiation, while differentiable-physics frameworks in JAX handle explicit ODEs and PDEs but leave the algebraic-constraint layer untouched. Here we show that the forward DAE solve and its reverse-mode sensitivity can be unified in one JAX-native suite. jaxdae pairs adaptive BDF, Radau, and Rosenbrock integration with Pantelides index reduction and dummy derivatives, and makes the adaptive BDF path differentiable by freezing the accepted step grid and re-solving a variable-step BDF-2 on it for the backward pass. The full pipeline differentiates under one $\texttt{jax.grad}$ call, XLA fuses a batched parameter sweep into one program whose wall time stays nearly flat from batch~1 to~1000, and the DAE becomes a differentiable primitive for inference, control, and design.

论文原文

arXiv 摘要页 · PDF 原文 · HTML 原文

↑