How to add thickness to flat polygons plotted on Matplotlib 3d axes?

Viewed 129

This is a follow-on to my previous question [here]Python: how to plot a series of regions in a stacked/layered way?.

I am using code along the lines of the following example to plot different "layers" of polygons in a 3-d plot.

import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection

fig = plt.figure()
ax = plt.gca(projection='3d')

poly1 = PolyCollection([[[-2,-2], [-2,2], [0,3], [2,2], [2,-2]]], color='red')
poly2 = PolyCollection([[[-1,-1], [-1,1], [1,1], [1,-1]]], color='blue')
poly3 = PolyCollection([[[0,0], [1,1], [1,0]]],color='green')

ax.add_collection3d(poly1, zs=0, zdir='z')
ax.add_collection3d(poly2, zs=1, zdir='z')
ax.add_collection3d(poly3, zs=2, zdir='z')

plt.show()

which gives the the following:

enter image description here

Now I want to add some thickness to each of these polygons, preferably with a option to specify arbitrary thickness, but would also go for "complete" thickness, by which I mean the sides go all the way down to z=0 and makes the thing look like a tiered wedding cake. I also would like the sides to not necessarily be the same color as the corresponding polygon. Is there a way to accomplish this?

I would also be happy with Python solutions outside of matplotlib. I may need to switch over to plotly for this project anyway.

1 Answers

I looked through matlotlib's documentation, but I don't think there is a convenient way to add thickness to a closed polygon, so I came up with a method that's probably not the most elegant but seems to work.

Basically what you need to do is loop through the pairs of coordinates on each polygon, then draw the face that extends down in the z-direction by whatever thickness you want. I used Poly3DCollection instead of PolyCollection so that we can work with 3D coordinates, and wrote convenience functions to convert your 2D coordinates to 3D coordinates, as well as calculate the 3D coordinates of each face based on the coordinates of the polygons.

By specifying thickness of 1 for each of the polygons you used, you can get a "tiered wedding cake" as you asked.

from mpl_toolkits.mplot3d import Axes3D
from matplotlib.collections import PolyCollection
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import matplotlib.pyplot as plt

fig = plt.figure()
ax = Axes3D(fig)
ax.set_xlim3d(-5, 5)
ax.set_ylim3d(-5, 5)
ax.set_zlim3d(-5, 5)

"""
This method takes a list of 2D coordinates and a z-coordinate, 
and returns a list of 3D coordinates with the z-coordinate added on 
"""
def make_3D_coordinate_list(coordinate_list, z_val):
    return [coordinate+[z_val] for coordinate in coordinate_list]

"""
This method iterates over pairs of coordinates in a coordinate list (where the first and last coordinate are identical)
and constructs a "face" by calculating two new points from each pair of coordinates, adding them to the Axes3D object

Arguments:
coordinate_list --- 2D list of lists of coordinates in 3 dimensions
face_color --- a string specifying the color of the face
thickness --- a number specifying how far in the z-direction the face of the prism should extend down
"""
def add_faces(coordinate_list, face_color, thickness):
    for i in range(len(coordinate_list) - 1):
        new_coordinate_list = [coordinate_list[i], coordinate_list[i+1]]
        new_coordinate_list.append([x + y for x, y in zip(coordinate_list[i+1], [0, 0, -1*thickness])])
        new_coordinate_list.append([x + y for x, y in zip(coordinate_list[i], [0, 0, -1*thickness])])
        new_coordinate_list.append(coordinate_list[i])
        print(f"new coordinate list: {new_coordinate_list}")
        ax.add_collection(Poly3DCollection([new_coordinate_list], color=face_color))

## make sure the coordinates cycle around from beginning to end
poly1_coordinate_list = [[-2,-2], [-2,2], [0,3], [2,2], [2,-2], [-2,-2]]
poly2_coordinate_list = [[-1,-1], [-1,1], [1,1], [1,-1], [-1,-1]]
poly3_coordinate_list = [[0,0], [1,1], [1,0], [0,0]]

poly1_3dcoordinate_list = make_3D_coordinate_list(poly1_coordinate_list, 0)
poly2_3dcoordinate_list = make_3D_coordinate_list(poly2_coordinate_list, 1)
poly3_3dcoordinate_list = make_3D_coordinate_list(poly3_coordinate_list, 2)

poly1 = Poly3DCollection([poly1_3dcoordinate_list], color='red')
poly2 = Poly3DCollection([poly2_3dcoordinate_list], color='blue')
poly3 = Poly3DCollection([poly3_3dcoordinate_list], color='green')

ax.add_collection3d(poly1)
ax.add_collection3d(poly2)
ax.add_collection3d(poly3)

add_faces(poly1_3dcoordinate_list, 'pink', 1)
add_faces(poly2_3dcoordinate_list, 'lightblue', 1)
add_faces(poly3_3dcoordinate_list, 'lightgreen', 1)

plt.show()

enter image description here

Related