How can I plot 2d FEM results using matplotlib?

After thinking for a while, the following code is one of the easiest ways of plotting a FEM mesh (with a nodal scalar field) using matplotlib.

This solution is based on matplotlib.pyplot.tricontourf(). Unfortunately, matplotlib does not have an easy way to plot filled contours if you have quadrangles or higher order elements in your finite element mesh. In order to plot the contours, all elements must first be "cut" into triangles, for example, a quadrangle can be split or cut into 2 triangles, and so on...

A custom method for plotting the mesh lines must also be employed, since matplotlib.pyplot.tricontourf() only works with a triangular grid/mesh. For this, matplotlib.pyplot.fill() was used.

Here is the full code with a simple example:

import matplotlib.pyplot as plt
import matplotlib.tri as tri

# converts quad elements into tri elements
def quads_to_tris(quads):
    tris = [[None for j in range(3)] for i in range(2*len(quads))]
    for i in range(len(quads)):
        j = 2*i
        n0 = quads[i][0]
        n1 = quads[i][1]
        n2 = quads[i][2]
        n3 = quads[i][3]
        tris[j][0] = n0
        tris[j][1] = n1
        tris[j][2] = n2
        tris[j + 1][0] = n2
        tris[j + 1][1] = n3
        tris[j + 1][2] = n0
    return tris

# plots a finite element mesh
def plot_fem_mesh(nodes_x, nodes_y, elements):
    for element in elements:
        x = [nodes_x[element[i]] for i in range(len(element))]
        y = [nodes_y[element[i]] for i in range(len(element))]
        plt.fill(x, y, edgecolor='black', fill=False)

# FEM data
nodes_x = [0.0, 1.0, 2.0, 0.0, 1.0, 2.0, 3.0, 0.0, 1.0, 2.0, 3.0]
nodes_y = [0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0]
nodal_values = [1.0, 0.9, 1.1, 0.9, 2.1, 2.1, 0.9, 1.0, 1.0, 0.9, 0.8]
elements_tris = [[2, 6, 5], [5, 6, 10], [10, 9, 5]]
elements_quads = [[0, 1, 4, 3], [1, 2, 5, 4], [3, 4, 8, 7], [4, 5, 9, 8]]
elements = elements_tris + elements_quads

# convert all elements into triangles
elements_all_tris = elements_tris + quads_to_tris(elements_quads)

# create an unstructured triangular grid instance
triangulation = tri.Triangulation(nodes_x, nodes_y, elements_all_tris)

# plot the finite element mesh
plot_fem_mesh(nodes_x, nodes_y, elements)

# plot the contours
plt.tricontourf(triangulation, nodal_values)

# show
plt.colorbar()
plt.axis('equal')
plt.show()

Which outputs:

enter image description here

Just by changing the FEM data (nodes, nodal values, elements) the above code can be used for more complicated meshes, however, the code is only prepared to deal with meshes containing triangles and quadrangles:

enter image description here

You may notice that for large meshes, matplotlib will get slow. Also with matplotlib you can't visualize 3D-elements. So for better efficiency and more functionality, consider using instead VTK, for example.


A PolyCollection is a ScalarMappable. It can have an array of values, a colormap and a normalization set. Here you would supply the stresses array to the PolyCollection and choose some colormap to use. The rest is rearanging the function a bit, such that can take the additional data as input and creates a colorbar.

import matplotlib.pyplot as plt
import matplotlib.collections
import numpy as np


def showMeshPlot(nodes, elements, values):

    y = nodes[:,0]
    z = nodes[:,1]

    def quatplot(y,z, quatrangles, values, ax=None, **kwargs):

        if not ax: ax=plt.gca()
        yz = np.c_[y,z]
        verts= yz[quatrangles]
        pc = matplotlib.collections.PolyCollection(verts, **kwargs)
        pc.set_array(values)
        ax.add_collection(pc)
        ax.autoscale()
        return pc

    fig, ax = plt.subplots()
    ax.set_aspect('equal')

    pc = quatplot(y,z, np.asarray(elements), values, ax=ax, 
             edgecolor="crimson", cmap="rainbow")
    fig.colorbar(pc, ax=ax)        
    ax.plot(y,z, marker="o", ls="", color="crimson")

    ax.set(title='This is the plot for: quad', xlabel='Y Axis', ylabel='Z Axis')

    plt.show()

nodes = np.array([[0,0], [0,0.5],[0,1],[0.5,0], [0.5,0.5], [0.5,1], [1,0], 
                  [1,0.5],[1,1]])
elements = np.array([[0,3,4,1],[1,4,5,2],[3,6,7,4],[4,7,8,5]])
stresses = np.array([1,2,3,4])

showMeshPlot(nodes, elements, stresses)

enter image description here


I think your best option is to use tricontour. You already have the triangulation, right?

It creates plots like this one: enter image description here

(from here)

https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.tricontour.html