3D Finite Element Analysis in 100 Lines of Python Code (Alternative Assembly)

minimal-fea-3d_alternative.py

import numpy as np
from types import SimpleNamespace as SN
from scipy.sparse import csr_matrix as sparse
from scipy.sparse.linalg import spsolve
import meshio

def Mesh(npoints, a=0, b=1):
    grid = np.linspace(a, b, npoints)
    points = np.pad(grid[:, None], ((0, 0), (0, 2)))
    cells = np.arange(npoints).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(npoints)]
        points = np.vstack([points + np.insert(np.zeros(2), dim + 1, h) for h in grid])
        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(points):
    r, s, t = points
    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, quadrature):
    element = Element(quadrature.points)
    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=quadrature.weights * np.linalg.det(dXdr.T).T,
    )

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

def Assemble(field, lmbda, mu):
    ascontiguous = lambda *args: [np.ascontiguousarray(x) for x in args]
    sym = lambda x: (x + np.einsum("ij...->ji...", x)) / 2
    dya = lambda x, y: np.einsum("ij...,kl...->ijkl...", x, y)
    cdya_ik = lambda x, y: np.einsum("ik...,jl...->ijkl...", x, y)
    cdya_il = lambda x, y: np.einsum("il...,kj...->ijkl...", x, y)
    cdya = lambda x, y: (cdya_ik(x, y) + cdya_il(x, y)) / 2

    dhdX = ascontiguous(field.region.gradient)[0]
    dudX = field.gradient(field.values)
    dV = field.region.dx
    ε = sym(dudX)
    I = np.eye(3)[..., None, None]

    σ = 2 * mu * ε + lmbda * np.trace(ε) * I
    C = 2 * mu * cdya(I, I) + lmbda * dya(I, I)

    vector = np.einsum("ajcq,ijcq,cq->aic", dhdX, σ, dV, optimize=True)
    matrix = np.einsum("ajcq,ijklcq,blcq,cq->aibkc", dhdX, C, dhdX, dV, optimize=True)

    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=sparse((vector.transpose([2, 0, 1]).ravel(), vidx)),
        matrix=sparse((matrix.transpose([4, 0, 1, 2, 3]).ravel(), midx)),
    )

mesh = Mesh(npoints=16, a=2, b=5)
quadrature = SN(
    points=np.concatenate(np.meshgrid([-1, 1], [-1, 1], [-1, 1])).reshape(3, -1)
    / np.sqrt(3),
    weights=np.ones(8),
)
region = Domain(mesh, Hexahedron, quadrature)
field = VectorField(region, values=np.zeros_like(mesh.points))

extforce = np.zeros_like(mesh.points)
extforce[:, 0][mesh.points[:, 0] == 5] = -3**2 / 4 / 16**2

dofs = np.arange(mesh.points.size).reshape(mesh.points.shape)
dof = SN(fixed=dofs[mesh.points[:, 0] == 2].ravel())
dof.active = np.delete(dofs.ravel(), dof.fixed)

b = extforce.ravel()[dof.active]
for iteration in range(8):
    system = Assemble(field, lmbda=1.0, mu=2.0)
    A = system.matrix[dof.active, :][:, dof.active]
    field.values.ravel()[dof.active] += spsolve(A, b).ravel()
    b = (extforce.ravel() - system.vector.toarray().ravel())[dof.active]
    norm = np.linalg.norm(b)
    print(f"Iteration {iteration + 1} | norm(force)={norm:1.2e}")
    if norm < np.sqrt(np.finfo(float).eps):
        break

meshio.Mesh(
    mesh.points, [("hexahedron", mesh.cells)], point_data={"displacement": field.values}
).write("result.vtk")