import matplotlib.pyplot as plt
import torch
from tqdm import tqdm

from torchfem.examples import get_example_file
from torchfem.io import import_mesh
from torchfem.materials import IsotropicElasticityPlaneStress

torch.set_default_dtype(torch.float64)

# Material model (plane stress)
material = IsotropicElasticityPlaneStress(E=1000.0, nu=0.3)

# Import mesh
specimen = import_mesh(get_example_file("fillet.vtu"), material)

# Symmetry nodes
left = specimen.nodes[:, 0] < 0.1
specimen.constraints[left, 0] = True
bottom = specimen.nodes[:, 1] < 0.1
specimen.constraints[bottom, 1] = True

# Load at right end
right = specimen.nodes[:, 0] >= specimen.nodes[:, 0].max() - 0.1
specimen.constraints[right, 0] = True
specimen.displacements[right, 0] = 6.0

u, f, σ, ε, α = specimen.solve()
mises = torch.sqrt(σ[:, 0] ** 2 - σ[:, 0] * σ[:, 1] + σ[:, 1] ** 2 + 3 * σ[:, 2] ** 2)
specimen.plot(u, bcs=False, element_property=mises, cmap="inferno")
max_mises = torch.max(mises)
print(f"Maximum v.Mises stress is {max_mises:.2f}.")

Maximum v.Mises stress is 164.19.

# Radial basis function
r = torch.cdist(specimen.nodes, specimen.nodes)
epsilon = 0.2
phi = torch.exp(-((epsilon * r) ** 2))


def morph(fem, nids, x, dir):
    # Build matrix
    M = phi[:, nids]
    M = M[nids, :]
    # Solve linear equation system
    weights = torch.linalg.solve(M, x - fem.nodes[nids, dir])
    # Apply deformation
    fem.nodes[:, dir] += weights @ phi[nids, :]

control_nodes = [3, 5, 7, 58]

# Plot the base mesh
specimen.plot(bcs=False)

# Plot control nodes
cx = specimen.nodes[control_nodes, 0]
cy = specimen.nodes[control_nodes, 1]
plt.scatter(cx, cy, marker="o", color="black")
for node in control_nodes:
    plt.annotate(str(node), (specimen.nodes[node, 0] + 1, specimen.nodes[node, 1] + 1))

def f(x):
    # Detach previous gradients
    specimen.nodes = specimen.nodes.detach()
    # Update nodes
    morph(specimen, control_nodes, x, 1)
    # Solve fem with updated nodes
    _, _, σ, _, _ = specimen.solve()
    # Compute stress
    mises = torch.sqrt(
        σ[:, 0] ** 2 - σ[:, 0] * σ[:, 1] + σ[:, 1] ** 2 + 3 * σ[:, 2] ** 2
    )
    # Compute Kreisselmeier–Steinhauser function to aggregate stresses
    eta = 10.0
    agg = 1 / eta * torch.log(torch.sum(torch.exp(eta * (mises / max_mises - 1))))
    # Return average of top stresses
    return agg

# Initial values (= current y node positions)
x = torch.tensor(specimen.nodes[control_nodes, 1].numpy(), requires_grad=True)
# Optimizer settings
optimizer = torch.optim.Adam([x], lr=0.04)
# Variable to store intermediate values
stresses = []

for _ in tqdm(range(100)):
    optimizer.zero_grad()
    objective = f(x)
    stresses.append(objective.detach().item())
    objective.backward()
    optimizer.step()

  0%|          | 0/100 [00:00<?, ?it/s]

  2%|▏         | 2/100 [00:00<00:05, 17.14it/s]

  4%|▍         | 4/100 [00:00<00:05, 17.24it/s]

  6%|▌         | 6/100 [00:00<00:05, 17.63it/s]

  8%|▊         | 8/100 [00:00<00:05, 17.55it/s]

 10%|█         | 10/100 [00:00<00:05, 17.42it/s]

 12%|█▏        | 12/100 [00:00<00:05, 17.42it/s]

 14%|█▍        | 14/100 [00:00<00:04, 17.39it/s]

 16%|█▌        | 16/100 [00:00<00:04, 17.63it/s]

 18%|█▊        | 18/100 [00:01<00:04, 17.75it/s]

plt.plot(stresses, ".-k")
plt.title("Optimization history")
plt.xlabel("Iteration")
plt.ylabel("Averaged top stress")
plt.show()

u, _, σ, _, _ = specimen.solve()
mises = torch.sqrt(σ[:, 0] ** 2 - σ[:, 0] * σ[:, 1] + σ[:, 1] ** 2 + 3 * σ[:, 2] ** 2)
specimen.plot(u, bcs=False, element_property=mises, figsize=(8, 8), cmap="inferno")
print(f"Maximum v.Mises stress is {torch.max(mises):.2f}.")

Maximum v.Mises stress is 136.09.

Shape optimization of a fillet¶

Import reference model¶

Mesh morphing with radial basis functions¶

Control nodes¶

Shape optimization¶