import matplotlib.pyplot as plt
import meshio
import numpy as np
import pyvista
import torch
from scipy.optimize import bisect
from scipy.spatial import cKDTree
from tqdm import tqdm

from torchfem import Solid
from torchfem.data import get_data
from torchfem.materials import IsotropicElasticity3D

torch.set_default_dtype(torch.float64)

# Material model (Ti-6Al-4V) in imperial units
material = IsotropicElasticity3D(E=16500.0, nu=0.342)

mesh = meshio.read(get_data("ge_bracket.vtu"))
elements = torch.tensor(mesh.cells[0].data)
nodes = torch.tensor(mesh.points)
domain = torch.tensor(mesh.cell_data["gmsh:geometrical"][0])

model = Solid(nodes, elements, material)
cmap = plt.get_cmap("tab10", 7)
model.plot(element_property={"Domain": domain}, cmap=cmap)

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

# Constrain inner nodes at fixation holes
R = 6.0
for d in [1, 3, 4, 5]:
    dom = torch.unique(elements[domain == d])
    center = nodes[dom].mean(dim=0)
    con = (nodes[dom, 0] - center[0]) ** 2 + (nodes[dom, 1] - center[1]) ** 2 < R**2
    model.constraints[dom[con], :] = True


# Define load cases (this is not using multi-point constraints yet...)
dom = torch.unique(elements[(domain == 2) | (domain == 7)])
load_case_1 = torch.zeros_like(nodes)
load_case_1[dom, 2] = 8000 / len(dom)
load_case_2 = torch.zeros_like(nodes)
load_case_2[dom, 1] = -8500 / len(dom)
load_case_3 = torch.zeros_like(nodes)
load_case_3[dom, 1] = -9500 * np.sin(np.deg2rad(42)) / len(dom)
load_case_3[dom, 2] = 9500 * np.cos(np.deg2rad(42)) / len(dom)
load_case_4 = torch.zeros_like(nodes)
levers = nodes[dom] - torch.tensor([0.0, 0.0, 0.0])
load_case_4[dom, 1] = 5000 / levers[:, 0] / len(dom)

# Apply load case 1 for testing
model.forces = load_case_1

# Solve
u, f, σ, F, α = model.solve(rtol=0.01, verbose=True)

─── torch-fem · solve ──────────────────────────────────────────────────────────────────
 model    Solid · 214,524 elem · 127,644 dof · float64
 machine  Apple M1 Pro · 8 threads · 16 GB RAM
 solver   minres · iterative · AMG · scipy · cpu
 newton   rtol 1e-02 · atol 1e-06 · ≤10 it
────────────────────────────────────────────────────────────────────────────────────────
   Increment   Load factor      Steps  Iterations         Residual   Wall time
           1             1          1           1         6.10e-06      4.31 s
────────────────────────────────────────────────────────────────────────────────────────
 converged · 1 increment · 1 iteration · 4.32 s

# Plot
model.plot(node_property={"u": u})

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

volfrac = 0.15
p = 3
move = 0.2
R = 5.0
soft = 1e-3

# Design elements are only the ones in domain 6
design_elements = elements[domain == 6]
vols = model.integrate_field()[domain == 6]

# Initial, minimum, and maximum values of design variables.
rho_0 = volfrac * torch.ones(len(design_elements))
rho_min = 1e-3 * torch.ones_like(rho_0)
rho_max = torch.ones_like(rho_0)

# Volume fraction
V_0 = volfrac * vols.sum()

# Analytical gradient of the stiffness matrix
k0 = model.k0()[domain == 6].clone()
C0 = model.material.C[domain == 6].clone()
C_min = soft * C0

# Precompute filter weights using scipy's KDTree
ecenters = torch.mean(nodes[design_elements], dim=1)
tree = cKDTree(ecenters)
d = tree.sparse_distance_matrix(tree, R, output_type="coo_matrix")

# Convert to torch sparse tensor
H = torch.sparse_coo_tensor(
    indices=torch.stack([torch.as_tensor(d.row), torch.as_tensor(d.col)]),
    values=torch.as_tensor(d.data),
    size=(len(design_elements), len(design_elements)),
)
H_sum = H.sum(dim=0).to_dense()

rho = [rho_0]
history = []

# Iterate solutions
for k in tqdm(range(25)):
    # Interpolate stiffness between void and solid
    model.material.C[domain == 6] = C_min + torch.einsum(
        "n,nijkl->nijkl", rho[k] ** p, C0 - C_min
    )

    sensitivity = torch.zeros_like(rho[k])
    # Iterate over load cases
    for lc in [load_case_1, load_case_2, load_case_3, load_case_4]:
        # Apply load case
        model.forces = lc
        # Compute solution
        u_k, f_k, _, _, _ = model.solve(rtol=0.01)
        # Evaluation of compliance
        compliance = torch.inner(f_k.ravel(), u_k.ravel())
        # Compute analytical sensitivities
        u_j = u_k[design_elements].reshape(len(design_elements), -1)
        w_k = torch.einsum("...i, ...ij, ...j", u_j, k0, u_j)
        sensitivity += -p * rho[k] ** (p - 1.0) * w_k

    # Filter sensitivities (if r provided)
    sensitivity = H @ (rho[k] * sensitivity / vols) / H_sum / (rho[k] / vols)

    # 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[k])
        lower = torch.max(rho_min, (1 - move) * rho[k])
        rho_trial = G_k**0.5 * rho[k]
        return torch.maximum(torch.minimum(rho_trial, upper), lower)

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

    # Find the root of g(mu)
    mu = bisect(g, 1e-10, 100.0)

    # Update design history
    rho.append(make_step(mu))
    history.append(compliance)

100%|██████████| 25/25 [10:18<00:00, 24.73s/it]

plt.plot(history)
plt.xlabel("Iteration")
plt.ylabel("Compliance")
plt.grid()
plt.xlim(0, len(history) - 1)
plt.ylim(bottom=0)
plt.show()

# Non-design elements are solid
rho_final = 2 * torch.ones(len(elements))
rho_final[domain == 6] = rho[-1]

cmap = plt.get_cmap("tab10", 2)

plotter = pyvista.Plotter()
model.plot(plotter=plotter, opacity=0.1)
model.plot(
    plotter=plotter,
    element_property={"rho": rho_final},
    clip=("rho", 0.5),
    show_edges=False,
    cmap=cmap,
)
plotter.show(jupyter_backend="html")

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

Topology optimization of the GE jet engine bracket challenge¶

Mesh¶

Optimization parameters¶

Optimization with optimality constraints.¶

Result¶