Search This Blog

Thursday, May 31, 2012

Draw a 3D surface figure with matplotlib from a data file

To draw a 3D surface figure with matplotlib, we generate 3 2D matrix - X, Y, and Z. All of them needs to have the same shape since the element of each matrix in the same position represents the point in 3D canvas.

For example, my data file looks
sec id value
5 1 10
10 1 12

I want to draw sec as x-axis, id as y-axis and value as z-axis, repsectively.When I know the range of sec and their intervals, and id,

# to generate 3D plots
from mpl_toolkits.mplot3d.axes3d import Axes3D
from matplotlib.pyplot as plt

# to generate 3D surface plots
from matplotlib import cm
from matplotlib.ticker import LinearLocator, formatStrFormatter

# to manipulate matrices
import numpy as np

# read data file
xs =[]
ys=[]
zs=[]
fd = open(inputfile,'r')
fd.readline() # remove the header
for line in fd:
    p = line.split()
    xs.append(float(p[0]))
    ys.append(float(p[1]))
    zs.append(float(p[2]))
fd.close()

#generate 2D arrays  
X = np.arange(0, float(timerange), 1)
Y = np.arange(0, int(float(id_range)), 1)
X, Y = np.meshgrid(X,Y) # this makes the shape of matrix X and Y the same.
Z = np.zeros(shape=X.shape)
for i in range(0, len(zs)):
    Z[ys[i],xs[i]] = zs[i]

# draw surface plot
fig=plt.figure()
ax=Axes3D(fig)
surf = ax.plot_surface(X,Y,Z,rstride=1, cstride=1, cmap=cm.jet, linewidth=0, antialiased=False)
ax.set_zlim(0.0, 100.0) # whatever  you want for the range of Z values

ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.2f'))
fig.colorbar(surf,ax=ax,shrink=0.5, aspect=5)
plt.show()

No comments:

Post a Comment