Parametric shape optimization with differentiable FEM simulation

Summary

Parametric shape optimization with differentiable FEM simulation In this example, you will learn how to: Build a Tesseract that wraps a differentiable finite-element solver from jax-fem. Build a Tesseract that uses finite differences under the hood to enable differentiability of a non-autodifferentiable geometry operation (computing a signed distance field from a 3D model). Compose both Tesseracts with Tesseract-JAX to create a pipeline that can be used for differentiable shape optimization. Perform gradient-based optimization using optax on the Tesseract-JAX pipeline. Introduction In this notebook, we explore the optimization of a parametric structure made of a linear elastic material. The structure is parametrized by N bars, each of which has M piecewise linear segments. We seek the ideal configuration of the \(y\)-coordinates of the vertices that connect those bar segments. This notebook is based on the 2D topology optimization example from jax-fem, but we solve the problem using a parametric approach instead. That is, we use end-to-end automatic differentiation (AD) through several components to optimize the design variables directly with respect to (simulated) performance of the design. The design space is defined using a geometry library called PyVista, which does not support automatic differentiation. However, we can enable differentiability of this operation by using a finite difference approximation of the Jacobian matrix. We denote the design space as a function \(g\) that maps the design variables to a signed distance field. Then, we can then define the density field \(\rho(\mathbf{x})\) as a function of a signed distance field (SDF) value \(g(\mathbf{x})\). Finally we denote the differentiable finite element method (FEM) solver as \(f\), which takes the density field as input and returns the structure’s compliance. Therefore, the optimization problem can be formulated as follows: \[ \begin{equation} \min_{\theta} f(\rho(g(\theta))). \end{equation} \] Her...