deepmodeling/jax-fem

Differentiable Finite Element Method with JAX

What it solves

JAX-FEM provides a way to perform finite element method (FEM) simulations that are fully differentiable. This allows researchers and engineers to solve inverse problems and perform design optimization (such as topology optimization) without having to manually derive complex mathematical sensitivities.

How it works

Built on top of JAX, the package combines Automatic Differentiation (AD) with the Finite Element Method. It supports a wide range of element types (2D and 3D) and boundary conditions, and integrates with PETSc for solver options. It can handle both linear and nonlinear analysis, including heat equations, elasticity, and plasticity.

Who it’s for

It is designed for scientists and engineers working in mechanistic data science, structural analysis, and optimal design who need GPU-accelerated, differentiable simulations.

Highlights

  • Differentiable Programming: Enables solving inverse and design problems without manual sensitivity derivation.
  • Broad Physics Support: Handles heat equations, linear/hyperelasticity, and crystal plasticity.
  • Multi-physics Capabilities: Supports solving multiple variables monolithically.
  • Flexible Elements: Supports 2D (quadrilateral/triangle) and 3D (hexahedron/tetrahedron) elements of first and second order.

Related

  • Project
    jax-md/jax-md

    A hardware-accelerated, differentiable molecular dynamics library built on JAX for simulating materials and small-scale interactions.

  • Project
    ymahlau/fdtdx

    FDTDX is a JAX-based Python package for high-performance 3D electromagnetic simulations of photonic nanostructures, featuring automatic differentiation and GPU acceleration.

  • Project
    tumaer/JAXFLUIDS

    A fully-differentiable CFD solver for 3D compressible flows, enabling the integration of machine learning and numerical fluid dynamics through automatic differentiation.

  • Project
    dynamiqs/dynamiqs

    A GPU-accelerated and differentiable Python library for high-performance simulation of quantum systems, designed for large-scale problems and gradient-based parameter estimation.

  • Project
    ekzhang/jax-js

    A high-performance machine learning framework for the browser that brings JAX-style numerical computing, automatic differentiation, and GPU acceleration via WebGPU and Wasm.