Embedded 1D Fibers¤

# Only run this if we are in Google Colab
if 'google.colab' in str(get_ipython()):
    print("Installing dependencies using uv...")
    # Install uv if not available
    !pip install -q uv
    # Install system dependencies
    !apt-get install -qq gmsh
    # Use uv to install Python dependencies
    !uv pip install --system matplotlib meshio
    !uv pip install --system "git+https://github.com/smec-ethz/tatva-docs.git"
    print("Installation complete!")
else:
    import pyvista as pv

import os

import gmsh
import matplotlib.pyplot as plt
import meshio
import numpy as np


def generate_plate_mesh(
    Lx: float, Ly: float, mesh_size: float, work_dir: str = "."
) -> Mesh:
    """
    Generates a 2D unstructured triangular mesh for a rectangular plate.

    Args:
        Lx (float): Length of the plate in the x-direction.
        Ly (float): Length of the plate in the y-direction.
        mesh_size (float): Target mesh size for the mesh generation.
        work_dir (str): Directory to store temporary mesh files.

    Returns:
        Mesh: The generated plate mesh.
    """
    if not os.path.exists(work_dir):
        os.makedirs(work_dir)

    filename = os.path.join(work_dir, "plate_2d.msh")

    gmsh.initialize()
    gmsh.model.add("plate")

    p1 = gmsh.model.geo.addPoint(0, 0, 0, mesh_size)
    p2 = gmsh.model.geo.addPoint(Lx, 0, 0, mesh_size)
    p3 = gmsh.model.geo.addPoint(Lx, Ly, 0, mesh_size)
    p4 = gmsh.model.geo.addPoint(0, Ly, 0, mesh_size)

    l1 = gmsh.model.geo.addLine(p1, p2)
    l2 = gmsh.model.geo.addLine(p2, p3)
    l3 = gmsh.model.geo.addLine(p3, p4)
    l4 = gmsh.model.geo.addLine(p4, p1)

    loop = gmsh.model.geo.addCurveLoop([l1, l2, l3, l4])
    _ = gmsh.model.geo.addPlaneSurface([loop])

    # 2. Mesh Generation
    gmsh.model.geo.synchronize()
    gmsh.model.mesh.generate(2)
    gmsh.write(filename)
    gmsh.finalize()

    # 3. Read back with meshio
    m = meshio.read(filename)
    if os.path.exists(filename):
        os.remove(filename)

    points = m.points[:, :2]  # Drop z-coordinate for 2D
    triangles = m.cells_dict["triangle"]

    return Mesh(points, triangles)


def generate_honeycomb_mesh(
    start_x: float,
    start_y: float,
    n_x: int,
    n_y: int,
    side_length: float,
    segments_per_side: int = 1,
) -> Mesh:
    """
    Generates a 2D honeycomb mesh (hexagonal grid) with specified parameters.
    Each hexagon side can be subdivided into smaller segments.

    Args:
        start_x (float): Starting x-coordinate of the honeycomb grid.
        start_y (float): Starting y-coordinate of the honeycomb grid.
        n_x (int): Number of hexagons along the x-direction.
        n_y (int): Number of hexagons along the y-direction.
        side_length (float): Length of each side of the hexagon.
        segments_per_side (int): Number of subdivisions per hexagon side.

    Returns:
        Mesh: The generated honeycomb mesh.
    """

    dx = np.sqrt(3) * side_length
    dy = 1.5 * side_length
    row_offset = (np.sqrt(3) * side_length) / 2.0

    node_map = {}
    coords_list = []
    lines_list = []
    edge_set = set()

    def get_or_create_node(x, y):
        key = (round(x, 6), round(y, 6))
        if key not in node_map:
            idx = len(coords_list)
            coords_list.append([x, y])
            node_map[key] = idx
            return idx
        return node_map[key]

    angles = np.deg2rad([30, 90, 150, 210, 270, 330])

    for row in range(n_y):
        cols_in_this_row = n_x if (row % 2 == 0) else (n_x - 1)
        current_offset = 0.0 if (row % 2 == 0) else row_offset

        for col in range(cols_in_this_row):
            cx = start_x + (col * dx) + current_offset
            cy = start_y + (row * dy)

            corners = []
            for theta in angles:
                vx = cx + side_length * np.cos(theta)
                vy = cy + side_length * np.sin(theta)
                corners.append((vx, vy))

            for k in range(6):
                # Start and End coordinates of the current side
                start_pt = np.array(corners[k])
                end_pt = np.array(corners[(k + 1) % 6])

                # Get the Node Index for the start of the side
                current_node_idx = get_or_create_node(start_pt[0], start_pt[1])

                # Vector along the side
                side_vector = end_pt - start_pt

                for i in range(1, segments_per_side + 1):
                    # Calculate fraction of distance (e.g., 1/3, 2/3, 3/3)
                    t = i / segments_per_side

                    # Calculate next coordinate
                    next_pt = start_pt + t * side_vector

                    # Get index for this new point
                    next_node_idx = get_or_create_node(next_pt[0], next_pt[1])

                    # Create the small segment
                    edge_key = tuple(sorted((current_node_idx, next_node_idx)))
                    if edge_key not in edge_set:
                        edge_set.add(edge_key)
                        lines_list.append([current_node_idx, next_node_idx])

                    # Move forward
                    current_node_idx = next_node_idx

    return Mesh(np.array(coords_list), np.array(lines_list))

plt.figure(figsize=(3, 3))
ax = plt.gca()

ax.tripcolor(
    plate_mesh.coords[:, 0],
    plate_mesh.coords[:, 1],
    plate_mesh.elements,
    facecolors=jnp.ones(plate_mesh.elements.shape[0]),
    edgecolors="k",
    cmap="managua",
    lw=0.2,
)

for i, el in enumerate(fiber_mesh.elements):
    p0 = fiber_mesh.coords[el[0]]
    p1 = fiber_mesh.coords[el[1]]
    ax.plot(
        [p0[0], p1[0]],
        [p0[1], p1[1]],
        "k-",
        lw=1,
    )

ax.set_xlabel("x")
ax.set_ylabel("y")
ax.axis("equal")
ax.margins(0.0, 0.0)
plt.show()

class SpringElement(element.Line2):
    """
    A 2-node linear element embedded in 3D space.
    Reference domain: [-1, 1]
    """

    def shape_function(self, xi: Array) -> Array:
        return jnp.array([0.5 * (1.0 - xi[0]), 0.5 * (1.0 + xi[0])])

    def shape_function_derivative(self, xi: Array) -> Array:
        return jnp.array([[-0.5, 0.5]])

    def get_jacobian(self, xi: Array, nodal_coords: Array) -> tuple[Array, Array]:
        """
        nodal_coords: (2, 3) -> Two nodes in 3D space
        """
        dN_dxi = self.shape_function_derivative(xi)

        J_vec = dN_dxi @ nodal_coords  # (1, 2) @ (2, 3) -> (1, 3)

        detJ = jnp.linalg.norm(J_vec)
        return J_vec, detJ

    def gradient(self, xi: Array, nodal_values: Array, nodal_coords: Array) -> Array:
        """
        Returns the 3D Gradient vector.
        """
        J_vec, detJ = self.get_jacobian(xi, nodal_coords)
        dN_dxi = self.shape_function_derivative(xi)
        du_dxi = dN_dxi @ nodal_values

        du_ds = du_dxi / detJ

        tangent = J_vec / detJ
        grad_u_3d = jnp.vdot(du_ds, tangent)

        return grad_u_3d

def newton_krylov_solver(
    u,
    fext,
    lifter_plate: lifter.Lifter,
):
    fint = gradient(u, lifter_plate)
    iiter = 0
    norm_res = 1.0
    tol = 1e-8
    max_iter = 80
    while norm_res > tol and iiter < max_iter:
        residual = fext - fint
        A = hessian_fn(u, lifter_plate)
        du = jax.experimental.sparse.linalg.spsolve(
            A.data, A.indices, A.indptr, residual
        )
        u = u.at[:].add(du)

        fint = gradient(u, lifter_plate)
        residual = fext - fint
        norm_res = jnp.linalg.norm(residual)

        print(f"  Residual: {norm_res:.2e}")
        iiter += 1
    return u, norm_res

import pyvista as pv


def set_size(fraction=1, height_ratio="golden", width="two-column", subplots=(1, 1)):
    if width == "two-column":
        width_pt = 180  # mm
    elif width == "one-column":
        width_pt = 90  # mm
    else:
        width_pt = width

    if height_ratio == "golden":
        ratio_pt = (np.sqrt(5) - 1.0) / 2.0
    else:
        ratio_pt = height_ratio

    fig_width_pt = width_pt * fraction
    inches_per_pt = 1.0 / 25.4
    fig_width_in = fig_width_pt * inches_per_pt
    fig_height_in = fig_width_in * ratio_pt * (subplots[0] / subplots[1])
    fig_dim = (fig_width_in, fig_height_in)

    return fig_dim


def get_pv_grid(mesh: Mesh) -> pv.UnstructuredGrid:
    """Convert Tatva mesh to PyVista UnstructuredGrid."""
    if mesh.coords.shape[1] == 2:
        pv_points = np.hstack((mesh.coords, np.zeros(shape=(mesh.coords.shape[0], 1))))
    else:
        pv_points = np.array(mesh.coords)
    cells = np.hstack(
        [
            np.full((mesh.elements.shape[0], 1), 3, dtype=np.int64),
            np.array(mesh.elements, dtype=np.int64),
        ]
    )
    grid = pv.UnstructuredGrid(
        cells, np.full(mesh.elements.shape[0], pv.CellType.TRIANGLE), pv_points
    )
    return grid


def find_domain_boundary(elements):
    # 1. Get all edges (3 per triangle)
    edges = np.concatenate(
        [elements[:, [0, 1]], elements[:, [1, 2]], elements[:, [2, 0]]], axis=0
    )

    # 2. Sort node IDs within each edge to handle directionality
    edges_sorted = np.sort(edges, axis=1)

    # 3. Find unique edges and their counts
    unique_edges, indices, counts = np.unique(
        edges_sorted, axis=0, return_index=True, return_counts=True
    )

    # 4. Boundary edges are those that appear only once
    boundary_edges = edges[indices[counts == 1]]

    return boundary_edges


def get_pv_line_grid(mesh: Mesh) -> pv.PolyData:
    """Convert a Line2 (segment) Tatva mesh to a PyVista PolyData of lines."""
    coords = np.asarray(mesh.coords)
    if coords.shape[1] == 2:
        pv_points = np.hstack((coords, np.zeros((coords.shape[0], 1))))
    else:
        pv_points = coords
    elems = np.asarray(mesh.elements, dtype=np.int64)
    # VTK line connectivity: [n_points_in_cell, id0, id1, ...] per cell
    lines = np.hstack([np.full((elems.shape[0], 1), 2, dtype=np.int64), elems])
    return pv.PolyData(pv_points, lines=lines)


def _pad3(u: Array) -> np.ndarray:
    """Pad a 2D nodal (n, 2) field to (n, 3) so it can warp PyVista points."""
    u = np.asarray(u)
    return np.hstack([u, np.zeros((u.shape[0], 1))])


def _neo_hookean_energy(F: Array) -> Array:
    """Scalar strain-energy density for a single 2x2 deformation gradient."""
    C = F.T @ F
    I1 = jnp.trace(C)
    J = jnp.linalg.det(F)
    return (
        0.5 * mat.mu * (I1 - 2 - 2 * jnp.log(J)) + (mat.lmbda / 2) * (jnp.log(J)) ** 2
    )


@autovmap(F=2)
def von_mises_from_F(F: Array) -> Array:
    """In-plane von Mises stress from the Cauchy stress sigma = J^-1 P F^T."""
    P = jax.grad(_neo_hookean_energy)(F)  # 1st Piola-Kirchhoff (2, 2)
    J = jnp.linalg.det(F)
    sigma = (P @ F.T) / J  # Cauchy stress (2, 2)
    sxx, syy, sxy = sigma[0, 0], sigma[1, 1], sigma[0, 1]
    return jnp.sqrt(sxx**2 - sxx * syy + syy**2 + 3 * sxy**2)


def plate_von_mises(u2d: Array) -> np.ndarray:
    """Per-element (per quadrature point) von Mises stress for a nodal field."""
    grad_u = op_plate.grad(u2d)
    F = compute_deformation_gradient(grad_u)
    return np.asarray(von_mises_from_F(F)).reshape(-1)


def animate_deformation(
    save_gif: str | None = None,
    plate_cmap: str = "viridis",
    fiber_cmap: str = "magma",
):
    """Step through ``u_sol_per_step``, warping the plate and fibers.

    With ``save_gif`` a GIF is written; otherwise an interactive window opens
    with a slider to scrub through the load steps. Colour limits are fixed to
    the global range so steps are directly comparable.
    """
    plate_grid = get_pv_grid(plate_mesh)
    fiber_grid = get_pv_line_grid(fiber_mesh)
    plate_pts0 = plate_grid.points.copy()
    fiber_pts0 = fiber_grid.points.copy()

    # lift_from_zeros returns a flat DOF vector; reshape to (n_nodes, 2).
    u2d_steps = [
        jnp.asarray(u).reshape(plate_mesh.coords.shape[0], 2) for u in u_sol_per_step
    ]

    # Precompute per-step warp fields and scalars.
    plate_disp = [_pad3(u) for u in u2d_steps]
    plate_vm = [plate_von_mises(u) for u in u2d_steps]  # per-element von Mises
    fiber_disp = [_pad3(interpolater(u)) for u in u2d_steps]
    fiber_strain = [
        np.asarray(op_line.grad(interpolater(u)).flatten()) for u in u2d_steps
    ]

    plate_clim = [0.0, float(max(s.max() for s in plate_vm))]
    fiber_clim = [
        float(min(s.min() for s in fiber_strain)),
        float(max(s.max() for s in fiber_strain)),
    ]

    def set_step(i):
        i = int(round(i))
        plate_grid.points = plate_pts0 + plate_disp[i]
        plate_grid["von Mises"] = plate_vm[i]
        fiber_grid.points = fiber_pts0 + fiber_disp[i]
        fiber_grid["strain"] = fiber_strain[i]

    pl = pv.Plotter(off_screen=save_gif is not None)
    set_step(0)
    pl.add_mesh(
        plate_grid,
        scalars="von Mises",
        cmap=plate_cmap,
        clim=plate_clim,
        show_edges=False,
        edge_color="gray",
        line_width=0.0,
        scalar_bar_args={
            "title": "von Mises",
            "vertical": False,
            "position_x": 0.05,
            "position_y": 0.05,
            "width": 0.4,
            "height": 0.05,
        },
    )
    pl.add_mesh(
        fiber_grid,
        scalars="strain",
        cmap=fiber_cmap,
        clim=fiber_clim,
        line_width=6,
        render_lines_as_tubes=True,
        scalar_bar_args={
            "title": "fiber strain",
            "vertical": False,
            "position_x": 0.55,
            "position_y": 0.05,
            "width": 0.4,
            "height": 0.05,
        },
    )
    pl.view_xy()
    pl.enable_parallel_projection()

    n = len(u_sol_per_step)
    if save_gif is not None:
        pl.open_gif(save_gif)
        for i in range(n):
            set_step(i)
            pl.write_frame()
        pl.close()
    else:
        pl.add_slider_widget(
            set_step,
            [0, n - 1],
            value=0,
            title="load step",
            fmt="%.0f",
        )
        pl.show()
    return pl