Full Finite Element Solver in 100 Lines of Python

https://polymerfem.com/full-finite-element-solver-in-100-lines-of-python/

Full Finite Element Solver in 100 Lines of Python - PolymerFEM.com

 

#!/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)

In [15]:

###############################
print('create mesh')
# input
# How many element we expected by each direction
mesh_ex = 9
mesh_ey = 49
# How Big is the block which we are simulation
mesh_lx = 10.0
mesh_ly = 50.0

# derived
# Node number : 500ea
mesh_nx      = mesh_ex + 1
mesh_ny      = mesh_ey + 1
num_nodes    = mesh_nx * mesh_ny
num_elements = mesh_ex * mesh_ey

# Element size 
mesh_hx      = mesh_lx / mesh_ex
mesh_hy      = mesh_ly / mesh_ey
nodes = []

# Assign the elements
# Assigne node coord
# nodes : [[x_coord1,y_coord1], [x_coord2,y_coord12, ...]
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)


# Define Element - Node Connectivity
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])

create mesh

In [16]:

# Material definition - plane strain linear elastic
###############################
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]])
###############################

material model - plane strain

In [17]:

# Sets up the global Stiffness Matrix (of this particular problem)
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,:] # Coords of every node of a one element
	Ke = np.zeros((8,8))
    
    # Ke : Stiffness of Element
	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)
    
    # K : Global Stiffness Matrix
	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]

create global stiffness matrix

In [79]:

# Assign nodal forces boundary condition
###############################
print('assign nodal forces and boundary conditions')

f = np.zeros((2*num_nodes))
for i in range(num_nodes):
    
    # Give Fixed B.C Where Y Coord is Zero
	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
    
    # Give Load Where Y Coord is mesh_ly
	if nodes[i,1] == mesh_ly:
		x = nodes[i,0]
		f[2*i+1] = 20.0
        
        # Half Load where end of x-axis (why..??)
		if x == 0.0 or x == mesh_lx:
			f[2*i+1] *= 0.5

assign nodal forces and boundary conditions

In [97]:

# Solve the Displacement field (for the given condition specified)
print('solving linear system')
u = np.linalg.solve(K, f)
print('max u=', max(u))

solving linear system
max u= 6.745400245724015

In [ ]:

# Plotting (Postprocess)
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()