Exercise 06 - Shape optimization of trusses

Shape optimization means that given a fixed topology of a truss, we want to optimize its stiffness by modifying some node positions. In this particular example, we investigate the optimal shape of a railway bridge like in the photograph here:

from math import sqrt

import matplotlib.pyplot as plt
import torch
from tqdm import tqdm
from torchfem import Truss
from torchfem.materials import IsotropicElasticity1D

torch.set_default_dtype(torch.double)

So let's start by defining the base truss topology of the bridge without considering the exact shape for now. We create a simple rectangular bridge that has all the bars seen in the photo. The truss is fixed at the bottom left side and simply supported at the bottom right side. The load is distributed along the bottom edge of the bridge, which represents the train track.

# Dimensions
A = 17
B = 2

# Nodes
n1 = torch.linspace(0.0, 5.0, A)
n2 = torch.linspace(0.0, 0.5, B)
n1, n2 = torch.stack(torch.meshgrid(n1, n2, indexing="xy"))
nodes = torch.stack([n1.ravel(), n2.ravel()], dim=1)

# Elements
connections = []
for i in range(A - 1):
    for j in range(B):
        connections.append([i + j * A, i + 1 + j * A])
for i in range(A):
    for j in range(B - 1):
        connections.append([i + j * A, i + A + j * A])
for i in range(A - 1):
    for j in range(B - 1):
        if i >= (A - 1) / 2:
            connections.append([i + j * A, i + 1 + A + j * A])
        else:
            connections.append([i + 1 + j * A, i + A + j * A])
elements = torch.tensor(connections)

# Create material
material = IsotropicElasticity1D(E=500.0)

# Create Truss
bridge = Truss(nodes.clone(), elements, material)

# Forces at bottom edge
bridge.forces[1 : A - 1, 1] = -0.1

# Constraints by the supports
bridge.constraints[0, 0] = True
bridge.constraints[0, 1] = True
bridge.constraints[A - 1, 1] = True

# Plot
bridge.plot()

We introduce some helper functions from the last exercises. These are just copies of functions you are already familiar with.

def compute_lengths(truss):
    start_nodes = truss.nodes[truss.elements[:, 0]]
    end_nodes = truss.nodes[truss.elements[:, 1]]
    dx = end_nodes - start_nodes
    return torch.linalg.norm(dx, dim=-1)


def bisection(f, a, b, max_iter=50, tol=1e-10):
    i = 0
    while (b - a) > tol:
        c = (a + b) / 2.0
        if i >= max_iter:
            raise Exception(f"Bisection did not converge in {max_iter} iterations.")
        if f(a) * f(c) > 0:
            a = c
        else:
            b = c
        i += 1
    return c


def MMA(func, x_k, L_k, U_k):
    x_lin = x_k.clone().requires_grad_()
    grads = torch.autograd.grad(func(x_lin), x_lin)[0]
    pg = grads >= 0.0
    ng = grads < 0.0
    f_k = func(x_k)

    def approximation(x):
        p = torch.zeros_like(grads)
        p[pg] = (U_k[pg] - x_k[pg]) ** 2 * grads[pg]
        q = torch.zeros_like(grads)
        q[ng] = -((x_k[ng] - L_k[ng]) ** 2) * grads[ng]
        return (
            f_k
            - torch.sum(p / (U_k - x_k) + q / (x_k - L_k))
            + torch.sum(p / (U_k - x) + q / (x - L_k))
        )

    return approximation, grads

Task 1 - Preparation of design variables

We want to restrict the shape optimization problem to the vertical displacement of the top nodes. Other nodes should not be modified - the train track should remain a flat line.

a) Create a boolean tensor mask of the same shape as nodes. It should be True for the vertical degrees of freedom of the top nodes and False for every other degree of freedom. Essentially it should mask out those nodal degrees of freedom which should be optimized.

Hints: Take a look at how the boolean constraints tensor is created in a cell above. Take a look at the plot of the bridge to see which node numbers are the top nodes.)

# Mask for design variables.
mask = torch.zeros_like(nodes, dtype=torch.bool)
mask[A : 2 * A, 1] = True

b) Create the initial design vector x_0 from the masked top-node vertical coordinates. Define lower and upper bounds (x_min, x_max) so each design variable may move down by 0.4 and up by 0.5 relative to its initial value.

Hint: Use mask to extract the top-node coordinates and ravel() to obtain a 1D design vector.

# Limits on design variables
x_0 = nodes[mask].ravel()
x_min = x_0 - 0.4
x_max = x_0 + 0.5

c) Compute the current volume of the truss V0. We will use this as a constraint in the optimization problem such that the optimized solution does not exceed this initial volume.

Hint: The current volume is the inner product of the current bar lengths and the cross-sectional areas of the bars.

# Compute current volume
l = compute_lengths(bridge)
V0 = torch.inner(bridge.areas, l)

Task 2 - Optimization

In this task, we will optimize the shape of the truss by minimizing the compliance of the bridge by adjusting the top nodes while keeping the volume constant:

$$ \min_{\mathbf{x}} \quad C(\mathbf{x}) = \mathbf{f} \cdot \mathbf{u}(\mathbf{x}) \\ \text{s.t.} \quad \mathbf{a} \cdot \mathbf{l}(\mathbf{x}) - V_0 \le 0\\ \quad \quad x \in [\mathbf{x}^{-}, \mathbf{x}^{+}] $$

a) Complete the objective function f(x).

Hint: Replace bridge.nodes[mask] with the design variables x to update the nodal positions of the top nodes. Use truss.solve() to compute the displacements and forces of the updated truss. Use ravel() to flatten the tensors and compute the compliance as the inner product of the forces and displacements.

def f(x):
    # Update nodes
    bridge.nodes[mask] = x
    # Solve truss with updated nodes
    u_k, f_k, _, _, _ = bridge.solve(differentiable_parameters=x)
    # Return compliance
    return torch.inner(f_k.ravel(), u_k.ravel())

b) Complete the constraint function g(x).

Hint: Replace bridge.nodes[mask] with the design variables x to update the nodal positions of the top nodes.

def g(x):
    # Update nodes
    bridge.nodes[mask] = x
    # Return constraint function
    return torch.inner(bridge.areas, compute_lengths(bridge)) - V0

c) Implement the updates of upper and lower asymptotes as in Exercise 04.

d) Implement the lower and upper move limits called x_min_k and x_max_k.

e) Compute the approximated functions f_tilde and g_tilde.

Hints: The MMA function returns the approximated functions f_tilde as well as the gradients of the original functions. You can assign them as follows: f_tilde, f_grad = MMA(...).

f) Implement the analytical solution of x_star(mu).

Hints: This is similar to the analytical solutions in Exercise 04. Initialize x_hat with x[k].clone() and overwrite components according to the four sign cases. Define the conditions for the four cases of the analytical solution and fill the x_hat tensor accordingly by setting x_hat[cond1] = solution1 etc. Finally, clamp with x_min_k and x_max_k.

h) Implement the gradient of the dual function $\frac{\partial \underbar{L}(\mu)}{\partial \mu} = \tilde{g}(x^*(\mu))$

i) Solve for the maximum of the dual function with bisection(...).

j) Append the solution to the list of solutions x = [$\mathbf{x}^0, \mathbf{x}^1, \mathbf{x}^2, ...$] containing the iteration steps of the optimization procedure.

s = 0.8

# Set up lists for L, U, x
L = []
U = []
x = [x_0]

# Define lower and upper bounds of "mu" for bisection
mu_lower = torch.tensor([1e-10])
mu_upper = torch.tensor([1.0])

for k in tqdm(range(80), delay=0.1):
    # Update asymptotes with heuristic procedure (see Exercise 04)
    if k <= 1:
        L.append(x[k] - s * (x_max - x_min))
        U.append(x[k] + s * (x_max - x_min))
    else:
        L_k = torch.zeros_like(L[k - 1])
        U_k = torch.zeros_like(U[k - 1])
        # Shrink all oscillating asymptotes
        osci = (x[k] - x[k - 1]) * (x[k - 1] - x[k - 2]) < 0.0
        L_k[osci] = x[k][osci] - s * (x[k - 1][osci] - L[k - 1][osci])
        U_k[osci] = x[k][osci] + s * (U[k - 1][osci] - x[k - 1][osci])
        # Expand all non-oscillating asymptotes
        L_k[~osci] = x[k][~osci] - 1 / sqrt(s) * (x[k - 1][~osci] - L[k - 1][~osci])
        U_k[~osci] = x[k][~osci] + 1 / sqrt(s) * (U[k - 1][~osci] - x[k - 1][~osci])
        L.append(L_k)
        U.append(U_k)

    # Compute lower and upper move limit in this step
    x_min_k = torch.max(x_min, 0.9 * L[k] + 0.1 * x[k])
    x_max_k = torch.min(x_max, 0.9 * U[k] + 0.1 * x[k])

    # Compute the current approximation function and save gradients
    f_tilde, f_grad = MMA(f, x[k], L[k], U[k])
    g_tilde, g_grad = MMA(g, x[k], L[k], U[k])

    # Define x_star by minimizing the Lagrangian w. r. t. x analytically
    def x_star(mu):
        x_hat = x[k].clone()
        # Case 1
        cond1 = (f_grad > 0.0) & (g_grad > 0.0)
        x_hat[cond1] = x_min[cond1]
        # Case 2
        cond2 = (f_grad < 0.0) & (g_grad < 0.0)
        x_hat[cond2] = x_max[cond2]
        # Case 3
        cond3 = (f_grad > 0.0) & (g_grad < 0.0)
        roots = torch.sqrt(
            (-mu * g_grad[cond3] * (x[k][cond3] - L[k][cond3]) ** 2)
            / (U[k][cond3] - x[k][cond3]) ** 2
            / f_grad[cond3]
        )
        x_hat[cond3] = (U[k][cond3] * roots + L[k][cond3]) / (1 + roots)
        # Case 4
        cond4 = (f_grad < 0.0) & (g_grad > 0.0)
        roots = torch.sqrt(
            (-mu * g_grad[cond4] * (U[k][cond4] - x[k][cond4]) ** 2)
            / (x[k][cond4] - L[k][cond4]) ** 2
            / f_grad[cond4]
        )
        x_hat[cond4] = (U[k][cond4] + L[k][cond4] * roots) / (1 + roots)
        return torch.clamp(x_hat, x_min_k, x_max_k)

    # Gradient of dual function
    def dual_grad(mu):
        return g_tilde(x_star(mu))

    # Solve for optimal "mu" with bisection method
    mu_star = bisection(dual_grad, mu_lower, mu_upper)

    # Compute current optimal point with dual solution and append to list of iterates
    with torch.no_grad():
        x.append(x_star(mu_star))

100%|██████████| 80/80 [00:02<00:00, 35.66it/s]

# Plot the development of design variables
plt.plot(torch.stack(x).detach())
plt.xlabel("Iteration")
plt.ylabel("Values $x_i$")
plt.show()

# Plot the optimized bridge
bridge.nodes[mask] = x[-1]
bridge.plot(node_labels=False)