3D Hyperelastic Finite Element Analysis in 150 Lines of Python Code

minimal-fea-3d-hyperelastic.py

import numpy as np
from types import SimpleNamespace as SN
from scipy.sparse import csr_matrix
from scipy.sparse.linalg import spsolve
import matplotlib.pyplot as plt


def Cube(a=(0, 0, 0), b=(1, 1, 1), n=(2, 2, 2)):
    grid = np.linspace(a[0], b[0], n[0])
    points = np.pad(grid[:, None], ((0, 0), (0, 2)))
    cells = np.arange(n[0]).repeat(2)[1:-1].reshape(-1, 2)
    for dim, sl in enumerate([slice(None, None, -1), slice(None)]):
        c = [cells + len(points) * a for a in np.arange(n[dim + 1])]
        expand = np.linspace(a[dim + 1], b[dim + 1], n[dim + 1])
        points = np.vstack(
            [points + np.insert(np.zeros(2), dim + 1, h) for h in expand]
        )
        cells = np.vstack([np.hstack((a, b[:, sl])) for a, b in zip(c[:-1], c[1:])])
    return SN(points=points, cells=cells)


def Hexahedron():
    quadrature = np.meshgrid([-1, 1], [-1, 1], [-1, 1])
    r, s, t = np.concatenate(quadrature).reshape(3, -1) / np.sqrt(3)
    a = np.array([[-1, 1, 1, -1, -1, 1, 1, -1]]).T
    b = np.array([[-1, -1, 1, 1, -1, -1, 1, 1]]).T
    c = np.array([[-1, -1, -1, -1, 1, 1, 1, 1]]).T
    ar, bs, ct = 1 + a * r, 1 + b * s, 1 + c * t
    gradient = np.stack([a * bs * ct, ar * b * ct, ar * bs * c], axis=1)
    return SN(function=ar * bs * ct / 8, gradient=gradient / 8)


def Domain(mesh, element):
    dXdr = np.einsum("cpi,pjq->ijcq", mesh.points[mesh.cells], element.gradient)
    drdX = np.linalg.inv(dXdr.T).T
    return SN(
        mesh=mesh,
        element=element,
        gradient=np.einsum("piq,ijcq->pjcq", element.gradient, drdX),
        dx=np.linalg.det(dXdr.T).T,  # Gauss-Legendre quadrature-weights are all 1
    )


def VectorField(region, values):
    def grad(values):
        return np.einsum(
            "cpi,pjcq->ijcq", values[region.mesh.cells], region.gradient
        ).copy()

    return SN(region=region, values=values, gradient=grad)


def SolidBody(field, umat, **kwargs):
    mesh = field.region.mesh
    dudX = field.gradient(field.values)
    dhdX = field.region.gradient
    dF = np.einsum("im,aj...->ijam...", np.eye(3), dhdX)
    δF = dF[..., None, None, :, :, :, :]
    ΔF = dF[..., :, :, None, None, :, :]
    F = (np.eye(3)[..., None, None] + dudX)[:, :, None, None, None, None, ...]

    vector, matrix = umat(δF, F, ΔF, dV=field.region.dx, **kwargs)

    idx = 3 * np.repeat(mesh.cells, 3) + np.tile(np.arange(3), mesh.cells.size)
    idx = idx.reshape(*mesh.cells.shape, 3)
    vidx = (idx.ravel(), np.zeros_like(idx.ravel()))
    midx = (
        np.repeat(idx, 3 * idx.shape[1]),
        np.tile(idx, (1, idx.shape[1] * 3, 1)).ravel(),
    )
    return SN(
        vector=csr_matrix((vector.sum(-1).transpose([4, 0, 1, 2, 3]).ravel(), vidx))
        .toarray()
        .ravel(),
        matrix=csr_matrix((matrix.sum(-1).transpose([4, 0, 1, 2, 3]).ravel(), midx)),
    )


ascontiguous = lambda *args: [np.ascontiguousarray(x) for x in args]
transpose = lambda x: np.einsum("ij...->ji...", x)
dot = lambda x, y, z: np.einsum("ik...,kl...,lj...->ij...", *ascontiguous(x, y, z))
ddot = lambda x, y: np.einsum("ij...,ij...->...", *ascontiguous(x, y))


def neo_hooke(δF, F, ΔF, dV, mu=1, bulk=30):
    lnJ = np.log((np.linalg.det(F.T).T))
    iFT = transpose(np.linalg.inv(F.T).T)
    ΔiFT = -dot(iFT, transpose(ΔF), iFT)
    δW = mu * ddot(F, δF) + (bulk * lnJ - mu) * ddot(iFT, δF)
    ΔδW1 = mu * ddot(ΔF, δF) + (bulk * lnJ - mu) * ddot(ΔiFT, δF)
    return δW * dV, (ΔδW1 + bulk * ddot(iFT, δF) * ddot(iFT, ΔF)) * dV


cube = Cube(a=(0, 0, 0), b=(25, 100, 100), n=(6, 21, 21))
region = Domain(cube, Hexahedron())
field = VectorField(region, values=np.zeros_like(cube.points))

dofs = np.arange(cube.points.size).reshape(cube.points.shape)
dof = SN(
    left=dofs[cube.points[:, 0] == 0].ravel(),
    right=dofs[cube.points[:, 0] == 25][:, [1, 2]].ravel(),
    move=dofs[cube.points[:, 0] == 25][:, 0].ravel(),
)
dof.active = np.delete(
    dofs.ravel(), np.unique(np.concatenate([dof.left, dof.right, dof.move]))
)

moves = -np.linspace(0, 1, 6) * 5
forces = np.zeros_like(moves)
ext = np.ones(len(dof.move))
solid = SolidBody(field, umat=neo_hooke, mu=1, bulk=30)
b1 = -solid.vector[dof.active]

for increment, move in enumerate(moves):
    print(f"Increment {increment + 1} | move={move:1.2f}")

    for iteration in range(8):
        A11 = solid.matrix[dof.active, :][:, dof.active]
        A10 = solid.matrix[dof.active, :][:, dof.move]
        du1 = spsolve(A11, b1 - A10.dot(ext * move - field.values.ravel()[dof.move]))

        field.values.ravel()[dof.active] += du1
        field.values.ravel()[dof.move] = ext * move

        solid = SolidBody(field, umat=neo_hooke, mu=1, bulk=30)
        b1 = -solid.vector[dof.active]
        norm = np.linalg.norm(b1)
        print(f"  Iteration {iteration + 1} | norm(force)={norm:1.2e}")

        if norm < np.sqrt(np.finfo(float).eps):
            forces[increment] = solid.vector[dof.move].sum()
            break

import meshio
import pyvista

mesh = meshio.Mesh(
    points=cube.points + field.values,
    cells=[("hexahedron", cube.cells)],
    point_data={"displacement": field.values},
)
pyvista.wrap(mesh).plot()

plt.plot(moves, forces / 1000)
plt.xlabel("Displacement $d_1$ in mm")
plt.ylabel("Force $F_1$ in kN")