「FEAX」で有限要素解析の勾配がjax.gradに渡すだけで取れた

どんなアプリか

A simulation node in your JAX workflow.

Naruki Ichihara さん、GitHub リポジトリ Naruki-Ichihara/feax の説明

FEAX は、JAX で有限要素解析（FEA）を行うための Python のライブラリです。行列の組み立て、境界条件、線形・非線形の求解、後処理までの解析全体を純粋な関数で書き、一つの静的な XLA のプログラムにまとめます。

まとめたプログラムは jax.jit、jax.grad、jax.vmap と組み合わせられ、解析の全体を通して勾配を求められます。README では、これによってトポロジー最適化、逆問題、設計の最適化ができると説明しています。

連立方程式の解き方は、GPU で cuDSS を使う疎行列の直接法、行列を作らない Krylov 法（CG / BiCGSTAB / GMRES）、代数的マルチグリッド（AMG）から選べます。トポロジー最適化用の gene と、周期的な単位セルの均質化用の flat という道具も入っています。

README の例では、表面に荷重をかけた3次元の片持ち梁を解いたあと、解を使った目的関数を jax.grad で微分しています。ソルバーの設定を差し替えれば、ほかのコードを変えずに解き方を切り替えられます。

入れる前に

README のバッジでは、Python 3.12 以上と JAX 0.11 以上が前提になっています。AMG を使うときは pip install feax[amg] で入れます。インストールとソルバーの設定はインストールガイドで説明されています。

ライセンスは GNU General Public License v3.0 です。