Simple open source topology optimization example

Topology optimization is generating a lot of hype lately. More powerful computational resources and the demand for more automation in the product development workflow are pushing the field to new areas of application.

Despite its apparent magic, topology optimization is accessible to anyone thanks to the modern open source finite element libraries. Two main tools facilitate this task: the UFL language¹ and Pyadjoint². The former lets you code the problem in a mathematical language and the latter calculates the necessary gradients for the optimization. Both Firedrake and FEniCS use UFL as their frontend interface, but their computational backend is different. Firedrake is more integrated with the PETSc library, which is very helpful for setting solver options required for large scale problems. For this reason, I will use Firedrake in this post.

There exist other open source alternatives for topology optimization. The "Efficient topology optimization in MATLAB using 88 lines of code."³ and the "Large scale topology optimization code using PETSc"⁴ are popular in the academic community. The former because of its simplicity and the latter because of its power to solve large scale problems. There is demand for an alternative that is easy to use, extend and scale. Firedrake helps us to achieve this goal. On a related note, there is an alternative in the Julia world that looks promising (TopOpt.jl). I will be watching this package closely as the scientific computing ecosystem in Julia grows. From a 10,000-foot view, they still lack a diverse set of solver options for large scale problems (specially for building tailored preconditioners.)

The goal of this post is to show how one can solve a topology optimization problem with few lines of easy to read, modify and extend code. Here, I design a bridge to sustain a load with a constrained total mass. All the code used is available in this Github repo. It contains an MIT license, so please feel free to use it at your disposal.

We first create the mesh with the Gmsh script bridge.geo. To speed up the simulation, we only simulate one quarter of the domain and apply the corresponding symmetries on the boundaries.

Please refer to bridge.geo for the dimensions. The markers are: 0, deck; 1, design domain; 2, bridge support; 3, applied load; 4, YZ symmetry plane; 5, XZ symmetry plane.

Our goal is to minimize the compliance

$$ \begin{align} \int_{\Gamma_3} \mathbf{t} \cdot \mathbf{u} ~d \partial\Omega,\\ \end{align} $$

where $\mathbf{t}$ is the applied load and the displacement $\mathbf{u}$ satisfies the linear elasticity equations in their weak form

$$ \begin{align} a(\nu ; \mathbf{u}, \mathbf{v})&=L(\mathbf{v}) ~\text{for all} ~\mathbf{v} \in W, \\ \end{align} $$

with a constrained total mass of 15% of the design domain $\Omega_1$

$$ \begin{align} \int_{\Omega_1} \nu d \Omega \leq 0.15 \int_{\Omega_1} d \Omega \end{align} $$

The volume fraction/density field $\nu$ varies between 0 and 1 and defines the geometry. The bilinear form of the elasticity problem is split in two regions: the design domain $\Omega_1$, where the geometry can vary and the deck $\Omega_0$, completely solid.

$$ \begin{align} a(\nu ; \mathbf{u}, \mathbf{v}) = \int_{\Omega_1} r(\nu)\mathbb{C}[\nabla \mathbf{u}]: \nabla \mathbf{v}d\Omega + \int_{\Omega_0} \mathbb{C}[\nabla \mathbf{u}]: \nabla \mathbf{v} d\Omega \end{align} $$

the right hand side $L( \mathbf{v})$ is simply a vertical traction on top of the deck

$$ \begin{align} L( \mathbf{v}) = \int_{\Gamma_3} \mathbf{t} \cdot \mathbf{v}~d\partial\Omega \end{align} $$

We start the Python script by importing the packages we are using

import firedrake as fd
import firedrake_adjoint as fda
from firedrake import inner, sqrt, jump, dx, ds, dS, sym, nabla_grad, tr, Identity
import itertools
import argparse

Here I am employing Python implementation of the MMA optimization algorithm⁵

from pyMMAopt import MMASolver, ReducedInequality

We separate the solver options in a different file (solver_parameters.py) for tidyness.

from penalization import ramp
from solver_parameters import gamg_parameters, hypre_params

The main function starts with the directory to output the results, the boundary and domain markers constants and the material properties.

def compliance_bridge():
    output_dir = "./"
    DECK = 0
    DOMAIN = 1
    SUPPORT = 2
    LOAD = 3
    SYMMETRY_X = 4
    SYMMETRY_Y = 5
    # Elasticity parameters
    E, nu = 1.0, 0.3
    mu, lmbda = (
        fd.Constant(E / (2 * (1 + nu))),
        fd.Constant(E * nu / ((1 + nu) * (1 - 2 * nu))),
    )

We then read the mesh and create the volume fraction function space. We use element based volume fraction variables, i.e. a zeroth-order Discontinuous Galerkin space.

    mesh = fd.Mesh("./bridge.msh")

    RHO = fd.FunctionSpace(mesh, "DG", 0)
    rho = fd.Function(RHO).assign(fd.Constant(0.5))

The compliance problem is mathematically ill-posed because the optimal solution is a geometry that oscillates between void and material with an infinite frequency. This causes the checkerboard problem⁶ and the optimal solution keeps changing as we refine the mesh. The mathematical details behind this problem are out of the scope of this post, but can be found in the literature⁷.

To regularize the problem, we use the popular Helmholtz filter⁸. It is actually a diffusion-reaction equation and not a Helmholtz equation, but the name has stuck in the community. There are other alternatives to regularize the problem. They mostly uncouple the discretization of the geometry and the displacement. However, for convenience we will simply use the Helmholtz filter. Because of our zeroth-order discretization, we solve it with a Two-Point Flux Approximation scheme:

    af, b = fd.TrialFunction(RHO), fd.TestFunction(RHO)

    filter_radius = fd.Constant(2e-4)
    x, y, z = fd.SpatialCoordinate(mesh)
    with fda.stop_annotating():
        x_ = fd.interpolate(x, RHO)
        y_ = fd.interpolate(y, RHO)
        z_ = fd.interpolate(z, RHO)
    Delta_h = sqrt(jump(x_) ** 2 + jump(y_) ** 2 + jump(z_) ** 2)
    aH = filter_radius * jump(af) / Delta_h * jump(b) * dS + af * b * dx
    LH = rho * b * dx

    rhof = fd.Function(RHO)
    problem_filter = fd.LinearVariationalProblem(aH, LH, rhof)
    solver_filter = fd.LinearVariationalSolver(
        problem_filter, solver_parameters=hypre_params
    )
    solver_filter.solve()
    rhofControl = fda.Control(rhof)

There is a bit of an error here in the flux approximation because we are using an unstructured mesh. Indeed, the line connecting the two elements centroids is not orthogonal to the face they share and therefore, the flux approximation is slightly off. Here we will assume that the quality of the mesh is good enough to limit the error in the approximation. In any case, we are solving this equation to simply regularize the geometry and not to obtain an accurate solution. As an exact alternative, we could have also a mixed formulation as in this Firedrake tutorial, where the fluxes are also resolved using Brezzi-Douglas-Marini finite elements. We could have also used linear Lagrange elements to discretize the volume fraction, but pyMMAopt performs better for zeroth-order discretization.

The displacement in the elasticity problem is discretized with a linear Lagrange basis function.

    H1 = fd.VectorElement("CG", mesh.ufl_cell(), 1)
    W = fd.FunctionSpace(mesh, H1)

To solve large scale problems, we need an iterative solver. Direct solvers scale poorly and eventually run out of memory. We use the Conjugate Gradient solver preconditioned with the Geometric Algebraic Multigrid (GAMG), readily available in PETSc. GAMG needs the null space of the elasticity operator to construct the prolongation operator between the grids. As such, we assemble the three translation and three rotation modes.

    x, y, z = fd.SpatialCoordinate(mesh)
    modes = [fd.Function(W) for _ in range(6)]
    modes[0].interpolate(fd.Constant([1, 0, 0]))
    modes[1].interpolate(fd.Constant([0, 1, 0]))
    modes[2].interpolate(fd.Constant([0, 0, 1]))
    modes[3].interpolate(fd.as_vector([0, z, -y]))
    modes[4].interpolate(fd.as_vector([-z, 0, x]))
    modes[5].interpolate(fd.as_vector([y, -x, 0]))
    nullmodes = fd.VectorSpaceBasis(modes)
    # Make sure they're orthonormal.
    nullmodes.orthonormalize()

The formulation of the elasticity problem is straightforward

    u = fd.TrialFunction(W)
    v = fd.TestFunction(W)

    def epsilon(u):
        return sym(nabla_grad(u))

    def sigma(v):
        return 2.0 * mu * epsilon(v) + lmbda * tr(epsilon(v)) * Identity(3)

    # Variational forms
    a = inner(ramp(rhof, ramp_p=10.0, val_0=1e-5) * sigma(u), nabla_grad(v)) * dx(
        DOMAIN
    ) + inner(sigma(u), nabla_grad(v)) * dx(DECK)
    t = fd.Constant((0.0, 0.0, -1.0))
    L = inner(t, v) * ds(LOAD)

We are using the RAMP function⁹ (defined in penalization.py) to penalize intermediate densities and drive the solution to a void/solid only design. The key here is that we are using different schemes for the volume constraint (defined below) and the stiffness. Indeed, as seen in the figure above, intermediate densities have a high volume/stiffness ratio. The optimization algorithm will deem these intermediate densities undesirable because they add volume, which is limited by the constraint, without providing much stiffness that we want to maximize (i.e. minimize compliance). Accordingly, the algorithm will remove them.

We then apply the strong boundary conditions: fixed at the support 2 and symmetry boundary conditions at the XZ and YZ planes (4 and 5).

    # Dirichlet BCs
    bc1 = fd.DirichletBC(W, fd.Constant((0.0, 0.0, 0.0)), SUPPORT)
    bc2 = fd.DirichletBC(W.sub(0), fd.Constant(0.0), SYMMETRY_X)
    bc3 = fd.DirichletBC(W.sub(1), fd.Constant(0.0), SYMMETRY_Y)

To solve the problem, we use the GAMG preconditioner with the options specified in the solver_parameters.py file

    u_sol = fd.Function(W)
    problem = fd.LinearVariationalProblem(a, L, u_sol, bcs=[bc1, bc2, bc3])
    solver = fd.LinearVariationalSolver(
        problem, near_nullspace=nullmodes, solver_parameters=gamg_parameters
    )
    solver.solve()

Our cost function is the compliance, which I am scaling here with a factor of 10 (just to accelerate convergence). The design volume is limited to 15% of the total volume in the design domain $\Omega_1$.

    # Cost function
    J = fd.assemble(fd.Constant(1e1) * inner(t, u_sol) * ds(LOAD))
    # Constraint
    VolPen = fd.assemble(rhof * dx(DOMAIN))
    with fda.stop_annotating():
        total_vol = fd.assemble(fd.Constant(1.0) * dx(DOMAIN, domain=mesh))
        Vlimit = 0.15 * total_vol
    VolControl = fda.Control(VolPen)

Some routines to plot the design evolution

    # Plotting
    global_counter1 = itertools.count()
    phi_pvd = fd.File(f"{output_dir}/bridge_evolution.pvd", target_continuity=fd.H1)
    rho_viz_f = fd.Function(RHO, name="rho")

    def deriv_cb(design):
        iter = next(global_counter1)
        if iter % 10 == 0:
            rho_viz_f.assign(rhofControl.tape_value())
            # Plot the deck as solid (1.0)
            fd.par_loop(
                ("{[i] : 0 <= i < f.dofs}", "f[i, 0] = 1.0"),
                dx(DECK),
                {"f": (rho_viz_f, fd.WRITE)},
                is_loopy_kernel=True,
            )
            phi_pvd.write(rho_viz_f)

Finally, we write the optimization problem

    c = fda.Control(rho)
    Jhat = fda.ReducedFunctional(J, c, derivative_cb_pre=deriv_cb)
    Volhat = fda.ReducedFunctional(VolPen, c)

    lb = 0.0
    ub = 1.0
    problem = fda.MinimizationProblem(
        Jhat,
        bounds=(lb, ub),
        constraints=[ReducedInequality(Volhat, Vlimit, VolControl)],
    )

and the parameters used for the MMA optimization algorithm

    parameters_mma = {
        "move": 0.2,
        "maximum_iterations": 100,
        "m": 1,
        "IP": 0,
        "tol": 1e-6,
        "accepted_tol": 1e-4,
        "norm": "L2",
        "gcmma": False,
    }
    solver = MMASolver(problem, parameters=parameters_mma)

    rho_opt = solver.solve()

After 100 iterations, the changes in the geometry were minimum. I have found hard to use a convergence criteria based on the KKT optimality system within the MMA. Even though the MMA implementation accepts the Globally Convergent MMA variant, I decided not to use it for this simple problem as it was not necessary. We launch the script to obtain the optimized design

if __name__ == "__main__":
    compliance_bridge()

and plot the design evolution using Paraview to extract the volume fraction 0.8 isocontour, which I then smoothed to alleviate the coarseness of the discretization.

We have finished a scalable topology optimization code for linear elasticity in less than 150 lines of code. One could argue that commercial softwares let you solve this problem in a couple of clicks. That is correct, however, it is important to still have access to the right level of abstraction for your problem. Here we are at the mathematical level which gives us flexibility and convenience to express our problem. We can use this template to rapidly switch to another problem formulation, e.g. to design a compliant mechanism or a stress constrained structure, while keeping the code readable and maintainable. Furthermore, for convenience one could wrap this script into a lightweight module and maybe use a Jupyter notebook to provide a better user experience. One benefit of staying in Python is that we can interact with its rich scientific ecosystem. For instance, we can interact with JAX through jax-fenics-adjoint if we were interested in modeling one of the PDE coefficient with a Neural Network.

I hope you found this post useful. Feel free to comment here with your Github account or in this LinkedIn post.