1D Finite Element Analysis in 100 Lines of Python Code (named tuple)

minimal-fea-1d_namedtuple.py

import numpy as np
from collections import namedtuple
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 namedtuple("Mesh", ("points", "cells"))(points, 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 namedtuple("Element", ("function", "gradient"))(function, gradient)


def Region(mesh, Element):
    quadrature_points = np.array([-1, 1]) / np.sqrt(3)
    quadrature_weights = np.ones([1, 1])
    element = Element(quadrature_points)
    gradient = mesh.points[mesh.cells] @ element.gradient
    return namedtuple("Region", ("mesh", "element", "gradient", "dx"))(
        mesh=mesh,
        element=element,
        gradient=element.gradient[..., None, :] / gradient,
        dx=quadrature_weights * gradient,
    )


def Field(region, values):
    return namedtuple("Field", ("region", "values"))(region, values)


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 namedtuple("SolidBody", ("force", "stiffness", "field"))(
        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 namedtuple("Assemble", ("vector", "matrix"))(
        vector=sparray((vector.sum(axis=-1).ravel(), vidx)).toarray().ravel(),
        matrix=sparray((matrix.sum(axis=-1).ravel(), midx)),
    )


mesh = Mesh(npoints=10001)
region = Region(mesh, Element)
field = Field(region, values=np.zeros_like(mesh.points))
dof = namedtuple("dof", ("fixed", "active"))(
    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-")