Differentiable Simulation Tutorials

Companion code for the review paper A Review of Differentiable Programming in Computational Mechanics. The two notebooks illustrate the two common ways of differentiable programming: explicit differentiation through an explicit time-stepping recurrence, and implicit differentiation of an implicitly defined equilibrium. Both implementations are written in JAX, with the FEA support from JAX-FEM.

Repository layout

tutorial_diffsim/
├── README.md                     
├── assets/        
├── data/                         # input and output files      
├── requirements.txt              # additional tutorial dependencies
├── tutorial_explicit.ipynb       # Tutorial 1: explicit differentiation + checkpointing
├── tutorial_implicit.ipynb       # Tutorial 2: implicit differentiation of neural elastoplasticity            

Installation

First follow the JAX-FEM installation instructions, including its JAX setup. Then clone this repository and install the additional tutorial dependencies in the same environment:

git clone https://github.com/xwpken/tutorial_diffsim.git
cd tutorial_diffsim

conda activate jax-fem-env  # replace with your JAX-FEM environment name

pip install -r requirements.txt
jupyter lab tutorial_explicit.ipynb

Tutorial examples

	Tutorial 1	Tutorial 2
Notebook	tutorial_explicit.ipynb	tutorial_implicit.ipynb
Problem	Free vibration of a 3D elastic waveguide	Quasi-static elastoplasticity of a dog-bone specimen
Forward solver	Explicit central difference with row-sum mass lumping	Newton–Raphson with a return-mapping constitutive update
Discretization	1,280 HEX8 elements (2,025 nodes, 6,075 DOF)	7,900 TET4 elements (1,706 nodes, 5,118 DOF)
Target	Recover the spatial Young's modulus field	Identify the neural yield function
Unknowns	80 layerwise Young's moduli	4,417 neural-network weights and biases
Gradient	Explicit differentiation through the time-stepping recurrence	Implicit differentiation of the equilibrium equations

Tutorial 1 — Explicit differentiation of transient elastodynamics

_{Three-dimensional elastic waveguide under longitudinal free vibration. (a) Model setup. (b) Extended and (c) contracted configurations.}

tutorial_explicit.ipynb identifies the axial distribution of Young's modulus in a three-dimensional elastic waveguide from transient displacement observations. Reverse-mode AD differentiates through explicit time integration, and a checkpointing comparison illustrates the memory–runtime trade-off.

_{Free-end displacement histories and identification convergence.}

_{Memory and runtime versus checkpoint interval.}

Tutorial 2 — Implicit differentiation of neural elastoplasticity

_{Dog-bone specimen with displacement loading and the resulting synthetic surface-displacement observations.}

tutorial_implicit.ipynb identifies a neural yield function from displacement observations of a dog-bone specimen under tension, compression, and shear. Gradients are obtained by implicitly differentiating the equilibrium equations. This example adapts the setup of Gavris and Sun (2025) to a finite element formulation.

_{Target, initial and identified yield surfaces for the three cases.}

Citation

If you find this repository useful, please cite the review paper:

To be added