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, σ, F, α = specimen.solve()
mises = torch.sqrt(
    σ[:, 0, 0] ** 2 - σ[:, 0, 0] * σ[:, 1, 1] + σ[:, 1, 1] ** 2 + 3 * σ[:, 0, 1] ** 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, 0] ** 2
        - σ[:, 0, 0] * σ[:, 1, 1]
        + σ[:, 1, 1] ** 2
        + 3 * σ[:, 0, 1] ** 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]

  3%|▎         | 3/100 [00:00<00:03, 25.99it/s]

  6%|▌         | 6/100 [00:00<00:03, 26.55it/s]

  9%|▉         | 9/100 [00:00<00:03, 26.34it/s]

 12%|█▏        | 12/100 [00:00<00:03, 26.17it/s]

 15%|█▌        | 15/100 [00:00<00:03, 26.40it/s]

 18%|█▊        | 18/100 [00:00<00:03, 26.51it/s]

 21%|██        | 21/100 [00:00<00:02, 26.50it/s]

 24%|██▍       | 24/100 [00:00<00:02, 26.60it/s]

 27%|██▋       | 27/100 [00:01<00:02, 26.68it/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, 0] ** 2 - σ[:, 0, 0] * σ[:, 1, 1] + σ[:, 1, 1] ** 2 + 3 * σ[:, 0, 1] ** 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¶