JAX-FEM

JAX-FEM is a differentiable finite element package based on JAX.

Documentation

For installation and user guide, please visit our documentation for details.

Key features

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

Examples

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.

JAX-FEM Express

We have built an LLM Agent for JAX-FEM: Try JAX-FEM Express! https://www.bohrium.com/apps/jax-fem-express

Click to watch the JAX-FEM Express video

License

This project is licensed under the GNU General Public License v3 - see the LICENSE for details. For commercial use, contact Tianju Xue.

Citations

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}
}