shuhengliao / jax-fem

Differentiable Finite Element Method with JAX

A GPU-accelerated differentiable finite element analysis package based on JAX. Part of the suite of open-source python packages for Additive Manufacturing (AM) research, JAX-AM.

Finite Element Method (FEM)

FEM is a powerful tool for thermal-mechanical analysis, especially in Additive manfacturing (AM). We support the following features

• 2D quadrilateral/triangle elements
• 3D hexahedron/tetrahedron elements
• First and second order elements
• Dirichlet/Neumann/Cauchy/periodic boundary conditions
• Linear and nonlinear analysis including
  • Heat equation
  • Linear elasticity
  • Hyperelasticity
  • Plasticity (macro and crystal plasticity)
• Differentiable simulation for solving inverse/design problems without human deriving sensitivities, e.g.,
  • Topology optimization
  • Optimal thermal control
• Integration with PETSc for solver choices

Thermal profile in direct energy deposition.

Linear static analysis of a bracket.

Crystal plasticity: grain structure (left) and stress-xx (right).

Topology optimization with differentiable simulation.

Installation

Clone the repository:

git clone [email protected]:tianjuxue/jax-fem.git
cd jax-fem

Create a conda environment from the given environment.yml file:

conda env create -f environment.yml

Activate the environment and install the package:

conda activate jax-fem-env
pip install -e .

Dependencies

Install JAX

• See jax installation instructions

Documentation

Please see the web documentation for the installation and use of this project.

License

This project is licensed under the GNU General Public License v3 - see the LICENSE for details.

JAX-AM

JAX-AM is a collection of several numerical tools, currently including Discrete Element Method (DEM), Lattice Boltzmann Methods (LBM), Computational Fluid Dynamics (CFD), Phase Field Method (PFM) and Finite Element Method (FEM), that cover the analysis of the Process-Structure-Property relationship in AM.

Our vision is to share with the AM community a free, open-source (under the GPL-3.0 License) software that facilitates the relevant computational research. In the JAX ecosystem, we hope to emphasize the potential of JAX for scientific computing. At the same time, AI-enabled research in AM can be made easy with JAX-AM.

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

@article{xue2022physics,
  title={Physics-embedded graph network for accelerating phase-field simulation of microstructure evolution in additive manufacturing},
  author={Xue, Tianju and Gan, Zhengtao and Liao, Shuheng and Cao, Jian},
  journal={npj Computational Materials},
  volume={8},
  number={1},
  pages={201},
  year={2022},
  publisher={Nature Publishing Group UK London}
}