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)

/home/runner/.local/lib/python3.14/site-packages/IPython/core/display.py:448: UserWarning: Consider using IPython.display.IFrame instead
  warnings.warn("Consider using IPython.display.IFrame instead")

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

────────────────────────────────────────────────────────────────────────────────────────
 model    Solid · 120 elem · 900 dof · float64
 machine  x86_64 · 2 threads · 16 GB RAM
 solver   direct · scipy · cpu
 newton   rtol 1e-08 · atol 1e-06 · <=10 it · finite strain
────────────────────────────────────────────────────────────────────────────────────────
  Increment  Load factor  Steps  Newton  Krylov  Residual  Wall time
          1          0.1      1       4       -  1.63e-11     0.11 s
          2          0.2      1       3       -  7.19e-09     0.08 s

          3          0.3      1       3       -  9.79e-09     0.08 s
          4          0.4      1       3       -  1.20e-07     0.08 s
          5          0.5      1       3       -  1.06e-07     0.11 s

          6          0.6      1       3       -  1.34e-08     0.09 s
          7          0.7      1       3       -  5.62e-09     0.08 s
          8          0.8      1       3       -  2.98e-09     0.08 s

          9          0.9      1       3       -  3.43e-08     0.08 s
         10            1      1       3       -  1.79e-07     0.09 s
────────────────────────────────────────────────────────────────────────────────────────
 converged · 10 increments · 31 Newton iterations · 0.89 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:01<00:13,  1.21s/it]

 17%|█▋        | 2/12 [00:02<00:11,  1.18s/it]

 25%|██▌       | 3/12 [00:03<00:11,  1.26s/it]

 33%|███▎      | 4/12 [00:04<00:09,  1.22s/it]

 42%|████▏     | 5/12 [00:06<00:08,  1.20s/it]

 50%|█████     | 6/12 [00:07<00:07,  1.22s/it]

 58%|█████▊    | 7/12 [00:08<00:06,  1.22s/it]

 67%|██████▋   | 8/12 [00:09<00:04,  1.20s/it]

 75%|███████▌  | 9/12 [00:10<00:03,  1.22s/it]

Rubber stretch¶