Plotting using PolyCollection in matplotlib

Karthik Venkatesh picture Karthik Venkatesh · Jul 2, 2015 · Viewed 9k times · Source

I am trying to plot a 3 dimensional plot in matplotlib. I have to plot Frequency vs Amplitude Distribution for four (or multiple) Radii in a single 3D plot. I was looking at PolyCollection command available in matplotlib.collections and I also went through the example but I do not know how to use the existing data to arrive at the plot.

The dimensions of the quantities that I have are,

Frequency : 4000 x 4, Amplitude : 4000 x 4, Radius : 4

I would like to plot something like, enter image description here

With X axis being Frequencies, Y axis being Radius, and Z axis being Amplitudes. How do I go about solving this problem?

Answer

Ajean picture Ajean · Jul 2, 2015

PolyCollection expects a sequence of vertices, which matches your desired data pretty well. You don't provide any example data, so I'll make some up for illustration (my dimension of 200 would be your 4000 .... although I might consider a different plot than this if you have so many data points):

import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection
from mpl_toolkits.mplot3d import axes3d
import numpy as np

# These will be (200, 4), (200, 4), and (4)
freq_data = np.linspace(0,300,200)[:,None] * np.ones(4)[None,:]
amp_data = np.random.rand(200*4).reshape((200,4))
rad_data = np.linspace(0,2,4)

verts = []
for irad in range(len(rad_data)):
    # I'm adding a zero amplitude at the beginning and the end to get a nice
    # flat bottom on the polygons
    xs = np.concatenate([[freq_data[0,irad]], freq_data[:,irad], [freq_data[-1,irad]]])
    ys = np.concatenate([[0],amp_data[:,irad],[0]])
    verts.append(list(zip(xs, ys)))

poly = PolyCollection(verts, facecolors = ['r', 'g', 'c', 'y'])
poly.set_alpha(0.7)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# The zdir keyword makes it plot the "z" vertex dimension (radius)
# along the y axis. The zs keyword sets each polygon at the
# correct radius value.
ax.add_collection3d(poly, zs=rad_data, zdir='y')

ax.set_xlim3d(freq_data.min(), freq_data.max())
ax.set_xlabel('Frequency')
ax.set_ylim3d(rad_data.min(), rad_data.max())
ax.set_ylabel('Radius')
ax.set_zlim3d(amp_data.min(), amp_data.max())
ax.set_zlabel('Amplitude')

plt.show()

Most of this is straight from the example you mention, I just made it clear where your particular datasets would lie. This yields this plot: example PolyCollection plot