import matplotlib.pyplot as plt
import numpy as np
import torch
from scipy.optimize import bisect
from tqdm import tqdm

torch.set_default_dtype(torch.double)


from torchfem import Planar
from torchfem.rotations import planar_rotation

# Orthotropic material parameters
E_1 = 100000.0
E_2 = 10000.0
nu_12 = 0.3
G_12 = 5000.0

# Mesh parameters
Nx = 60
Ny = 30

# Optimization parameters (SIMP penalization, volume fraction, move limit, step size)
p = 3
volfrac = 0.4
move = 0.1
ori_step = 1000.0

# Volume constraint
V_0 = volfrac * Nx * Ny

from torchfem.materials import OrthotropicElasticityPlaneStress

mat = OrthotropicElasticityPlaneStress(E_1, E_2, nu_12, G_12)

# Create nodes
n1 = torch.linspace(0.0, Nx, Nx + 1)
n2 = torch.linspace(0.0, Ny, Ny + 1)
n1, n2 = torch.stack(torch.meshgrid(n1, n2, indexing="ij"))
nodes = torch.stack([n1.ravel(), n2.ravel()], dim=1)

# Create elements connecting nodes
elements = []
for i in range(Nx):
    for j in range(Ny):
        n0 = j + i * (Ny + 1)
        elements.append([n0, n0 + Ny + 1, n0 + Ny + 2, n0 + 1])
elements = torch.tensor(elements)

# Create the beam model
beam = Planar(nodes, elements, mat)

# Load at top
beam.forces[(Nx + 1) * (Ny + 1) - Ny - 1, 1] = -1

# Constrained displacement at left end
beam.constraints[nodes[:, 0] == 0, :] = True

# Create and plot the domain
beam.plot()

def compliance(rho, phi):
    # Rotate
    R = planar_rotation(phi)
    beam.material = mat.vectorize(beam.n_elem).rotate(R)
    # Apply density
    beam.material.C = rho[:, None, None] ** p * beam.material.C
    u, f, _, _, _ = beam.solve()
    return 0.5 * torch.inner(u.ravel(), f.ravel())

filter_radius = 1.5
ecenters = torch.stack([torch.mean(nodes[e], dim=0) for e in elements])
dist = torch.cdist(ecenters, ecenters)
H = filter_radius - dist
H[dist > filter_radius] = 0.0

def optimize(rho, phi, n_iter=100):
    # Bounds
    rho_min = 0.001 * torch.ones_like(rho)
    rho_max = torch.ones_like(rho)

    # Storage for compliance
    energies = []

    for _ in tqdm(range(n_iter)):
        C = compliance(rho, phi)
        dC_dphi = torch.autograd.grad(C, phi, retain_graph=True)[0]
        dC_drho = torch.autograd.grad(C, rho)[0]
        dC_drho = H @ (rho * dC_drho) / H.sum(dim=0) / rho

        # For a certain value of mu, apply the iteration scheme
        def make_step(mu):
            G_k = -dC_drho / mu
            upper = torch.min(rho_max, (1 + move) * rho)
            lower = torch.max(rho_min, (1 - move) * rho)
            rho_trial = G_k**0.5 * rho
            return torch.max(torch.min(rho_trial, upper), lower)

        # Constraint function
        def g(mu):
            rho_k = make_step(mu)
            return rho_k.sum() - V_0

        # Find the root of g(mu)
        with torch.no_grad():
            mu = bisect(g, 1e-10, 100.0)

        # Advance rho and phi
        rho.data = make_step(mu)
        phi.data -= ori_step * dC_dphi.data

        # Track compliance
        with torch.no_grad():
            energies.append(C.item())
    return energies

@torch.no_grad()
def plot_solution(rho, phi, ptype="element", color="lightskyblue"):
    # Compute properties
    centers = beam.nodes[beam.elements, :].mean(dim=1)
    dir = torch.stack([torch.cos(phi), -torch.sin(phi), torch.zeros_like(phi)]).T

    # Plot orientations
    beam.plot(element_property=rho, linewidth=0.2, cmap="gray_r")
    if ptype == "element":
        plt.quiver(
            centers[:, 0],
            centers[:, 1],
            dir[:, 0],
            dir[:, 1],
            pivot="middle",
            headlength=0,
            headaxislength=0,
            headwidth=0,
            width=0.003,
            color=color,
            alpha=rho,
        )
    if ptype == "stream":
        dir[rho < 0.75] = 0.0
        plt.streamplot(
            centers[:, 0].reshape(Nx, Ny).numpy()[:, 0],
            centers[:, 1].reshape(Nx, Ny).numpy()[0, :],
            dir[:, 0].reshape(Nx, Ny).T.numpy(),
            dir[:, 1].reshape(Nx, Ny).T.numpy(),
            color=color,
            broken_streamlines=False,
            arrowstyle="-",
            cmap="coolwarm",
        )
    plt.show()

# Design variables
phi = 0.0 * torch.ones((len(beam.elements)), requires_grad=True)
rho = volfrac * torch.ones((len(beam.elements)), requires_grad=True)

energies = optimize(rho, phi, 200)
plt.semilogy(energies, ".-k")
plt.grid()
plt.title("Optimization history")
plt.xlabel("Iteration")
plt.ylabel("Compliance")
plt.show()

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

  0%|          | 1/200 [00:00<00:38,  5.12it/s]

  1%|          | 2/200 [00:00<00:37,  5.23it/s]

  2%|▏         | 3/200 [00:00<00:37,  5.22it/s]

  2%|▏         | 4/200 [00:00<00:36,  5.31it/s]

  2%|▎         | 5/200 [00:00<00:36,  5.31it/s]

  3%|▎         | 6/200 [00:01<00:36,  5.29it/s]

  4%|▎         | 7/200 [00:01<00:36,  5.26it/s]

  4%|▍         | 8/200 [00:01<00:36,  5.22it/s]

  4%|▍         | 9/200 [00:01<00:36,  5.18it/s]

plot_solution(rho, phi)

u, f, σ, _, _ = beam.solve()
e = torch.stack(
    [torch.linalg.eigvalsh(torch.tensor([[s[0], s[2]], [s[2], s[1]]])) for s in σ]
)
mag, pos = torch.max(torch.abs(e), dim=1)
pval = torch.tensor([ev[p] for ev, p in zip(e, pos)])
pval = pval.reshape(Nx, Ny).T.numpy()
plot_solution(rho, phi, ptype="stream", color=pval)

Concurrent optimization of orientation and topology¶

Parameters¶

Material model¶

Mesh generation¶

Objective function¶

Filter¶

The optimizer¶

Optimize with fibers initially aligned in x¶

Generate fiber paths with streamlines¶