Bringing PDEs to JAX with forward and reverse modes automatic differentiation

Bringing PDEs to JAX with forward and reverse modes automatic differentiation
复制标题

通过前向和反向模式自动微分将偏微分方程引入 JAX

DOI:
10.48550/arxiv.2309.07137
复制
发表时间:
2020
期刊:
ArXiv
影响因子:
--
通讯作者:
I. Yashchuk
I. Yashchuk
中科院分区:
--
文献类型:
--
作者:
I. Yashchuk

文献摘要

参考文献

被引文献

相似文献

偏微分方程 (PDE) 用于描述各种物理现象。通常这些方程没有解析解,而是使用数值近似。求解偏微分方程的常用方法之一是有限元法。在科学计算的许多任务中,计算解相对于输入参数的导数信息非常重要。我们通过 Firedrake 有限元库的接口扩展了 JAX 自动微分库。偏微分方程的高级符号表示允许通过底层非线性求解器的低级可能的多次迭代来绕过微分。 Firedrake 求解器的微分是使用切线线性方程和伴随方程完成的。这使得有限元求解器与任意可微分程序的有效组合成为可能。该代码可在 github.com/IvanYashchuk/jax-firedrake 获取。
Partial differential equations (PDEs) are used to describe a variety of physical phenomena. Often these equations do not have analytical solutions and numerical approximations are used instead. One of the common methods to solve PDEs is the finite element method. Computing derivative information of the solution with respect to the input parameters is important in many tasks in scientific computing. We extend JAX automatic differentiation library with an interface to Firedrake finite element library. High-level symbolic representation of PDEs allows bypassing differentiating through low-level possibly many iterations of the underlying nonlinear solvers. Differentiating through Firedrake solvers is done using tangent-linear and adjoint equations. This enables the efficient composition of finite element solvers with arbitrary differentiable programs. The code is available at github.com/IvanYashchuk/jax-firedrake.
DOI: 10.1137/18m1209465
发表时间: 2019
影响因子: 3.1
作者:
Maddison J
通讯作者: Maddison J
DOI: 10.1021/acscatal.5b02014
发表时间: 2016-01-01
期刊: ACS CATALYSIS
影响因子: 12.9
作者:
Nicholls, Thomas P.;Constable, Grace E.;Bissember, Alex C.
通讯作者: Bissember, Alex C.