Topology optimization of a 3D heat sink

The model is defined on the cube $\Omega$ with length $L=1\,\text{m}$. An homogenous heat source is applied to the solution domain. Simple isotropic Fourier's law connects diffusive heat flux $\boldsymbol{h}$ and temperature gradient via the scalar conductivity $\kappa=400\,\mathrm{W}/\mathrm{m\,K}$.

A constant Dirichlet temperature BC $T_D= 0\,\mathrm{K}$ is prescribed on portion on the bottom face. Homogenous Neumann boundary conditions are prescribed at the remaining faces.

Then, the Poisson-like PDE is defined as:

$$\begin{align} \mathrm{div}\left( \kappa \, \mathrm{grad}\left( T \right) \right) + q = 0 \quad &\in \Omega,\\ \boldsymbol{h}\cdot\boldsymbol{n} = 0 \quad &\in \partial\Omega_\text{N},\\ T = T_\text{D} \quad &\in \partial\Omega_\text{D}. \end{align}$$

Definition of constants

L = 1.0
n_elements = 60  # elements per edge
heat_source = 1000.0  # W in entire domain
kappa = 400.0  # W / m / K
kappa_min = 1.0  # W / m / K

import matplotlib.pyplot as plt
import torch
from scipy.optimize import bisect
from scipy.spatial import KDTree
from tqdm import tqdm

from torchfem import SolidHeat
from torchfem.materials import IsotropicConductivity3D
from torchfem.mesh import cube_hexa

torch.set_default_dtype(torch.float64)

Model setup

We start by defining the base problem without considering the optimization yet.

# Material model
material = IsotropicConductivity3D(kappa=kappa)

# Create mesh
nodes, elements = cube_hexa(
    Nx=n_elements + 1, Ny=n_elements + 1, Nz=n_elements + 1, Lx=L, Ly=L, Lz=L
)
model = SolidHeat(nodes, elements, material)

# Prescribed temperature on a square patch in the center of the bottom face
bottom = torch.isclose(nodes[:, 2], nodes[:, 2].min())
dirichlet = (
    bottom
    & (nodes[:, 0] >= 0.45 * L)
    & (nodes[:, 0] <= 0.55 * L)
    & (nodes[:, 1] >= 0.45 * L)
    & (nodes[:, 1] <= 0.55 * L)
)
model.constraints[dirichlet] = True

# Distribute the total heat source uniformly over the domain
element_volume = model.integrate_field()
model.heat_flux = model.integrate_body_load(heat_source / element_volume.sum())

Optimization parameters

We define the optimization parameters, i.e. the volume fraction, the penalization factor, the move limit, the filter radius, and the number of iterations.

vol_frac = 0.1  # volume fraction
p = 3.0  # SIMP penalization exponent
move = 0.1  # move limit of the optimality criteria update
R = 2.0 * L / n_elements  # filter radius, spanning about two elements
TORCH_SENS = False  # compute sensitivities with autograd instead of analytically

# Initial, minimum, and maximum values of design variables
rho = vol_frac * torch.ones(len(elements))
rho_min = kappa_min / kappa * torch.ones_like(rho)
rho_max = torch.ones_like(rho)

# Target volume
V_0 = vol_frac * element_volume.sum()

# Element conductivity matrices at rho = 1 for the analytical sensitivity
k0 = model.k0()
KAPPA_0 = model.material.KAPPA.clone()

# Precompute linear filter weights between neighboring element centroids
ecenters = nodes[elements].mean(dim=-2)
tree = KDTree(ecenters)
distances = tree.sparse_distance_matrix(tree, R, output_type="coo_matrix")
H = torch.sparse_coo_tensor(
    indices=torch.stack(
        [torch.as_tensor(distances.row), torch.as_tensor(distances.col)]
    ),
    values=R - torch.as_tensor(distances.data),
    size=(len(elements), len(elements)),
    check_invariants=False,
)
H_sum = H.sum(dim=0).to_dense()

Optimization with optimality constraints.

This takes about two minutes to run.

compliance = []

# Iterate solutions
for k in tqdm(range(50)):
    # Make rho a differentiable parameter if using autograd for sensitivities
    rho.requires_grad_(True)

    # SIMP interpolation of the element conductivities
    model.material.KAPPA = rho[:, None, None] ** p * KAPPA_0

    # Compute solution
    u_k, f_k, *_ = model.solve(method="cg", differentiable_parameters=rho)

    # Compliance of the current design
    c_k = torch.inner(f_k.ravel(), u_k.ravel())

    # Sensitivity of the compliance, either by autograd or analytically
    if TORCH_SENS:
        sensitivity = torch.autograd.grad(c_k, rho)[0]
    else:
        u_j = u_k[elements].reshape(model.n_elem, -1)
        w_k = torch.einsum("...i, ...ij, ...j", u_j, k0, u_j)
        sensitivity = -p * rho ** (p - 1.0) * w_k

    # Filter sensitivities
    sensitivity = H @ (rho * sensitivity) / H_sum / rho

    with torch.no_grad():
        # For a certain value of mu, apply the iteration scheme
        def make_step(mu):
            G_k = -sensitivity / mu
            upper = torch.min(rho_max, (1 + move) * rho)
            lower = torch.max(rho_min, (1 - move) * rho)
            rho_trial = G_k**0.5 * rho
            return torch.maximum(torch.minimum(rho_trial, upper), lower)

        # Constraint function
        def g(mu):
            return torch.inner(make_step(mu), element_volume) - V_0

        # Find the root of g(mu)
        mu = bisect(g, 1e-5, 1e5)

        # Update design variables
        rho = make_step(mu)

        # Track the compliance history
        compliance.append(c_k.item())

100%|██████████| 50/50 [02:01<00:00,  2.43s/it]

plt.semilogy(compliance)
plt.title("Optimization history")
plt.xlabel("Iteration")
plt.ylabel("Thermal Compliance")
plt.grid()
plt.show()

model.plot(
    element_property={"rho": rho},
    clip=("rho", 0.5),
    show_edges=False,
    show_outline=True,
    color="lightblue",
)

EmbeddableWidget(value='<iframe srcdoc="<!doctype html>\n<html lang=&quot;en&quot;>\n  <head>\n    <meta chars…