1D Finite Element Analysis in 100 Lines of Python Code

minimal-fea-1d.py

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


def Mesh(npoints):
    points = np.linspace(0, 1, npoints)
    cells = np.arange(len(points)).repeat(2)[1:-1].reshape(-1, 2)
    return SN(points=points, cells=cells)


def Element(points):
    function = np.array([1 - points, 1 + points]) / 2
    gradient = np.array([-np.ones_like(points), np.ones_like(points)]) / 2
    return SN(function=function, gradient=gradient)


def Region(mesh, Element, quadrature):
    element = Element(quadrature.points)
    gradient = mesh.points[mesh.cells] @ element.gradient
    return SN(
        mesh=mesh,
        element=element,
        gradient=element.gradient[..., None, :] / gradient,
        dx=quadrature.weights * gradient,
    )


def Field(region, values):
    return SN(values=values, region=region)


def SolidBody(field, E, A):
    def gradient(field):
        region = field.region
        return (field.values[region.mesh.cells] * region.gradient.T).sum(-1).T

    stretch = 1 + gradient(field)
    return SN(
        force=E * A * (stretch - 1 / stretch**2),
        stiffness=E * A * (1 + 2 / stretch**3),
        field=field,
    )


def Assemble(solid):
    dx = solid.field.region.dx
    gradv = solid.field.region.gradient[:, None]
    gradu = solid.field.region.gradient[None, :]

    vector = gradv * solid.force * dx
    matrix = gradv * solid.stiffness * gradu * dx

    vidx = (mesh.cells.T.ravel(), np.zeros(mesh.cells.size, dtype=int))
    midx = (
        np.tile(mesh.cells.T.ravel(), len(mesh.cells.T)),
        np.repeat(mesh.cells.T.ravel(), len(mesh.cells.T)),
    )

    return SN(
        vector=sparray((vector.sum(axis=-1).ravel(), vidx)).toarray().ravel(),
        matrix=sparray((matrix.sum(axis=-1).ravel(), midx)),
    )


mesh = Mesh(npoints=10001)
quadrature = SN(points=np.array([-1, 1]) / np.sqrt(3), weights=np.ones([1, 1]))
region = Region(mesh, Element, quadrature)
field = Field(region, values=np.zeros_like(mesh.points))

dof = SN(fixed=np.arange(1), active=np.arange(len(mesh.points))[1:])
extforce = np.append(np.zeros(len(mesh.points) - 1), 1)

for iteration in range(8):
    system = Assemble(SolidBody(field, E=1, A=1))

    A = system.matrix[dof.active, :][:, dof.active]
    b = extforce[dof.active] - system.vector[dof.active]

    norm = np.linalg.norm(b)
    print(f"Iteration {iteration + 1} | norm(force)={norm:1.2e}")
    if norm < 1e-3:
        break

    x = spsolve(A, b)
    field.values[dof.active] += x

plt.plot(field.values + mesh.points, np.zeros_like(mesh.points), "C0o-")