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

from torchfem.data import get_data
from torchfem.io import import_mesh
from torchfem.materials import Hyperelastic3D

torch.set_default_dtype(torch.float64)

# Load per node in top region
F_max = 2.0

# Increments for loading
increments = torch.linspace(0.0, 1.0, 11)

def psi(F, params):
    """Neo-Hookean strain energy density function."""
    # Extract material parameters
    G = params[0]
    D = params[1]
    # Compute right Cauchy-Green tensor
    C = F.transpose(-1, -2) @ F
    # Stable computation of Jacobian
    J = torch.exp(0.5 * torch.logdet(C))
    C_bar = C * J ** (-2.0 / 3.0)
    return G * (torch.trace(C_bar) - 3.0) + 1 / D * (J - 1) ** 2

material = Hyperelastic3D(psi, params=[0.5, 1.0])


# Import mesh
sample = import_mesh(get_data("iso37.vtu"), material)


# Define boundary sets
top = sample.nodes[:, 2] >= -16.0
bottom = sample.nodes[:, 2] <= -99.0
back = sample.nodes[:, 1] == 0.0
front = sample.nodes[:, 1] == 4.0

# Apply boundary conditions
sample.constraints[bottom, :] = True
sample.constraints[back, 1] = True
sample.constraints[top, 0] = True
sample.constraints[top, 1] = True
sample.forces[top, 2] = F_max

sample.plot(bcs=True)

# Solve
u, f, σ, F, α = sample.solve(
    increments=increments,
    return_intermediate=True,
    verbose=True,
)

----------------------------------------------------------------------------------------
 model    Solid | 120 elem | 900 dof | float64
 machine  AMD EPYC 9V74 80-Core Processor | 2 threads | 16 GB RAM
 solver   direct | scipy | cpu
 newton   rtol 1e-08 | atol 1e-06 | <=10 it | finite strain
----------------------------------------------------------------------------------------
   Increment   Load factor      Steps  Iterations         Residual   Wall time
           1           0.1          1           4         1.63e-11      1.08 s
           2           0.2          1           3         7.19e-09      0.19 s
           3           0.3          1           3         9.79e-09      0.19 s
           4           0.4          1           3         1.20e-07      0.19 s
           5           0.5          1           3         1.06e-07      0.19 s
           6           0.6          1           3         1.34e-08      0.19 s
           7           0.7          1           3         5.62e-09      0.19 s
           8           0.8          1           3         2.98e-09      0.19 s
           9           0.9          1           3         3.43e-08      0.19 s
          10             1          1           3         1.79e-07      0.19 s
----------------------------------------------------------------------------------------
 converged | 10 increments | 31 iterations | 2.80 s

# Plot results
sample.plot(
    u=u[-1],
    element_property={"σ_zz": σ[-1, :, 2, 2]},
    cmap="inferno",
    show_undeformed=True,
)

t_ref = torch.linspace(0.0, 1.0, 11)
u_ref = torch.tensor(
    [0.0, 10.16, 23.44, 40.91, 63.55, 91.96, 126.32, 166.61, 212.75, 264.67, 322.31]
)
u_max, _ = u[:, :, 2].detach().max(dim=1)
plt.plot(t_ref, u_ref, ".-", label="Reference", color="gray")
plt.plot(increments, u_max, ".--", label="torch-fem", color="deeppink")
plt.title("Maximum displacement at top")
plt.legend()
plt.xlabel("Increment")
plt.ylabel("Displacement")
plt.grid()
plt.show()

# Set up optimization
G = torch.tensor(1.0, requires_grad=True)
D = torch.tensor(1.0)
optimizer = torch.optim.LBFGS([G], lr=0.5, max_iter=1)
history = []


def closure():
    # Reset gradients
    optimizer.zero_grad()

    # Build params as a tensor to preserve autograd path to G
    params = torch.stack((G, D))
    sample.material = Hyperelastic3D(psi, params=params).vectorize(sample.n_elem)

    # Solve with differentiable parameters and compute loss
    u, _, _, _, _ = sample.solve(
        increments=increments,
        return_intermediate=True,
        differentiable_parameters=G,
    )
    u_top = u[:, top, 2].mean(dim=1)
    loss = torch.mean(((u_top - u_ref) / (u_ref + 1e-8)) ** 2)
    loss.backward()

    # Append training loss
    history.append(loss.item())
    return loss


for epoch in tqdm(range(12)):
    optimizer.step(closure)

print(f"Optimized G: {G.item():.4f}.")

plt.semilogy(history)
plt.xlabel("Epoch")
plt.ylabel("MSE Loss")
plt.title("Training History")
plt.grid()
plt.show()

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

  8%|▊         | 1/12 [00:02<00:27,  2.46s/it]

 17%|█▋        | 2/12 [00:04<00:24,  2.44s/it]

 25%|██▌       | 3/12 [00:07<00:23,  2.56s/it]

 33%|███▎      | 4/12 [00:10<00:20,  2.54s/it]

 42%|████▏     | 5/12 [00:12<00:17,  2.55s/it]

 50%|█████     | 6/12 [00:15<00:15,  2.54s/it]

 58%|█████▊    | 7/12 [00:17<00:12,  2.52s/it]

 67%|██████▋   | 8/12 [00:20<00:10,  2.57s/it]

 75%|███████▌  | 9/12 [00:22<00:07,  2.55s/it]

Rubber stretch¶