2D Timoshenko Beam¤

import jax

jax.config.update("jax_enable_x64", True)  # use double-precision

import time
from functools import partial
from typing import NamedTuple

import equinox as eqx
import jax.experimental.sparse as jsp
import jax.numpy as jnp
import numpy as np
from jax import Array
from jax_autovmap import autovmap
from tatva import Mesh, Operator, element, sparse

L = 1  # Total length of the beam
nb_elem = 20

mesh = create_beam_structure(nb_elem=nb_elem, length=L)
n_nodes = mesh.coords.shape[0]
nb_nodes_per_cell = mesh.elements.shape[1]
nb_dofs_per_node = 3
n_dofs = n_nodes * nb_dofs_per_node

class Line2In3D(element.Element):
    """
    A 2-node linear element embedded in 2D/3D space.
    Should be used for Timoshenko beam formulations.
    """


    def _default_quadrature(self):
        # 1-point Gauss quadrature for line element
        quad_points = jnp.array([0.0])  # Natural coordinate at the center
        quad_weights = jnp.array([2.0])  # Weight for the quadrature point
        return quad_points, quad_weights


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

    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]:
        """
        Assumes nodal_coords is of shape (2, 3) for the two nodes in 3D space.
        Args:
            xi: Natural coordinate (-1 to 1)
            nodal_coords: Array of shape (2, 3) containing the coordinates of the two nodes.

        Returns:
            J_vec: The tangent vector (unnormalized) dx/dxi
            detJ: The scaling factor (ds/dxi) = norm(J_vec)
        """
        dN_dxi = self.shape_function_derivative(xi)
        J_vec = dN_dxi @ nodal_coords  # Shape: (3,) or (2, )
        detJ = jnp.linalg.norm(J_vec)
        return J_vec, detJ

    def gradient(self, xi: Array, nodal_values: Array, nodal_coords: Array) -> Array:
        """
        Returns the derivative of the field w.r.t physical arc length s.

        If nodal_values is (2, 3) [Displacement], returns (3,) vector [du/ds, dv/ds, dw/ds]
        If nodal_values is (2, 1) [Scalar], returns (1,) scalar [dT/ds]

        Args:
            xi: Natural coordinate (-1 to 1)
            nodal_values: Array of shape (2, dofs_per_node) containing the values
            nodal_coords: Array of shape (2, 3) containing the coordinates of the two nodes.
        Returns:
            Array of shape (dofs_per_node,) containing the gradient of the field w.r.t physical arc length s.
        """
        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

        return du_ds

@autovmap(du_ds=1, dtheta_ds=1, theta=1, tangent=1, E=0, A=0, I=0, G=0, J=0)
def timoshenko_density(du_ds, dtheta_ds, theta, tangent, E, A, I, G, J):
    """
    Computes strain energy density at a single integration point.

    Args:
        du_ds: (3,) Derivative of displacement w.r.t arc length s
        dtheta_ds: (3,) Derivative of rotation w.r.t arc length s
        theta: (3,) Rotation vector value
        tangent: (3,) Unit tangent vector of the beam
        E, A, I, G, J: Scalar stiffness properties

    Returns:
        Scalar energy density
    """

    epsilon = jnp.dot(du_ds, tangent)

    normal = jnp.array([-tangent[1], tangent[0]])

    geometric_slope = jnp.dot(du_ds, normal)
    gamma = geometric_slope - theta

    w_axial = 0.5 * E * A * epsilon**2
    w_shear = 0.5 * G * kappa * A * jnp.dot(gamma, gamma)
    w_bend = 0.5 * E * I * jnp.dot(dtheta_ds, dtheta_ds)

    return w_axial + w_shear + w_bend

# material and geometrical parameters
E = 500
h = 1
A = h**2
I = h * h**3 / 12

G = E / 10.
kappa = 5 / 6  # shear correction coefficient (for rectangular cross section )


class Material(NamedTuple):
    E: float
    A: float
    I: float
    G: float
    kappa: float


mat = Material(E=E, A=A, I=I, G=G, kappa=kappa)

x_max = jnp.max(mesh.coords[:, 0])
x_min = jnp.min(mesh.coords[:, 0])
y_max = jnp.max(mesh.coords[:, 1])
y_min = jnp.min(mesh.coords[:, 1])

left_nodes = jnp.where(jnp.isclose(mesh.coords[:, 0], x_min, atol=1e-6))[0]

right_nodes = jnp.where(jnp.isclose(mesh.coords[:, 0], x_max, atol=1e-6))[0]


fixed_dofs = jnp.concatenate(
    [
        left_nodes * nb_dofs_per_node,
        left_nodes * nb_dofs_per_node + 1,
        left_nodes * nb_dofs_per_node + 2,
    ]
)

applied_dofs = jnp.array(right_nodes * nb_dofs_per_node + 1)

free_dofs = jnp.setdiff1d(jnp.arange(n_dofs), fixed_dofs)

beam_element = Line2In3D()
op = Operator(mesh=mesh, element=beam_element)

J_vecs, detJs = op.map(beam_element.get_jacobian)(mesh.coords)

tangents = J_vecs / detJs[..., None]

@jax.jit
def _total_energy_timoshenko(z_flat: Array, tangents: Array) -> Array:
    """
    Computes total Timoshenko beam energy.
    Args:
        z_flat: (n_nodes * 3,) Flattened displacement and rotation vector (ux1, uy1, theta1, ...)
        tangents: (n_elements, 3) Unit tangent vectors for each element
    Returns:
        Scalar total energy of the beam
    """

    z = z_flat.reshape(-1, 3)  # ux, uy, theta

    u = z[:, [0, 1]]  # ux, uy
    theta = z[:, [2]]  # theta

    du_ds = op.grad(u)
    dtheta_ds = op.grad(theta)

    theta_quad = op.eval(theta)

    density = timoshenko_density(
        du_ds, dtheta_ds, theta_quad, tangents, mat.E, mat.A, mat.I, mat.G, mat.kappa
    )

    return op.integrate(density)

total_energy_timoshenko = eqx.Partial(_total_energy_timoshenko, tangents=tangents)

@eqx.filter_jit
def total_energy_timoshenko_free(u_free, applied_disp):
    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_timoshenko(u_full)


compute_internal_free = jax.jacrev(total_energy_timoshenko_free)

sparsity_pattern_csr = sparse.create_sparsity_pattern(
    mesh, n_dofs_per_node=nb_dofs_per_node
)
reduced_sparsity_pattern_csr = sparse.reduce_sparsity_pattern(sparsity_pattern_csr, free_dofs)
colored_matrix = sparse.ColoredMatrix.from_csr(reduced_sparsity_pattern_csr)

def fn(u, fext, applied_disp):
    fint = compute_internal_free(u, applied_disp)
    return fext - fint

u = jnp.zeros(n_dofs)
fext = jnp.zeros(n_dofs)

nsteps = 20
P = -1.0
applied_loading = jnp.linspace(0, P, num=nsteps)

dumped_nodes = jnp.zeros((nsteps, n_nodes * (nb_dofs_per_node - 1)))

dumped_loads = jnp.zeros(nsteps)

errors = jnp.zeros(nsteps)

u_free = u.at[free_dofs].get()

fn_partial = jax.jit(
    partial(fn, fext=fext.at[free_dofs].get(), applied_disp=jnp.zeros(len(fixed_dofs)))
)
start_time = time.time()
hessian_fn = sparse.jacfwd(
    fn=fn,
    colored_matrix=colored_matrix,
    color_batch_size=len(jnp.unique(colored_matrix.colors)) + 1,
)

print(f"Time to compute Hessian: {time.time() - start_time:.2f} seconds")


max_iter = 20
for step, load in enumerate(applied_loading):
    fext = fext.at[applied_dofs].set(load)
    fext_free = fext.at[free_dofs].get()
    fn_partial = eqx.Partial(
        fn, fext=fext_free, applied_disp=jnp.zeros(len(fixed_dofs))
    )

    hessian_partial = eqx.Partial(
        hessian_fn, fext=fext_free, applied_disp=jnp.zeros(len(fixed_dofs))
    )

    residual = fn_partial(u_free)
    K_sparse = hessian_partial(u_free)

    du = jnp.zeros(len(free_dofs))
    rnorm = jnp.linalg.norm(residual)
    iiter = 0

    while rnorm > 1e-8 and iiter < max_iter:
        du = du + jsp.linalg.spsolve(
            K_sparse.data,
            indices=K_sparse.indices,
            indptr=K_sparse.indptr,
            b=-residual,
        )

        u_curr = u_free + du
        residual = fn_partial(u_curr)

        K_sparse = hessian_partial(u_curr)

        rnorm = jnp.linalg.norm(residual)

        iiter += 1

    u_free = u_free + du
    u = u.at[free_dofs].set(u_free)

    new_nodes = mesh.coords + u.reshape((n_nodes, nb_dofs_per_node))[:, :2]
    dumped_nodes = dumped_nodes.at[step].set(new_nodes.reshape(-1))
    dumped_loads = dumped_loads.at[step].set(load)
    errors = errors.at[step].set(rnorm)

    print(load, rnorm)