Matrix-Free approach with PETSc¤

from typing import NamedTuple

import gmsh
import jax

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

import jax.numpy as jnp
from jax import Array
from jax_autovmap import autovmap
from tatva import Mesh, Operator

from petsc4py import PETSc

lx = 1.0
ly = 1.0
mesh = generate_refined_plate_with_hole(
    lx, ly, hole_radius=0.2, mesh_size_fine=0.01, mesh_size_coarse=0.05
)

n_dofs_per_node = 2
n_dofs = mesh.coords.shape[0] * n_dofs_per_node

plot_mesh(mesh)

from tatva.element import Tri3

op = Operator(mesh, Tri3())


boundary_left = jnp.where(jnp.isclose(mesh.coords[:, 0], 0.0))[0]
boundary_right = jnp.where(jnp.isclose(mesh.coords[:, 0], lx))[0]
point_at_y_0 = jnp.where(
    jnp.isclose(mesh.coords[:, 0], lx) & jnp.isclose(mesh.coords[:, 1], 0.0)
)[0][0]
assert point_at_y_0

fixed_dofs = jnp.concatenate(
    [
        boundary_left * n_dofs_per_node,
    ]
)
free_dofs = jnp.setdiff1d(jnp.arange(n_dofs), fixed_dofs)

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

    mu: float  
    lmbda: float

    @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)


mat = Material.from_youngs_poisson_2d(1, 0.3)


@autovmap(grad_u=2)
def compute_strain(grad_u):
    return 0.5 * (grad_u + grad_u.T)


@autovmap(eps=2, mu=0, lmbda=0)
def compute_stress(eps, mu, lmbda):
    return 2 * mu * eps + lmbda * jnp.trace(eps) * jnp.eye(2)


@autovmap(grad_u=2, mu=0, lmbda=0)
def strain_energy_density(grad_u, mu, lmbda):
    eps = compute_strain(grad_u)
    sigma = compute_stress(eps, mu, lmbda)
    return 0.5 * jnp.einsum("ij,ij->", sigma, eps)

@jax.jit
def total_energy_full(u_flat: Array) -> Array:
    """Compute the total energy of the system."""
    u = u_flat.reshape(-1, 2)
    u_grad = op.grad(u)
    e_density = strain_energy_density(u_grad, mat.mu, mat.lmbda)
    return op.integrate(e_density)


@jax.jit
def total_energy(u_free: Array, applied_disp: Array) -> Array:
    """Compute the total energy of the system."""
    u_full = jnp.zeros(n_dofs).at[free_dofs].set(u_free)
    u_full = u_full.at[fixed_dofs].set(applied_disp)
    return total_energy_full(u_full)


compute_internal = jax.jacrev(total_energy)

@jax.jit
def compute_tangent(u: Array, v: Array, applied_disp: Array) -> Array:
    """Compute the tangent stiffness matrix."""
    tangent = jax.jvp(compute_internal, (u, applied_disp), (v, applied_disp))[1]
    return tangent

sig_loading = 1e-2

f_ext_0 = jnp.zeros(n_dofs)

idx_right = n_dofs_per_node * boundary_right

f_ext_0 = f_ext_0.at[idx_right].add(sig_loading)
f_ext = f_ext_0.at[free_dofs].get()

ctx = {}  # context for the Jacobian-vector product


class MatMultCtx:
    def __init__(self):
        self.x_cur = None # J(u) which is the last iterate state
        self.applied_disp = None

    def set_x(self, x: Array):
        """Set the current iterate for the Jacobian-vector product."""
        self.x_cur = x

    def set_applied_disp(self, applied_disp: Array):
        """Set the applied displacement for the Jacobian-vector product."""
        self.applied_disp = applied_disp

    def mult(self, A, V, Y):
        "Y = J(u) * V with current x stored in ctx['x']"
        v_np = V.getArray(readonly=True)
        jvp = compute_tangent(self.x_cur, v_np, self.applied_disp)
        Y.setArray(jvp)


ctx = MatMultCtx()
J = PETSc.Mat().createPython(
    [len(free_dofs), len(free_dofs)], comm=PETSc.COMM_SELF, context=ctx
)
J.setUp()

class ElasticitySNES:
    def __init__(self, applied_disp: Array, f_ext: Array):
        self.applied_disp = applied_disp
        self.f_ext = f_ext

    def residual(self, snes, u_petsc, r_petsc):
        """Compute the residual of the nonlinear system."""
        u = u_petsc.getArray(readonly=True)
        r = compute_internal(u, self.applied_disp)
        res = r - self.f_ext
        r_petsc.setArray(res)

    def jacobian(self, snes, u_petsc, J_mat: PETSc.Mat, P_mat: PETSc.Mat):
        """Compute the Jacobian of the nonlinear system."""
        u = u_petsc.getArray(readonly=True)
        ctx.set_x(u)
        ctx.set_applied_disp(self.applied_disp)
        return True

problem = ElasticitySNES(applied_disp=jnp.zeros(len(fixed_dofs)), f_ext=f_ext)


snes = PETSc.SNES().create(comm=PETSc.COMM_SELF)
x_sol = PETSc.Vec().createSeq(len(free_dofs))

snes.setFunction(problem.residual, x_sol)
snes.setJacobian(problem.jacobian, J, J)

snes.setType("newtonls")
snes.setTolerances(atol=1e-8, rtol=1e-10)
snes.setConvergenceHistory()
snes.setConvergedReason(reason=PETSc.SNES.ConvergedReason.CONVERGED_FNORM_ABS)

ksp = snes.getKSP()
ksp.setType("cg")  # Use Conjugate Gradient method
ksp.getPC().setType("none")  # No preconditioner
ksp.setConvergedReason(reason=PETSc.KSP.ConvergedReason.CONVERGED_ATOL)
ksp.setTolerances(rtol=1e-2, atol=1e-8)

convergence_history = np.zeros(100)  # Preallocate convergence history array

def monitor_fn(_snes, it, residual):
    convergence_history[it] = residual
    print(it, residual)


snes.setMonitor(monitor_fn)

du = x_sol.duplicate()
du.setArray(jnp.zeros(len(free_dofs)))  # Initial guess (zero displacement)
snes.solve(None, du)