JAX-FEM is a differentiable finite element package based on JAX.
For installation and user guide, please visit our documentation for details.
JAX-FEM is Automatic Differentiation (AD) + Finite Element Method (FEM), and we support the following features:
- 2D quadrilateral/triangle elements
- 3D hexahedron/tetrahedron elements
- First and second order elements
- Dirichlet/Neumann/Robin boundary conditions
- Linear and nonlinear analysis including
- Heat equation
- Linear elasticity
- Hyperelasticity
- Plasticity (macro and crystal plasticity)
- Multi-physics problems
- Integration with PETSc for solver options
- Differentiable programming for solving inverse/design problems without deriving sensitivities by hand, e.g.,
- Topology optimization
- Optimal thermal control
Thermal profile in direct energy deposition.
Linear static analysis of a bracket.
Crystal plasticity: grain structure (left) and stress-xx (right).
Stokes flow: velocity (left) and pressure(right).
Topology optimization with differentiable simulation.
This project is licensed under the GNU General Public License v3 - see the LICENSE for details. For commercial use, contact Tianju Xue.
If you found this library useful in academic or industry work, we appreciate your support if you consider 1) starring the project on Github, and 2) citing relevant papers:
@article{xue2023jax,
title={JAX-FEM: A differentiable GPU-accelerated 3D finite element solver for automatic inverse design and mechanistic data science},
author={Xue, Tianju and Liao, Shuheng and Gan, Zhengtao and Park, Chanwook and Xie, Xiaoyu and Liu, Wing Kam and Cao, Jian},
journal={Computer Physics Communications},
pages={108802},
year={2023},
publisher={Elsevier}
}
