Full Finite Element Solver in 100 Lines of Python

Introduction

Back in the day when I started to study the finite element method at MIT, I was taught the material using a bottom-up approach. That is, the focus was on the theory behind the method. This is, of course, a useful and logical approach, but in my opinion it also makes the topic unnecessarily abstract and hard to follow. To make things worse, the code examples were provided in Fortran 77, one of the worst languages I have seen (I much prefer Fortran 90 and later versions).

I think a much better approach is to start with a full Finite Element (FE) solver written in a modern language, like Python, and then work backwards to explain how the code actually work. It is really amazing how about 100 lines of Python code can really solve advanced and interesting real problems. In this tutorial series I will show how this works.

Goals of Python FE Solver

To limit the amount of code I will focus this example on the following:

  • 2D plane strain
  • Linear FE solver
  • Linear elastic material
  • Linear full integration elements
  • Any geometry
  • Any boundary conditions
  • Any boundary forces

Problem Formulation

Consider a strip with a length of 50 mm, a width of 10 mm, and thickness of 1 mm. This strip is deformed in plane-strain tension in the vertical (2-dir) using a force of 200 N. The bottom of the strip is held fixed.

The elastic properties of the strip material are E=100 MPa, and \(\nu=0.48\).

The goal is to determine the displacement field after the load had been applied.

Approximate Answer

When possible, it is useful to approximate the solution using closed-form calculations. In this case I can use the linear elastic equation: \(\sigma_{ij} = 2\mu\varepsilon_{ij} + \lambda \varepsilon_{kk} \delta_{ij}\).

The Lame constants can be determined from the known elastic properties. By assigning plane strain, and assuming free displacement in the 1-dir, I can perform a quick calculation giving:

max top displacement = 8 mm.

This is an upper bound since the horizontal displacements is not free in the real applications.

Python Code

Here is the Python code that performs the FE calculation.

				
					#!/usr/bin/env python
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# Copyright 2022 Jorgen Bergstrom
# See also code from http://compmech.lab.asu.edu/codes.php
import numpy as np
import math
from matplotlib import pyplot as plt
def shape(xi):
	x,y = tuple(xi)
	N = [(1.0-x)*(1.0-y), (1.0+x)*(1.0-y), (1.0+x)*(1.0+y), (1.0-x)*(1.0+y)]
	return 0.25*np.array(N)
def gradshape(xi):
	x,y = tuple(xi)
	dN = [[-(1.0-y),  (1.0-y), (1.0+y), -(1.0+y)],
		  [-(1.0-x), -(1.0+x), (1.0+x),  (1.0-x)]]
	return 0.25*np.array(dN)
###############################
print('create mesh')
# input
mesh_ex = 9
mesh_ey = 49
mesh_lx = 10.0
mesh_ly = 50.0
# derived
mesh_nx      = mesh_ex + 1
mesh_ny      = mesh_ey + 1
num_nodes    = mesh_nx * mesh_ny
num_elements = mesh_ex * mesh_ey
mesh_hx      = mesh_lx / mesh_ex
mesh_hy      = mesh_ly / mesh_ey
nodes = []
for y in np.linspace(0.0, mesh_ly, mesh_ny):
	for x in np.linspace(0.0, mesh_lx, mesh_nx):
		nodes.append([x,y])
nodes = np.array(nodes)
conn = []
for j in range(mesh_ey):
	for i in range(mesh_ex):
		n0 = i + j*mesh_nx
		conn.append([n0, n0 + 1, n0 + 1 + mesh_nx, n0 + mesh_nx])
###############################
print ('material model - plane strain')
E = 100.0
v = 0.48
C = E/(1.0+v)/(1.0-2.0*v) * np.array([[1.0-v,     v,     0.0],
								      [    v, 1.0-v,     0.0],
								      [  0.0,   0.0,   0.5-v]])
###############################
print('create global stiffness matrix')
K = np.zeros((2*num_nodes, 2*num_nodes))
q4 = [[x/math.sqrt(3.0),y/math.sqrt(3.0)] for y in [-1.0,1.0] for x in [-1.0,1.0]]
B = np.zeros((3,8))
for c in conn:
	xIe = nodes[c,:]
	Ke = np.zeros((8,8))
	for q in q4:
		dN = gradshape(q)
		J  = np.dot(dN, xIe).T
		dN = np.dot(np.linalg.inv(J), dN)
		B[0,0::2] = dN[0,:]
		B[1,1::2] = dN[1,:]
		B[2,0::2] = dN[1,:]
		B[2,1::2] = dN[0,:]
		Ke += np.dot(np.dot(B.T,C),B) * np.linalg.det(J)
	for i,I in enumerate(c):
		for j,J in enumerate(c):
			K[2*I,2*J]     += Ke[2*i,2*j]
			K[2*I+1,2*J]   += Ke[2*i+1,2*j]
			K[2*I+1,2*J+1] += Ke[2*i+1,2*j+1]
			K[2*I,2*J+1]   += Ke[2*i,2*j+1]
###############################
print('assign nodal forces and boundary conditions')
f = np.zeros((2*num_nodes))
for i in range(num_nodes):
	if nodes[i,1] == 0.0:
		K[2*i,:]     = 0.0
		K[2*i+1,:]   = 0.0
		K[2*i,2*i]   = 1.0
		K[2*i+1,2*i+1] = 1.0
	if nodes[i,1] == mesh_ly:
		x = nodes[i,0]
		f[2*i+1] = 20.0
		if x == 0.0 or x == mesh_lx:
			f[2*i+1] *= 0.5
###############################
print('solving linear system')
u = np.linalg.solve(K, f)
print('max u=', max(u))
###############################
print('plotting displacement')
ux = np.reshape(u[0::2], (mesh_ny,mesh_nx))
uy = np.reshape(u[1::2], (mesh_ny,mesh_nx))
xvec = []
yvec = []
res  = []
for i in range(mesh_nx):
    for j in range(mesh_ny):
        xvec.append(i*mesh_hx + ux[j,i])
        yvec.append(j*mesh_hy + uy[j,i])
        res.append(uy[j,i])
t = plt.tricontourf(xvec, yvec, res, levels=14, cmap=plt.cm.jet)
plt.scatter(xvec, yvec, marker='o', c='b', s=2)
plt.grid()
plt.colorbar(t)
plt.axis('equal')
plt.show()

				
			

FEA Results

The figure to the right shows the FE results from running the Python code. The max displacement is 6.75 mm, which is consistent with the approximate calculation above.

This example showed the that Python FE code runs quickly and gives the expected results. 

In the next part of this series I will show another longer Python code example that uses an input file containing the geometry, displacements, and mesh.

Facebook
Twitter
LinkedIn

More to explore

1 thought on “Full Finite Element Solver in 100 Lines of Python”

  1. If you want the bottom of the specimen to be free to contract horizontally (i.e. to be on rollers) then replace lines 92 – 104 with:
    print(‘assign nodal forces and boundary conditions’)
    f = np.zeros((2*num_nodes))
    for i in range(num_nodes):
    x = nodes[i,0]
    y = nodes[i,1]
    if x==0 and y==0:
    K[2*i,:] = 0.0
    K[2*i+1,:] = 0.0
    K[2*i,2*i] = 1.0
    K[2*i+1,2*i+1] = 1.0
    if y==0:
    K[2*i+1,:] = 0.0
    K[2*i+1,2*i+1] = 1.0
    if y==mesh_ly:
    x = nodes[i,0]
    f[2*i+1] = 20.0
    if x == 0.0 or x == mesh_lx:
    f[2*i+1] *= 0.5

Leave a Comment