Peeling of a dimpled adhesive beam via Lagrange Multipliers¤

from dataclasses import replace
from typing import Literal, NamedTuple, Callable

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import scipy.sparse as sps
from jax import Array
from jax_autovmap import autovmap

from tatva import Mesh, Operator, compound, sparse
from tatva.element.base import Line2, Tri3
from tatva.lifter import Constraint, Fixed, Lifter, RuntimeValue

jax.config.update("jax_enable_x64", True)


def jit[**P, T](f: Callable[P, T], *args, **kwargs) -> Callable[P, T]:
    return jax.jit(f, *args, **kwargs)


def jacrev[**P, T](f: Callable[P, T], *args, **kwargs) -> Callable[P, T]:
    return jax.jacrev(f, *args, **kwargs)

class Material(NamedTuple):
    """Material properties for the elasticity operator."""

    mu: float  # Diffusion coefficient
    lmbda: float  # Diffusion coefficient

    @classmethod
    def from_youngs_poisson_2d(
        cls, E: float, nu: float, plane_stress: bool = False
    ) -> "Material":
        mu = E / 2 / (1 + nu)
        if plane_stress:
            lmbda = 2 * nu * mu / (1 - nu)
        else:
            lmbda = E * nu / (1 - 2 * nu) / (1 + nu)
        return cls(mu=mu, lmbda=lmbda)


@autovmap(grad_u=2)
def compute_deformation_gradient(grad_u: Array) -> Array:
    return jnp.eye(2) + grad_u


@autovmap(grad_u=2, mat=None)
def strain_energy_density(grad_u: Array, mat: Material) -> Array:
    F = compute_deformation_gradient(grad_u)
    C = F.T @ F
    J = jnp.linalg.det(F)
    return (
        mat.mu / 2 * (jnp.trace(C) - 2)  # 2D case
        - mat.mu * jnp.log(J)
        + (mat.lmbda / 2) * (jnp.log(J)) ** 2
    )


@autovmap(grad_u=2, mat=None)
def get_cauchy_stress(grad_u: Array, mat: Material) -> Array:
    F = compute_deformation_gradient(grad_u)
    C = F.T @ F
    J = jnp.linalg.det(F)

    C_inv = jnp.linalg.inv(C)
    S = mat.mu * (jnp.eye(2) - C_inv) + mat.lmbda * jnp.log(J) * C_inv  # 2nd PK
    P = F @ S  # 1st PK

    sigma = (P @ F.T) / J  # Cauchy
    return sigma


def von_mises_stress(sig):
    s_xx, s_yy, s_xy = sig[..., 0, 0], sig[..., 1, 1], sig[..., 0, 1]
    return np.sqrt(s_xx**2 - s_xx * s_yy + s_yy**2 + 3 * s_xy**2)

mesh = Mesh(jnp.asarray(nodes[:, :2]), jnp.asarray(elements))

bottom_nodes = jnp.where((mesh.coords[:, 1] == 0) & (mesh.coords[:, 0] > 1))[0]
# bottom_nodes = jnp.where(mesh.coords[:, 1] == 0)[0]
load_nodes = jnp.where(mesh.coords[:, 0] == 0)[0]
# we fix the bottom right corner to prevent rigid body motion when fully cracked
right_bottom = jnp.where((mesh.coords[:, 0] == l) & (mesh.coords[:, 1] == 0))[0]

# the line elements are those that have at least 2 nodes on the bottom boundary, and all
# of their nodes have y=0
line_element_indices = jnp.where(
    sum(
        jnp.isin(mesh.elements[:, dim], bottom_nodes)
        for dim in range(mesh.elements.shape[1])
    )
    >= 2
)[0]
line_elements = mesh.elements[line_element_indices]
mask_y_is_zero = mesh.coords[line_elements][..., 1] == 0
line_elements = line_elements[mask_y_is_zero].reshape(-1, 2)

op = Operator(mesh, Tri3())
op_line = Operator(replace(mesh, elements=line_elements), Line2())
mat = Material.from_youngs_poisson_2d(1, 0.3)


class Solution(compound.Compound):
    u_rigid = compound.field((3,))
    """x, y displacement of the rigid body"""
    u = compound.field((mesh.coords.shape))
    """Displacement field (u_x, u_y) at each node"""
    lm = compound.field((op_line.mesh.elements.shape[0],))  # DG space
    """Lagrange-multipliers for y-displacement at bottom nodes"""


Solution.size

class RigidBody(Constraint):
    """Constraint to enforce rigid body motion on a set of nodes."""

    def __init__(self, rigid_nodes: Array, x_ref: Array):
        self.dofs = Solution.u[rigid_nodes, :].flatten()
        self.rigid_nodes = rigid_nodes
        self.x_ref = x_ref

    def apply_lift(self, u_full: Array) -> Array:
        sol = Solution(u_full)
        u, u_rigid = sol.u, sol.u_rigid
        u_after_constraint = self.constraint_function(u_rigid, u)

        return sol.at("u").set(u_after_constraint).at("u_rigid").set(u_rigid).flatten()

    def constraint_function(self, u_rigid: Array, u: Array) -> Array:
        """Takes u_rigid and computes the constraint u_flat
        fn (3,) -> (self.dofs.size,)
        """
        vecs = mesh.coords[self.rigid_nodes] - self.x_ref
        ux, uy, theta = u_rigid
        cos, sin = jnp.cos(theta), jnp.sin(theta)
        rotation_matrix = jnp.array([[cos, -sin], [sin, cos]])
        rotated_vecs = vecs @ rotation_matrix.T
        u_rotation = rotated_vecs - vecs
        u_rigid = jnp.stack([ux, uy]) + u_rotation
        return u.at[self.rigid_nodes, :].set(u_rigid)


rigid_body_constraint = RigidBody(load_nodes, jnp.array([0.0, h / 2]))
lifter = Lifter(
    Solution.size,
    Fixed(Solution.u[right_bottom, 0]),
    Fixed(Solution.u_rigid[1], RuntimeValue("disp", default=0.0)),
    rigid_body_constraint,
)

def equality_constraint_potential(lm: Array, u: Array, active_set: Array) -> Array:
    density = jnp.where(
        active_set[..., None],
        lm[..., None] * op_line.eval(u)[..., 1],
        0.5 * lm[..., None] ** 2,
    )
    return op_line.integrate(density)


def total_lagrangian(z: Array, active_set: Array, mat: Material = mat) -> Array:
    u_rigid, u, lm = Solution(z)

    # strain energy contribution
    psi = strain_energy_density(op.grad(u), mat)
    Psi = op.integrate(psi)

    # constraint contribution
    G = equality_constraint_potential(lm, u, active_set)
    return Psi + G


def total_lagrangian_lifted(
    z_reduced: Array, active_set: Array, lifter: Lifter, mat: Material = mat
) -> Array:
    z_full = lifter.lift_from_zeros(z_reduced)
    return total_lagrangian(z_full, active_set, mat)

u_rigid, u, lm = Solution()
active_set = jnp.ones_like(lm, dtype=bool)  # all constraints active for now

# We use jacrev to get the gradient w.r.t. lm, and jacfwd to get the Jacobian w.r.t. u.
g_fn = jax.jacrev(lambda lm, u: equality_constraint_potential(lm, u, active_set))
# B_fn will give us the coupling Jacobian (shape=(num_constraints, num_dofs))
# need to flatten u for jacfwd and then reshape inside the lambda
B_fn = jax.jacfwd(
    lambda lm, u: g_fn(lm, u.reshape(Solution.u.shape)).flatten(), argnums=1
)
B = sps.csr_matrix(B_fn(lm, u.flatten()))

# these are the block-diagonal stiffness matrices for the rigid body, the displacement field, and the Lagrange multipliers
K_rigid = sps.csr_matrix(jnp.ones((3, 3)))
K_u = sparse.create_sparsity_pattern(mesh, 2)
K_lm = sps.eye(Solution.lm.size, format="csr")

# create a 'Mask' matrix that represents which nodes the rigid body 'touches'
# Shape: (Total_DOFs, 3)
B_pattern = sps.lil_matrix((Solution.u.size, 3))
B_pattern[rigid_body_constraint.rigid_nodes * 2] = 1
B_pattern[rigid_body_constraint.rigid_nodes * 2 + 1] = 1
B_pattern = B_pattern.tocsr()

# The coupling actually extends to anyone connected to a rigid node.
# We find this by multiplying the mesh connectivity by our rigid mask.
# (Mesh_Sparsity) @ (Rigid_Mask) gives the "influence zone"
extended_coupling_pattern = (K_u != 0) @ B_pattern
B_rigid = extended_coupling_pattern.tocsr()

block_matrix = [
    [K_rigid, B_rigid.T, None],
    [B_rigid, K_u, B.T],
    [None, B, K_lm],
]

sparsity_pattern = sps.bmat(block_matrix, format="csr")
sparsity_pattern = sparse.reduce_sparsity_pattern(sparsity_pattern, lifter.free_dofs)

# finally, color the matrix for sparse differentiation
colored_matrix = sparse.ColoredMatrix.from_csr(sparsity_pattern)

fig, ax = plt.subplots(figsize=(10, 8))
ax.spy(sparsity_pattern, markersize=0.3, color="k")

residual_fn = jit(jacrev(total_lagrangian_lifted))
residual_fn_full = jit(jacrev(total_lagrangian))

jacobian_fn = jit(sparse.jacfwd(residual_fn, colored_matrix=colored_matrix))

from bamboost import Collection

coll = Collection("data")
coll = coll.filter(tags="adhesion")
sim = coll.add("adhesion_example", tags={"adhesion", "with_sig"})
sim.meshes.add(mesh.coords, mesh.elements)

from typing import TypeAlias

from jax.experimental import checkify
from soldis import linear, newton
from soldis.newton import SolverState
from tatva.utils import make_project_function

Carry: TypeAlias = tuple[int, Array, SolverState[Array], Array]

solver = newton.NewtonSolver(
    residual_fn,
    linear.SparseTatva(),
    jacobian_fn,
    options=newton.NewtonSolverOptions(maxiter=50, tol=1e-10, verbose=False),
)

# initial guesses
y0 = jnp.zeros(lifter.size_reduced)
active_set = jnp.ones(Solution.lm.size, dtype=bool)  # all constraints active for now

# adhesive strength
strength = 1e-1

solver.root = jit(solver.root)
initial_state = solver.init(y0, active_set, lifter, mat)

project = make_project_function(mesh.coords.shape[0], elements=mesh.elements)


@jit
@checkify.checkify
def solve_with_active_set(
    state: SolverState[Array],
    active_set: Array,
    lifter: Lifter,
    mat: Material,
    strength: float,
) -> tuple[Array, SolverState[Array]]:

    carry: Carry = (0, active_set, state, jnp.asarray(False))

    def cond_fun(carry: Carry) -> Array:
        *_, converged = carry
        # continue as long as NOT converged
        return jnp.logical_not(converged)

    def body_fun(carry: Carry) -> Carry:
        i, active_set, _state, converged = carry

        jax.debug.print("    Active set loop iteration {}", i)

        sol = solver.root(_state.value, active_set, lifter, mat)
        checkify.check(
            sol.residual_norm < 1e-5,
            "Solver did not converge within the body of the loop",
        )
        lm = Solution(lifter.lift_from_zeros(sol.value)).lm
        new_active_set = active_set * (lm <= strength)
        # max_lm = jnp.argmax(lm)
        # new_active_set = active_set.at[max_lm].set(lm[max_lm] <= strength)

        # Check for convergence: has the active set stopped changing?
        converged = jnp.array_equal(new_active_set, active_set)

        return (i + 1, new_active_set, sol, converged)

    final_i, final_active_set, final_state, _ = jax.lax.while_loop(
        cond_fun, body_fun, carry
    )

    return final_active_set, final_state


d_previous = 0.0
state = initial_state
_lifter = lifter


def step(
    i: int,
    d: float,
    d_previous: float,
    state: SolverState[Array],
    active_set: Array,
    _lifter: Lifter,
):
    active_set_backup = active_set
    state_backup = state

    print(f"--> Solving for displacement: {d:.2f}")
    _lifter = _lifter.at("disp").set(d)
    err, (active_set, state) = solve_with_active_set(
        state, active_set, _lifter, mat, strength
    )
    # err.throw()

    if err.get() is not None:
        # solve at an intermediate displacement to get a better solution
        _lifter = _lifter.at("disp").set(d_previous + (d - d_previous) / 2)
        err, (active_set, state) = solve_with_active_set(
            state_backup, active_set_backup, lifter, mat, strength
        )
        err.throw()
        # solve again at the original displacement to get the correct active set for post-processing
        _lifter = _lifter.at("disp").set(d)
        err, (active_set, state) = solve_with_active_set(
            state_backup, active_set_backup, lifter, mat, strength
        )

    print(f"    Residual after solve: {state.residual_norm:.2e} [{state.iteration}]")
    d_previous = d

    # compute the reaction force at the loaded nodes using the full residual function
    _lifter = _lifter.at("disp").set(d)
    force_value = residual_fn_full(
        _lifter.lift_from_zeros(state.value), active_set, mat
    )[Solution.u[load_nodes, 1]].sum()

    # finally, write the results to the simulation
    u_rigid, u, lm = Solution(_lifter.lift_from_zeros(state.value))

    if i % 1 == 0:
        s = sim.data.require_step(d)
        s.add_field("u", u)
        s.add_field("u_rigid", u_rigid)
        s.add_field("lm", lm)
        s.add_field("active_set", active_set)
        sig = get_cauchy_stress(op.grad(u), mat)
        sig_vm = von_mises_stress(sig)
        # sig_vm_nodal = project(op, sig_vm)
        s.add_field("sig_vm", sig_vm)
        s.add_scalar("force", force_value)
    return d_previous, state, active_set, _lifter

for i, d in enumerate(jnp.arange(0, 14, 0.05)):
    d_previous, state, active_set, _lifter = step(
        i, d, d_previous, state, active_set, _lifter
    )

# After the solver fails, we continue solving with a smaller increment
# We do this manually here, because we can do that flexibly in a notebook setting. Would have to be automated in a script.
# Continue from previous solution
last_step = sim.data.last_step or 0
y_now = (
    Solution()
    .at("u")
    .set(sim.data["u"][last_step])
    .at("u_rigid")
    .set(sim.data["u_rigid"][last_step])
    .at("lm")
    .set(sim.data["lm"][last_step])
    .arr
)
y_now = lifter.reduce(y_now)
state = solver.init(y_now, sim.data["active_set"][last_step], lifter, mat)
d_previous = sim.data.values[last_step]
active_set = sim.data["active_set"][last_step]

increment = 0.02
while True:
    d_previous, state, active_set, _lifter = step(
        0, d_previous + increment, d_previous, state, active_set, _lifter
    )
    if d_previous > 13.5:
        break

from matplotlib.tri import Triangulation

i = -5
print(sim.data.values[i])
u, lm, active_set = (
    sim.data["u"][i],
    sim.data["lm"][i],
    sim.data["active_set"][i],
)
sig_vm = sim.data["sig_vm"][i].squeeze()

tri = Triangulation(mesh.coords[:, 0], mesh.coords[:, 1], mesh.elements)
x, y = mesh.coords[:, 0] + u[:, 0], mesh.coords[:, 1] + u[:, 1]

fig, ax = plt.subplots(figsize=(16, 4))
ax.set_aspect("equal")
ax.tripcolor(x, y, mesh.elements, sig_vm, cmap="managua_r")

fig, ax = plt.subplots(figsize=(6, 4))

ax.plot(sim.data.values, sim.data.globals["force"][:], "-")
ax.set(xlabel="Applied Displacement", ylabel="Reaction Force")

total_steps = len(sim.data.values)
indices = np.arange(total_steps)
n_frames = 150
equidistant_loading = np.linspace(0.0, sim.data.values[-1], n_frames)
sampled_float_indices = np.interp(equidistant_loading, sim.data.values, indices)
sampled_indices = np.unique(np.round(sampled_float_indices).astype(int))
sampled_indices = sampled_indices.tolist()