In-class activity: Hertzian contact

In this exercise, we will solve numerically the problem of contact between a elastic cylinder and a rigid plane. This is a classical problem in contact, called Hertz contact problem, for which we know the analytical solution. In this problem, a elastic cylinder is pressed against a rigid plane. The elastic cylinder is infinitely long and has a circular cross-section. Since the cylinder is infinitely long, we can consider a cross-section of the cylinder and solve a 2D problem. The schematic of the problem is shown below.

An animation of the elastic cylinder in contact with the rigid plane is shown below. The stress value shown is \(\sigma_{yy}\).

From Fundamentals, we know that the constraint condition for the contact between two bodies is given by:

\[ g_n = (\boldsymbol{r} - \boldsymbol{\rho}) \cdot \boldsymbol{n} \geq 0, \]

where \(\boldsymbol{r}\) is the position vector of a point on the contact surface, \(\boldsymbol{\rho}\) is the position vector of the point on the target surface, and \(\boldsymbol{n}\) is the normal vector to the target surface.

For this problem, the target surface is the rigid plane and the contact surface is the surface of the elastic cylinder. The normal vector \(\boldsymbol{n}\) to the target surface is the \(+y\)-axis i.e. \(\boldsymbol{n} = (0, 1)\). Therefore, the constraint condition for the contact between the elastic cylinder and the rigid plane is given by:

\[ g_n = y - y_\text{rigid} \geq 0, \]

where

The \(y-\)coordinate of the contact surface of the elastic cylinder is regularly updated as the elastic cylinder deforms. If the initial \(y-\)coordinate of the contact surface is \(y_0\), then the \(y-\)coordinate of the contact surface at any time \(t\) is given by \(y_0 + u_y\), where \(u_y\) is the \(y-\)displacement of the contact surface. Thus, the constraint condition for the contact between the elastic cylinder and the rigid plane is given by:

\[ g_{i}(u) = {y}^{i}_0 + u_y^{i} - y_\text{rigid} \geq 0, \quad i \in \mathcal{A}, \]

Since the rigid body is infinitely stiff, the deformable cylinder will deform so as to make contact with the rigid plane. The constraints are therefore on the displacements of the nodes of the contact surface of the deformable body. In this exercise, we will use the penalty method to enforce the constraints. The penalty method adds a penalty term to the objective function which penalizes the violation of the constraints. The objective function is given by:

\[ \Psi(u) = \Psi_\text{elastic}(u) + \int_{\Gamma} k_\text{pen}\Big\langle g_i(u) \Big\rangle^2 d\Gamma \]

where \(\Psi_\text{elastic}(u)\) is the elastic energy of the system and \(k_\text{pen}\) is the penalty parameter.

Importing essential libraries 
import jax

jax.config.update("jax_enable_x64", True)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
jax.config.update("jax_platforms", "cpu")
import jax.numpy as jnp
from jax import Array

from tatva import Mesh, Operator, element, sparse
from tatva.plotting import (
    STYLE_PATH,
    colors,
    plot_element_values,
    plot_nodal_values,
)

from jax_autovmap import autovmap
from functools import partial

import numpy as np
import matplotlib.pyplot as plt

Model setup

Below, we generate a triangular mesh of a half-circle. We also define a function to get the elements and nodes on the potential contact surface.

Function to generate a half-circle mesh and get the contact elements 
import gmsh
import meshio



def get_elements_on_curve(mesh, curve_func, tol=1e-3):
    coords = mesh.coords
    elements_2d = mesh.elements
    # Efficiently find all nodes on the curve using jax.vmap
    on_curve_mask = jax.vmap(lambda c: curve_func(c, tol))(coords)

    elements_1d = []
    # Iterate through all 2D elements to find edges on the curve
    for tri in elements_2d:
        # Define the three edges of the triangle
        edges = [(tri[0], tri[1]), (tri[1], tri[2]), (tri[2], tri[0])]
        for n_a, n_b in edges:
            # If both nodes of an edge are on the curve, add it to the set
            if on_curve_mask[n_a] and on_curve_mask[n_b]:
                # Sort to store canonical representation, e.g., (1, 2) not (2, 1)
                elements_1d.append(tuple(sorted((n_a, n_b))))

    if not elements_1d:
        return jnp.array([], dtype=int)

    return jnp.array(elements_1d)


def generate_half_circle_mesh(
    radius: float = 1.0,
    mesh_size_fine: float = 0.4,
    mesh_size_coarse: float = 0.4,
    distance_from_rigid_plane: float = 1e-3,
):
    import os

    mesh_dir = os.path.join(os.getcwd(), "../meshes")
    os.makedirs(mesh_dir, exist_ok=True)
    output_filename = os.path.join(mesh_dir, "half_circle.msh")

    # Parameters
    R = radius  # Radius of the half-circle

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

    # Create points
    p1 = gmsh.model.geo.addPoint(0, R + distance_from_rigid_plane, 0)  # Center
    p2 = gmsh.model.geo.addPoint(-R, R + distance_from_rigid_plane, 0)  # Left point
    p3 = gmsh.model.geo.addPoint(R, R + distance_from_rigid_plane, 0)
    p4 = gmsh.model.geo.addPoint(0, distance_from_rigid_plane, 0)  # Center

    # Create the arc (half-circle)
    arc = gmsh.model.geo.addCircleArc(p2, p1, p4)
    arc_2 = gmsh.model.geo.addCircleArc(p4, p1, p3)

    # Create the flat line to close the half-circle
    line = gmsh.model.geo.addLine(p3, p1)
    line_2 = gmsh.model.geo.addLine(p1, p2)

    # Create a line loop and surface
    loop = gmsh.model.geo.addCurveLoop([arc, arc_2, line, line_2])
    surface = gmsh.model.geo.addPlaneSurface([loop])

    # Synchronize and mesh
    gmsh.model.geo.synchronize()

    # Create a Distance field to control mesh size near p_mid
    field_id = gmsh.model.mesh.field.add("Distance", 1)
    gmsh.model.mesh.field.setNumbers(field_id, "NodesList", [p4])

    # Create a Threshold field to refine near p_mid and coarsen further away
    thresh_field = gmsh.model.mesh.field.add("Threshold", 2)
    gmsh.model.mesh.field.setNumber(thresh_field, "InField", field_id)
    gmsh.model.mesh.field.setNumber(thresh_field, "SizeMin", mesh_size_fine)
    gmsh.model.mesh.field.setNumber(thresh_field, "SizeMax", mesh_size_coarse)
    gmsh.model.mesh.field.setNumber(thresh_field, "DistMin", R / 8)
    gmsh.model.mesh.field.setNumber(thresh_field, "DistMax", R / 4)

    # Set it as background mesh size field
    gmsh.model.mesh.field.setAsBackgroundMesh(thresh_field)

    # Generate the mesh
    gmsh.model.mesh.generate(2)
    gmsh.write(output_filename)
    print(f"Mesh successfully generated and saved to '{output_filename}'")

    gmsh.finalize()

    _mesh = meshio.read(output_filename)

    mesh = Mesh(
        coords=_mesh.points[:, :2],
        elements=_mesh.cells_dict["triangle"],
    )


    contact_elements = get_elements_on_curve(mesh, contact_line, tol=1e-6)

    return mesh, contact_elements


# get contact elements
# function identifies nodes on the line
def contact_line(coord: jnp.ndarray, tol: float) -> bool:
    return jnp.isclose(
        jnp.linalg.norm(coord - jnp.array([0, radius + distance_from_rigid_plane])),
        radius,
        atol=1e-10,
    )

Below, we generate the mesh and get the contact elements. We assume the following parameters for the mesh:

  • radius of the half-circle: 1.0 \(\text{m}\)
  • initial distance from the rigid plane: 1e-6 \(\text{m}\)
  • fine mesh size along the expected contact region: 0.02, reducing this value will increase the number of elements along the contact region
  • coarse mesh size elsewhere: 0.05
Hint

The mesh plays a critical role in FEM solution. Since the contact is localized, we need to resolve the contact region properly. If the mesh is not fine enough, the solution will not give the correct results. You will see whether the above chosen value for the mesh size is fine enough as we proceed with the solution.

radius = 1.0
distance_from_rigid_plane = 1e-6

mesh_params = {
    "radius": radius,
    "mesh_size_fine": 0.02, 
    "mesh_size_coarse": 0.05,
    "distance_from_rigid_plane": distance_from_rigid_plane,
}

mesh, contact_elements = generate_half_circle_mesh(**mesh_params)


contact_nodes = jnp.unique(contact_elements.flatten())
n_contact_nodes = len(contact_nodes)


n_nodes = mesh.coords.shape[0]
n_dofs_per_node = 2
n_dofs = n_dofs_per_node * n_nodes
Code: Plot the mesh 
plt.style.use(STYLE_PATH)
plt.figure(figsize=(5, 5))
ax = plt.axes()
ax.tripcolor(
    *mesh.coords.T,
    mesh.elements,
    facecolors=jnp.ones(mesh.elements.shape[0]),
    edgecolors="k",
    lw=0.2,
    cmap="managua_r",
)
ax.scatter(
    mesh.coords[contact_nodes, 0], mesh.coords[contact_nodes, 1], color=colors.red, s=4
)
ax.set_aspect("equal")
ax.margins(0.0, 0.0)
plt.show()
Figure 11.1: Mesh of the half-circle. The potential contact nodes are highlighted in red.

Just checking if the contact nodes are on the correct position.

np.allclose(
    np.linalg.norm(
        mesh.coords[contact_nodes] - np.array([0, radius + distance_from_rigid_plane]),
        axis=1,
    ),
    radius,
)
True

Defining material parameters

Let us define the material parameters mainly the shear modulus and the first Lamé parameter. The Youngs modulus and the Poisson’s ratio are related for the material are:

\[ E = 1.0 ~\text{N/m}^2, \quad \nu = 0.3 \]

From the above, we can compute the shear modulus and the first Lamé parameter as:

\[ \mu = \frac{E}{2(1+\nu)}, \quad \lambda = \frac{E\nu}{(1+\nu)(1-2\nu)} \]

from typing import NamedTuple


class Material(NamedTuple):
    """Material properties for the linear elastic material"""

    mu: float  # Shear modulus
    lmbda: float  # First Lamé parameter


E = 1.0
nu = 0.3
mu = E / (2 * (1 + nu))
lmbda = E * nu / ((1 + nu) * (1 - 2 * nu))


mat = Material(mu=mu, lmbda=lmbda)

Defining total elastic energy of the system

\[ \Psi_\text{e}(\boldsymbol{u}) = \frac{1}{2} \int_\Omega \psi_e d\Omega \]

where

\[ \psi_\text{e}(x) = \boldsymbol{\sigma}(x) : \boldsymbol{\varepsilon}(x) \]

where \(\boldsymbol{\sigma}\) is the stress tensor and \(\boldsymbol{\varepsilon}\) is the strain tensor.

\[ \boldsymbol{\sigma} = \lambda \text{tr}(\boldsymbol{\varepsilon}) \mathbf{I} + 2\mu \boldsymbol{\varepsilon} \]

and

\[ \boldsymbol{\varepsilon} = \frac{1}{2} (\nabla \boldsymbol{u} + \nabla \boldsymbol{u}^T) \]

tri = element.Tri3()
op = Operator(mesh, tri)


@autovmap(grad_u=2)
def compute_strain(grad_u: Array) -> Array:
    """Compute the strain tensor from the gradient of the displacement."""
    return 0.5 * (grad_u + grad_u.T)


@autovmap(eps=2)
def compute_stress(eps: Array, mu: float, lmbda: float) -> Array:
    """Compute the stress tensor from the strain tensor."""
    I = jnp.eye(2)
    return 2 * mu * eps + lmbda * jnp.trace(eps) * I


@autovmap(grad_u=2)
def strain_energy(grad_u: Array, mu: float, lmbda: float) -> Array:
    """Compute the strain energy density."""
    eps = compute_strain(grad_u)
    sig = compute_stress(eps, mu, lmbda)
    return 0.5 * jnp.einsum("ij,ij->", sig, eps)


@jax.jit
def total_strain_energy(u_flat: Array) -> float:
    """Compute the total energy of the system."""
    u = u_flat.reshape(-1, n_dofs_per_node)
    grad_u = op.grad(u)
    energy_density = strain_energy(grad_u, mat.mu, mat.lmbda)
    return op.integrate(energy_density)

Defining the contact energy

The contact energy is defined as:

\[ \Psi_\text{contact}(u)=\frac{1}{2}k_\text{pen}\int_{\Gamma_c}\Big\langle g_i(u)\Big\rangle^2 \text{d}\Gamma \] where \(\Gamma_c\) is the contact surface, \(k_\text{pen}\) is the penalty parameter, and \(g(u)\) is the gap function.

For this exercise, we define the discrete penalization only at the nodes on the contact surface. Therefore, the contact energy is defined as:

\[ \Psi_\text{contact}(u)=\frac{1}{2}k_\text{pen}\sum_{i=1}^{N_c}\Big\langle g_i(u)\Big\rangle^2 a_{i} \]

where \(N_c\) is the number of nodes on the contact surface, \(a_i\) is the area of the node \(i\), and \(g_i(u)\) is the gap function at node \(i\).

Defining the penalty parameter

The penalty parameter \(k_\text{pen}\) is a parameter that controls the stiffness of the contact. It is a scalar parameter that is used to penalize the violation of the contact constraint.

Below define the penalty parameter \(k_\text{pen}\).

k_pen = 1e6

Computing the area associated with each contact node

For correct resolution of the contact problem, we need the area associated with each contact node. In order to compute the area of each contact node, we define a line mesh that contains the contact nodes and the edges of the contact elements. We then define a function which integrates a constant field of value 1 over the line mesh.

\[ a_\text{total} = \int_{\Gamma_{c}} q(x) \text{d}\Gamma \]

where \(q(x)\) is a constant field of value 1. We can then compute the area of each node by differentiating the total area with respect to the field \(q(x)\).

\[ a_i = \frac{\partial a_\text{total}}{\partial q(x_i)} \]

Note

Below we compute the area associated with all the nodes in the mesh and store in the array nodes_area.

Thus the size of the array nodes_area = n_nodes i.e number of nodes in the mesh.

Later, we will use nodes_area to extract the area associated with the contact nodes.

line_mesh = Mesh(mesh.coords, contact_elements)
line2 = element.Line2()
line_op = Operator(line_mesh, line2)


@jax.jit
def total_area(u_flat: Array) -> float:
    u = u_flat.reshape(-1, 1)
    u_quad = line_op.eval(u)
    return jnp.sum(line_op.integrate(u_quad))


ones = jnp.ones(n_nodes)
nodes_area = jax.jacrev(total_area)(ones)
len(nodes_area) == n_nodes
True

Computing the contact energy

In Figure 11.1, we consider all the potential contact nodes that can be in contact with the rigid plane highlighted in red. This will be our active set. In order to check which of these nodes violates the contact constraint \(g(u_i) \leq 0\), we will use the Macaulay’s bracket.

\[ \Big \langle g_i(u) \Big \rangle = \begin{cases} 0 & \text{if } g_i(u) > 0 \\ g_i(u) & \text{if } g_i(u) \leq 0 \end{cases} \]

Active set strategy for Penalty Method

In the previous exercise, we used a Lagrange multiplier to enforce the contact constraint and therefore, we had set the \(\lambda_i\) to zero for the nodes which were not active. However, in the penalty method, we add \(k_\text{pen}\) to the stiffness matrix to enforce the contact constraint. Therefore, we need to make sure that where ever the contact constraint is violated, we do not add \(k_\text{pen}\) to the stiffness matrix at those locations.

By using the Macaulay’s bracket, we set the contact energy to zero for the nodes which are not in contact with the rigid plane.

In order to compute the contact energy, we need to compute the Macaulay’s bracket.

The Macaulay’s bracket is defined as:

\[ \Big\langle \star\Big\rangle = \begin{cases} 0 & \text{if } \star > 0 \\ \star & \text{if } \star \leq 0 \end{cases} \]

The gap function is defined as:

\[ g_i(u)=y^i_0+u_y^i-0.0 \]

Below, define the functions to define the Macaulay’s bracket and then use that function to compute the contact energy.

Hint

The function for Macaulay’s bracket will look like this:


def macalauy_bracket(x: float) -> float:
    """
    Computes the Macaulay bracket for the penalty method.
    This function should return x if x < 0, and 0 otherwise.
    """
    # TODO: Implement the Macaulay bracket using jnp.where(....)

    return ...§

To compute the contact energy, we will have to do the following:

  • Extract the displacements, coordinates, and areas for all the potential contact nodes.
  • Compute the gap function for each node.
  • Compute the penetration using the macaulay_bracket function you just wrote. This is applying the active set.
  • Compute the penalty energy for this single node: 0.5 _ k_pen _ penetration² * area.
  • Sum up all the individual nodal energies to get the total contact energy.

The function for the contact energy will look like this:


def compute_contact_energy(
    u: Array, coords: Array, contact_nodes: Array, nodes_area: Array
) -> float:
    """Compute the contact energy for a given displacement field.
    Args:
        u: Displacement field, shape (n_nodes, 2).
        coords: Initial coordinates of the nodes, shape (n_nodes, 2).
        contact_nodes: Indices of all the nodes on the contact surface, shape (n_contact_nodes,).
        nodes_area: Area of all the nodes in the mesh, shape (n_nodes,).
    Returns:
        Contact energy, a scalar.
    """

    # TODO Extract the displacements, coordinates, and areas for ONLY the contact_nodes.
    u_nodes = ...
    x_nodes = ...
    contact_nodes_area = ...

    # TODO Define a helper function that calculates the energy for a SINGLE node.
    # This function will be vectorized in the next step.
    def contact_energy_node(u_node: tuple[float, float], x_node: tuple[float, float], contact_area : float) -> float:
        """
        Computes the contact energy for a single node

        Args:

        u_node: displacements for a node (u_x, u_y)
        x_node: initial coordinate of a node (x, y)
        contact_area: area associated with a node
        """

        # Calculate the gap 'g_n'.
        gap = ...

        # Calculate the penetration using the macaulay_bracket function you just wrote.
        penetration = ...

        # Compute the penalty energy for this single node: 0.5 * k_pen * penetration² * area.
        return ...

    # TODO Vectorize the single-node function using jax.vmap.
    # This creates a new function that can process all nodes at once.
    # The in_axes argument should be (0, 0, 0) because we are iterating over the
    # first axis of u_nodes, x_nodes, and area_nodes.
    contact_energy_vmap = jax.vmap(...)

    # TODO Call the vectorized function with the arrays of nodal data you extracted in Step 1.
    all_nodal_energies = contact_energy_vmap(...)

    # TODO Sum up all the individual nodal energies to get the total contact energy.
    return ...
@jax.jit
def macaulay_bracket(x):
    return jnp.where(x > 0, 0, x)


@jax.jit
def compute_contact_energy(
    u: Array, coords: Array, contact_nodes: Array, nodes_area: Array
) -> Array:
    """Compute the contact energy for a given displacement field.
    Args:
        u: Displacement field.
        coords: Coordinates of the nodes.
        contact_nodes: Indices of the nodes on the contact surface.
        nodes_area: Area of the nodes on the contact surface.
    Returns:
        Contact energy.
    """
    u_nodes = u[contact_nodes]
    x_nodes = coords[contact_nodes]
    contact_nodes_area = nodes_area[contact_nodes]

    # Loop over nodes on the potential contact surface
    def _contact_energy_node(
        u_node: tuple[float, float], x_node: tuple[float, float], contact_area: float
    ) -> float:
        gap = (x_node[1] + u_node[1]) - 0.0
        penetration = macaulay_bracket(gap)
        return 0.5 * k_pen * (penetration**2) * contact_area

    contact_energy_node = jax.vmap(_contact_energy_node, in_axes=(0, 0, 0))

    return jnp.sum(contact_energy_node(u_nodes, x_nodes, contact_nodes_area))

Compute the total energy

Now lets us define the total energy function.

\[ \Psi(u) = \Psi_\text{elastic}(u) + \Psi_\text{contact}(u) \]

Hint

The total energy function will look like this:


def total_energy(
    u_flat: Array,
    coords: Array,
    contact_nodes: Array,
    nodes_area: Array,
) -> float:
    """Compute the total energy for a given displacement field.
    Args:
        u_flat: Flattened displacement field, shape (n_dofs,).
        coords: Initial coordinates of the nodes, shape (n_nodes, 2).
        contact_nodes: Indices of the nodes on the contact surface, shape (n_contact_nodes,).
        nodes_area: Area of the nodes on the contact surface, shape (n_nodes,).
    Returns:
        Total energy, a scalar
    """

    # TODO compute the contact energy
    contact_energy = ...

    # TODO compute the elastic energy
    strain_energy = ...

    # TODO Sum the contact and strain energies to get the total energy of the system.
    return ...
@jax.jit
def _total_energy(
    u_flat: Array,
    coords: Array,
    contact_nodes: Array,
    nodes_area: Array,
) -> Array:
    """Compute the total energy for a given displacement field.
    Args:
        u_flat: Flattened displacement field.
        coords: Coordinates of the nodes.
        contact_nodes: Indices of the nodes on the contact surface.
        nodes_area: Area associated with all the nodes in the mesh
    Returns:
        Total energy.
    """
    u = u_flat.reshape(-1, n_dofs_per_node)
    contact_energy = compute_contact_energy(
        u, coords, contact_nodes, nodes_area
    )
    strain_energy = total_strain_energy(u)
    return contact_energy + strain_energy


total_energy = partial(
    _total_energy,
    coords=mesh.coords,
    contact_nodes=contact_nodes,
    nodes_area=nodes_area,
)

Computing the internal forces and the stiffness matrix

For this exercise, we will like to use the sparse solver, given the large number of degrees of freedom. We will use the sparse module to compute the sparsity pattern and then use the sparse.jacfwd to compute the Jacobian of the internal forces.

\[ \boldsymbol{f}_\text{int} = \frac{\partial \Psi}{\partial \boldsymbol{u}} \]

and

\[ \mathbf{K} = \frac{\partial \boldsymbol{f}_\text{int}}{\partial \boldsymbol{u}} \]

We make use of the sparse.create_sparsity_pattern to define the sparsity pattern of the problem. Below, we define steps to:

  • Create the sparsity pattern
  • Compute the internal force vector from the total energy using jax.jacrev
  • Compute the sparse stiffness matrix from the internal force vector using sparse.jacfwd and the sparsity pattern
Hint

Be careful when defining the sparsity pattern. Unlike the previous exercises, where we had extra degrees of freedom for lagrange multipliers, here we have only the displacement degrees of freedom. So the sparsity pattern would only be dependent on the mesh and the number of displacement dofs per node.

sparsity_pattern = sparse.create_sparsity_pattern(mesh, n_dofs_per_node=n_dofs_per_node)


gradient = jax.jacrev(total_energy)
hessian_sparse = sparse.jacfwd(gradient, sparsity_pattern=sparsity_pattern)

Applying Dirichlet boundary conditions

We push the deformable cylinder into the rigid plane by applying a displacement in the \(y\) direction on the top edge of the cylinder. To avoid rigid body motion, we fix the displacement in the \(x\) direction on the top edge of the cylinder. The boundary conditions thus applied are:

\[ \begin{aligned} u_y &= -R/500 \quad \text{on the top edge} \\ \end{aligned} \]

where \(R\) is the radius of the cylinder.

Hint

Get the nodes on the top edge of the cylinder and then get the dofs associated with those nodes.

top_nodes = # get unique nodes from the interface_elements

fixed_dofs = [dofs_associated_with_top_nodes_in_y_direction]
y_max = np.max(mesh.coords[:, 1])

top_nodes = jnp.where(jnp.isclose(mesh.coords[:, 1], y_max))[0]
fixed_dofs = jnp.concatenate([2 * top_nodes + 1])
free_dofs = jnp.setdiff1d(jnp.arange(n_dofs), fixed_dofs)

prescribed_values = jnp.zeros(n_dofs)
prescribed_values = prescribed_values.at[2 * top_nodes + 1].set(-radius / 500)

We apply boundary conditions using the lifting approach. To apply the lifting to a sparse stiffness matrix, we need the indices corresponding the rows and columns that are fixed (to be set to zero) and the indices of the diagonal elements that are set to one.

Below use the sparse.get_bc_indices to get the indices (row, col) where we set the stiffness entries to zero and the indices of the diagonal elements that are set to one.

See Sparse solvers for more details.

Hint

The indices will look like this:

zero_indices, one_indices = ... # get the indices from the sparse.get_bc_indices function
zero_indices, one_indices = sparse.get_bc_indices(sparsity_pattern, fixed_dofs)

Solving the problem

We will use the Newton’s method to solve the problem. Remember we use the lifting approach to apply the boundary conditions, therefore, we need to modify the stiffness matrix and the residual vector to account for the boundary conditions. See Sparse solvers for more details on how to use the lifting approach within a Newton solver.

Hint

The Newton solver will look like this:

import scipy.sparse as sp
from jax import Array
import jax.numpy as jnp

def newton_sparse_solver(
    u: Array,
    fext: Array,
    gradient: callable,
    hessian_sparse: callable,
    fixed_dofs: Array,
    zero_indices: Array,
    one_indices: Array,
    ...,
    ...

) -> tuple[Array, float]:
    """Solves the nonlinear system using Newton's method with a sparse direct solver.
    
    Args:
        u: Current displacement field, shape (n_dofs,).
        fext: External force vector, shape (n_dofs,).
        gradient: Function to compute the internal forces
        hessian_sparse: Function to compute the sparse stiffness matrix
        fixed_dofs: Indices of the fixed degrees of freedom where we apply the Dirichlet boundary conditions, shape (n_fixed_dofs,).
        zero_indices: Indices of the rows and columns to set to zero in the stiffness matrix, because of the Dirichlet boundary conditions,
        one_indices: Indices of the diagonal elements to set to one in the stiffness matrix, because of the Dirichlet boundary conditions.
        ...: Any other arguments required by the gradient and hessian_sparse functions.
        ...: Any other arguments required by the gradient and hessian_sparse functions.
    Returns:
        u: Displacement field, shape (n_dofs,).
        norm_res: Norm of the residual, a scalar.
    """

    fint = gradient(u)
    iiter = 0
    norm_res = 1.0
    tol = 1e-8
    max_iter = 10

    while norm_res > tol and iiter < max_iter:
        
        # TODO: Calculate the residual vector R = F_ext - F_int.
        # For the lifting method, the residual at the fixed DOFs must be set to zero.
        residual = ...
        
        # TODO: Compute the sparse stiffness matrix K at the current displacement u.
        K_sparse = ...

        # TODO: Apply the lifting method to modify the stiffness matrix data.
        # "Zero out" the rows/columns for the fixed DOFs and set the diagonal to 1.
        # Use the indices for these operations (zero_indices, one_indices)
        K_data_lifted = ...

        # TODO: Reconstruct the modified scipy.sparse.csr_matrix for the solver.
        # Use the modified data `K_data_lifted` and the original indices from `K_sparse`.
        K_csr = ...


        # TODO: Solve the modified linear system K_mod * du = R_mod for the update `du`.
        # Use the direct solver: scipy.sparse.linalg.spsolve.
        du = ...
        
        # TODO: Update the total displacement vector u.
        u = ...

        # TODO: Re-compute the internal forces `fint` at the new position `u`
        # to get the latest residual for the convergence check.
        fint = ...

        # TODO: Calculate the new residual and its norm.
        # Again, remember to zero out the residual at the fixed DOFs before taking the norm.
        residual = ...
        norm_res = ...

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

    return u, norm_res
Newton solver 
import scipy.sparse as sp


def newton_scipy_solver(
    u,
    fext,
    gradient,
    hessian_sparse,
    fixed_dofs,
    zero_indices,
    one_indices,
):
    fint = gradient(u)

    iiter = 0
    norm_res = 1.0

    tol = 1e-8
    max_iter = 10

    while norm_res > tol and iiter < max_iter:
        residual = fext - fint
        residual = residual.at[fixed_dofs].set(0.0)

        K_sparse = hessian_sparse(u)
        K_data_lifted = K_sparse.data.at[zero_indices].set(0)
        K_data_lifted = K_data_lifted.at[one_indices].set(1)

        K_csr = sp.csr_matrix(
            (K_data_lifted, (K_sparse.indices[:, 0], K_sparse.indices[:, 1]))
        )

        du = sp.linalg.spsolve(K_csr, residual)
        u = u.at[:].add(du)

        fint = gradient(u)
        residual = fext - fint
        residual = residual.at[fixed_dofs].set(0)
        norm_res = jnp.linalg.norm(residual)

        print(f"  Residual: {norm_res:.2e}")

        iiter += 1

    return u, norm_res

Implementing the Incremental Loading Loop

Now that we have the mesh, the energy functions, and the newton_sparse_solver, it’s time to set up the main simulation loop. We will apply the displacement in small increments and solve for equilibrium at each step. This is important because using penalty method, the problem becomes non-linear and the solution is sensitive to the initial guess.

Follow these steps to build the main part of your script.

Hint

Initialization

Before the loop, you need to initialize all the variables for the simulation.

  • Create the initial solution vector u_prev, which contains both displacements. Initialize it as a vector of zeros with the correct total size (n_dofs).
  • Initialize the external force vector fext as a vector of zeros.
  • Set the number of load increments, n_steps (e.g., 10). We would apply the total applied_displacement in n_steps.`
  • Calculate the displacement to apply at each step.

The Main Loading Loop

Create a for loop that iterates through each load step.

for step in range(n_steps):
    print(f"Step {step+1}/{n_steps}")
    # ... Your code for each step will go inside this loop ...

Inside the Loop: Apply Displacement and Prepare Solver

At the beginning of each loop iteration, you need to update the boundary conditions and prepare the functions for the solver.

  • Update the u_prev vector with the prescribed displacement for the current step.

Inside the Loop: Solve for Equilibrium

  • Call the newton_sparse_solver function with the initial guess and other arguments. Depending on what inputs your gradient and hessian_sparse functions require, you may need to pass additional arguments to the newton_sparse_solver function. You will also have to pass the fixed_dofs, zero_indices, and one_indices to the newton_sparse_solver function to apply the lifting method.
  • Store the new solution in u_new.
  • Update the state for the next iteration by setting u_prev = u_new. This is basically providing the solution from the previous step as the initial guess for the next step., which helps the solver converge faster.

After the Loop: Final Solution

Once the loop is complete, extract the final displacement field from the final u_prev vector for visualization.

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

n_steps = 10

applied_displacement = prescribed_values / n_steps  # displacement increment

u_per_step = []
for step in range(n_steps):
    print(f"Step {step + 1}/{n_steps}")
    u_prev = u_prev.at[fixed_dofs].add(applied_displacement[fixed_dofs])

    u_new, rnorm = newton_scipy_solver(
        u_prev,
        fext,
        gradient,
        hessian_sparse,
        fixed_dofs,
        zero_indices,
        one_indices,
    )

    u_prev = u_new
    u_per_step.append(u_prev.reshape(n_nodes, n_dofs_per_node))

u_solution = u_prev.reshape(n_nodes, n_dofs_per_node)

Visualizing the the stress field

Let’s plot the stress and the deformed configuration. We plot the \(\sigma_{yy}\) stress.

Plotting the stress field 
# squeeze to remove the quad point dimension (only 1 quad point)
grad_u = op.grad(u_solution).squeeze()
strains = compute_strain(grad_u)
stresses = compute_stress(strains, mat.mu, mat.lmbda)

plt.style.use(STYLE_PATH)
fig, axs = plt.subplots(1, 2, figsize=(6, 4), layout="constrained")

plot_element_values(
    u=u_solution,
    mesh=mesh,
    values=stresses[:, 1, 1].flatten(),
    ax=axs[0],
    label=r"$\sigma_{yy}$",
)

axs[0].axhline(y=0, color="black", linewidth=0.5)
axs[0].set_xlabel("$x$")
axs[0].set_ylabel("$y$")
axs[0].set_aspect("equal")
axs[0].margins(0.0, 0.0)

plot_nodal_values(
    u=u_solution,
    mesh=mesh,
    nodal_values=u_solution[:, 1].flatten(),
    ax=axs[1],
    label=r"$u_y$",
)

axs[1].axhline(y=0, color="black", linewidth=0.5)
axs[1].set_xlabel("$x$")
axs[1].set_ylabel("$y$")
axs[1].set_aspect("equal")
axs[1].margins(0.0, 0.0)

plt.show()
Figure 11.2: Stress field \(\sigma_{yy}\) after the contact is resolved. The undeformed configuration is shown in the background.

Checking the effect of penalty parameter on the gap function

Since, we used the penalty approach, the nodes are not exactly on the contact surface but depending on the penalty parameter, they will penetrate the surface. Let us see how the gap function looks like at the contact surface.

Code: Plot the zoomed in contact surface to see the node locations. 
new_position = mesh.coords + u_solution
sorted_idx = jnp.argsort(mesh.coords[contact_nodes, 0])
gap = macaulay_bracket(new_position[contact_nodes[sorted_idx], 1])


plt.style.use(STYLE_PATH)
fig = plt.figure(figsize=(4, 3), layout="constrained")
ax = plt.axes()
ax.plot(mesh.coords[contact_nodes[sorted_idx], 0], gap)
ax.set_xlabel(r"$x$")
ax.set_ylabel(r"$\mathrm{gap}, \langle g \rangle$")
ax.set_xlim(-radius / 5, radius / 5)
ax.set_aspect(1 / ax.get_data_ratio())
ax.margins(0.0, 0.0)
ax.grid(True)


plt.show()
plt.close()
Figure 11.3: Zoomed in contact surface to see the node locations. We see that there is some interpenetration due to the penalty parameter.

The above results looks good but we do not know whether the solution that we get is correct. We do have convergence at each applied displacement but we do not know whether the solution is correct.

Therefore, to verify the solution, we can compare the solution with Hertz theory that is explained in Fundamentals.

Comparing numerical results with Hertz contact theory

As mentioned earlier, the Hertz solution is characterized by a contact radius \(a\) and a contact pressure \(p_0\). The contact radius \(a\) is given by,

\[ a = \sqrt{\frac{4}{\pi} \frac{PR}{E^*}}, \quad p_0 = \frac{2P}{\pi a} \]

where \(P\) is the load and \(E^{\star}\) is the effective Young’s modulus of the cylinder. For plain strain condition, the effective Young’s modulus is given by, \[ E^* = \frac{E}{1-\nu^2} \]

Below compute the contact radius \(a\) and the contact pressure \(p_0\) from Hertz theory.

P = f_contact
E_eff = E / (1 - nu**2)  # effective Young's modulus for plain strain condition
a0 = np.sqrt(4 * P * radius / jnp.pi / E_eff)
p0 = 2 * P / jnp.pi / a0


x_hertz = np.linspace(-a0, a0, 100)
x_hertz = x_hertz[x_hertz != 0]
p_hertz = 2 * P * np.sqrt(a0**2 - x_hertz**2) / jnp.pi / a0**2


print(f"Contact radius from Hertz theory {a0}")
print(f"Contact pressure from Hertz theory {p0}")
Contact radius from Hertz theory 0.03022952413571173
Contact pressure from Hertz theory 0.016609628645995456

Comparing contact pressure profile with Hertz theory

Based on Hertz theory, the traction profile along the contact surface is given by,

\[ p(r) = -p_0 \sqrt{1-(\frac{r}{a})^2} \]

where \(r\) is the radial distance from the center of the contact surface and it varies from \(-a\) to \(a\).

Now, we will extract the surface tractions \(t_n = \boldsymbol{\sigma} : \boldsymbol{n}\) along the contact surface from our numerical solution, where \(\boldsymbol{n}\) is the normal vector to the contact surface and \(\boldsymbol{\sigma}\) is the stress tensor in the elements along the contact surface. Below we have defined a function to extract the surface tractions from the numerical solution. This function calculates the surface traction \(t_n\) along the contact surface for each element and returns the points, tractions, and normals of the surface elements.

Use this function to extract the surface tractions.

Code: Function to extract surface tractions 
def get_surface_tractions(
    surface_elements: np.ndarray,
    mesh: Mesh,
    element_values: jnp.ndarray,
) -> np.ndarray:
    """
    Finds triangle elements that share a common edge with given surface elements,
    For each triangle element, it computes the traction on the surface element.
    It returns the points, tractions, and normals of the surface elements.

    Args:
        surface_elements (np.ndarray): An array of surface connectivity, shape (num_surface_elements, 2).
        mesh (Mesh): A mesh object.
        element_values (jnp.ndarray): An array of element values, shape (num_elements,).

    Returns:
        np.ndarray: An array of points, shape (num_surface_elements, 2).
        np.ndarray: An array of tractions, shape (num_surface_elements, 2).
        np.ndarray: An array of normals, shape (num_surface_elements, 2).
    """
    line_edge_set = {tuple(sorted(edge)) for edge in surface_elements}

    tractions = []
    points = []
    normals = []

    for i, tri in enumerate(mesh.elements):
        p1, p2, p3 = tri

        edges = [
            tuple(sorted((p1, p2))),
            tuple(sorted((p2, p3))),
            tuple(sorted((p3, p1))),
        ]

        for edge in edges:
            if edge in line_edge_set:
                stress_value = element_values[i]
                coords = np.array([mesh.coords[edge[0]], mesh.coords[edge[1]]])
                tangent_vector = coords[1] - coords[0]
                normal = jnp.array([-tangent_vector[1], tangent_vector[0]])
                normal = normal / jnp.linalg.norm(normal)
                traction = stress_value @ normal

                mid_point = jnp.mean(coords, axis=0)

                tractions.append(traction)
                points.append(mid_point)
                normals.append(normal)
                break

    return np.array(points), np.array(tractions)

Below we call the function get_surface_tractions to extract the surface tractions. The surface tractions are stored in the variable surface_tractions and the surface points are stored in the variable surface_points.

surface_points, surface_tractions = get_surface_tractions(
    surface_elements=np.array(contact_elements),
    mesh=mesh,
    element_values=stresses,
)

Once you have the surface tractions, you should plot the traction profile along the contact surface. Also plot the \(p(r)\) from Hertz theory and see if they match.

Plot the numerical Contact Pressure (from your simulation):

  • X-axis: The x-coordinates of the nodes on the contact surface (surface_points).

  • Y-axis: The magnitude of the contact pressure from your simulation. This is the absolute value of the y-component of the surface_tractions.

Plot the Analytical Hertzian Pressure:

  • X-axis: The array of x-coordinates from the Hertzian calculation, you can either use the surface_points or define a new array of x-coordinates for smoother curve.

  • Y-axis: The array of corresponding pressure values from the Hertzian calculation.

Code: Plot the traction profile. 
plt.style.use(STYLE_PATH)
fig = plt.figure(figsize=(6, 3))
ax = plt.axes()
sorted_idx2 = np.argsort(surface_points[:, 0])

ax.plot(
    surface_points[sorted_idx2, 0] ,
    np.abs(surface_tractions[sorted_idx2, 1]) ,
    color=colors.blue,
    marker="o",
    zorder=20,
    label=r"$\mathbf{\sigma} : \mathbf{n}$",
)

ax.plot(x_hertz, np.abs(p_hertz) , color="gray", label="Hertz", zorder=-10)

ax.set_xlabel(r"$x$")
ax.set_ylabel(r"Contact Pressure/Traction")
ax.set_xlim(-1.2*a0, 1.2*a0)
ax.grid(True)
ax.set_ylim(bottom=0)
ax.legend(frameon=False)
plt.show()
Figure 11.4: Profile of traction normal (\(t_n = \boldsymbol{\sigma} : \boldsymbol{n}\)) to the curved surface computed from the numerical solution (blue) and compare it with Hertz theory (gray).

Does the numerical solution match with Hertz theory?

  • Do you think the numerical solution matches with Hertz theory? Why or why not?
  • If doesnot match, what is the reason for the mismatch? Is it due to the penalty parameter? Is it due to the fact the mesh is not fine enough at the contact surface?
  • Try to improve the match by changing the penalty parameter and the mesh.

Comparing stresses inside of the cylinder with Hertz theory

We also compare the stresses inside of the cylinder along the line \(x=0\) with Hertz theory. According to Hertz theory, the stresses are given by:

\[ \sigma_x = -\frac{p_0}{a_0} \left( m (1 + \frac{y^2 + n^2}{m^2 + n^2}) - 2y \right) \]

\[ \sigma_y = -\frac{p_0}{a_0} \left( m (1 - \frac{y^2 + n^2}{m^2 + n^2}) \right) \]

where \(m\) and \(n\) are given by:

\[ m = \sqrt{0.5 \left( \sqrt{(a_0^2 - x^2 + y^2)^2 + 4x^2y^2} + (a_0^2 - x^2 + y^2) \right)} \]

\[ n = \sqrt{0.5 \left( \sqrt{(a_0^2 - x^2 + y^2)^2 + 4x^2y^2} - (a_0^2 - x^2 + y^2) \right)} \]

The principal shear stresses are given by:

\[ \tau_1 = \frac{1}{2} \left( \sigma_y - \sigma_x \right) \]

Below define a function to compute the stresses inside of a body based on Hertz solution.

Hint

The function definition will be something like this:

def hertz_solution(x, y, a0, p0):
    """Compute the stresses based on Hertz solution.
    
    Args:
        x (float): The x-coordinate of the point.
        y (float): The y-coordinate of the point.
        a0 (float): The contact radius.
        p0 (float): The contact pressure.

    Returns:
        tuple: A tuple of three floats representing the $\sigma_{xx}$, $\sigma_{yy}$, and $\sigma_{xy}$ stresses.
    """

    # TODO:Compute the stresses based on Hertz solution
    ...

 
    # Return the stresses
    return sigma_x, sigma_y, sigma_xy
Code: Function to compute the stresses based on Hertz solution 
def hertz_solution(x, y, a0, p0):
    m = np.sqrt(
        0.5
        * (
            ((a0**2 - x**2 + y**2) ** 2 + 4 * x**2 * y**2) ** (1 / 2)
            + (a0**2 - x**2 + y**2)
        )
    )
    n = np.sqrt(
        0.5
        * (
            ((a0**2 - x**2 + y**2) ** 2 + 4 * x**2 * y**2) ** (1 / 2)
            - (a0**2 - x**2 + y**2)
        )
    )

    sigma_x = -p0 * (m * (1 + (y**2 + n**2) / (m**2 + n**2)) - 2 * y) / a0
    sigma_y = -p0 * m * (1 - (y**2 + n**2) / (m**2 + n**2)) / a0
    sigma_xy = -p0 * n * ((m**2 - y**2) / (m**2 + n**2)) / a0

    return sigma_x, sigma_y, sigma_xy
Code: Function to find the index of the containing polygon for each point. 
import numpy as np


@jax.jit
def find_containing_polygons(
    points: jnp.ndarray,
    polygons: jnp.ndarray,
) -> jnp.ndarray:
    """
    Finds the index of the containing polygon for each point.

    This function uses a vectorized Ray Casting algorithm and is JIT-compiled
    for maximum performance. It assumes polygons are non-overlapping.

    Args:
        points (jnp.ndarray): An array of points to test, shape (num_points, 2).
        polygons (jnp.ndarray): A 3D array of polygons, where each polygon is a
                                list of vertices. Shape (num_polygons, num_vertices, 2).

    Returns:
        jnp.ndarray: An array of shape (num_points,) where each element is the
                     index of the polygon containing the corresponding point.
                     Returns -1 if a point is not in any polygon.
    """

    # --- Core function for a single point and a single polygon ---
    def is_inside(point, vertices):
        px, py = point

        # Get all edges of the polygon by pairing vertices with the next one
        p1s = vertices
        p2s = jnp.roll(vertices, -1, axis=0)  # Get p_{i+1} for each p_i

        # Conditions for a valid intersection of the horizontal ray from the point
        # 1. The point's y-coord must be between the edge's y-endpoints
        y_cond = (p1s[:, 1] <= py) & (p2s[:, 1] > py) | (p2s[:, 1] <= py) & (
            p1s[:, 1] > py
        )

        # 2. The point's x-coord must be to the left of the edge's x-intersection
        # Calculate the x-intersection of the ray with the edge
        x_intersect = (p2s[:, 0] - p1s[:, 0]) * (py - p1s[:, 1]) / (
            p2s[:, 1] - p1s[:, 1]
        ) + p1s[:, 0]
        x_cond = px < x_intersect

        # An intersection occurs if both conditions are met.
        intersections = jnp.sum(y_cond & x_cond)

        # The point is inside if the number of intersections is odd.
        return intersections % 2 == 1

    # --- Vectorize and apply the function ---
    # Create a boolean matrix: matrix[i, j] is True if point i is in polygon j
    # Vmap over points (axis 0) and polygons (axis 0)
    # in_axes=(0, None) -> maps over points, polygon is fixed
    # in_axes=(None, 0) -> maps over polygons, point is fixed
    # We vmap the second case over all points
    is_inside_matrix = jax.vmap(
        lambda p: jax.vmap(lambda poly: is_inside(p, poly))(polygons)
    )(points)

    # Find the index of the first 'True' value for each point (row).
    # This gives the index of the containing polygon.
    # We add a 'False' column to handle points outside all polygons.
    # jnp.argmax will then return the index of this last column.
    padded_matrix = jnp.pad(
        is_inside_matrix, ((0, 0), (0, 1)), "constant", constant_values=False
    )
    indices = jnp.argmax(padded_matrix, axis=1)

    # If the index is the last one, it means the point was not in any polygon.
    # We map this index to -1 for clarity.
    return jnp.where(indices == is_inside_matrix.shape[1], -1, indices)

Next, we will compute the stresses along the line \(x=0\) from our numerical solution and compare them with the stresses from Hertz theory.

Below, we define a set of points along the line \(x=0\) and find the element indices of the containing polygons for each point. We use the function find_containing_polygons to find the element indices. Then we select the stresses from the stress tensor for the corresponding element indices.

The stresses are stored in the variable stresses_along_points which is a 3D array of shape \((n, 2, 2)\) where \(n\) is the number of points. The first dimension corresponds to the points, the second dimension corresponds to the stress components, and the third dimension corresponds to the stress tensor. We will plot different components of the stress tensor along the line \(x=0\) and compare them with solution from Hertz theory as discussed above.

y_min = jnp.min(mesh.coords[:, 1])
y_max = jnp.max(mesh.coords[:, 1])
y = np.linspace(y_min, radius / 2, 100)
x = jnp.full_like(y, fill_value=0.0)

points = np.stack([x, y], axis=1)
element_indices = find_containing_polygons(points, mesh.coords[mesh.elements])

grad_u = op.grad(u_solution).squeeze()
strains = compute_strain(grad_u)
stresses = compute_stress(strains, mat.mu, mat.lmbda)

stresses_along_points = stresses[element_indices]

Below, we plot the stresses (\(\sigma_{yy}, \sigma_{xx}, \tau_{max}\)) along the line \(x=0\) from the numerical solution and compare them with the stresses from Hertz theory.

In the plot, plot the stresses along the \(x-\) axis and the \(y-\) coordinates along the \(y-\) axis.

Code: Plot stresses along the line \(x=0\) 
plt.style.use(STYLE_PATH)
fig = plt.figure(figsize=(4, 4), layout="constrained")
ax = plt.axes()
ax.plot(
    stresses_along_points[:, 1, 1] ,
    y ,
    "o-",
    label=r"$\sigma_{yy}$",
    color=colors.blue,
    markevery=2,
)

sigma_x, sigma_y, sigma_xy = hertz_solution(x, y, a0, p0)
ax.plot(sigma_y , y , color=colors.blue, ls="--", lw=1.0)

ax.plot(
    stresses_along_points[:, 0, 0] ,
    y ,
    "o-",
    label=r"$\sigma_{xx}$",
    color=colors.red,
    markevery=2,
)
ax.plot(sigma_x , y , color=colors.red, ls="--", lw=1.0)


tau_max = 0.5 * jnp.abs(
    stresses_along_points[:, 1, 1] - stresses_along_points[:, 0, 0]
)

ax.plot(
    tau_max ,
    y ,
    "o-",
    label=r"$\tau_{max}=\frac{1}{2}(\sigma_{yy}-\sigma_{xx})$",
    color=colors.green,
    markevery=2,
)

ax.plot(
    p0 * (y - y**2 * np.power(a0**2 + y**2, -1 / 2)) / a0 ,
    y ,
    color=colors.green,
    ls="--",
    lw=1.0,
)

ax.axvline(x=0, color="k", linewidth=0.5, zorder=-10)
ax.axvline(x=0.3*p0, color="gray", linewidth=0.5)
ax.axhline(y=0.78*a0, color="gray", linewidth=0.5)
ax.set_ylabel(r"$y$")
ax.set_xlabel(r"$\sigma$")
ax.set_ylim(0, 2.0*a0)
ax.set_xlim(-1.5*p0, 0.5*p0)
ax.set_aspect(1 / ax.get_data_ratio())
ax.grid(True)
ax.margins(0.0, 0.0)
ax.legend(frameon=False)

plt.show()
Figure 11.5: Stresses along the line \(x=0\) from the numerical solution (in solid line) and comparison with Hertz theory (in dashed line)

Does the numerical solution match with Hertz theory?

  • Do you think the numerical solution matches with Hertz theory? Why or why not? Try to improve the match by changing the penalty parameter and the mesh.