Search code examples
pythonmatplotlibmatplotlib-3d

Color plot from 3D arrays

I have a numpy array of arrays such as: Y[:,0] is a list representing the x axis coordinates. For each point on this axis, marked by the index indx, there is Y[indx,1][0] is a list of y axis coordinates, and Y[indx,1][1] is a list of of the z axis data I can now use the code below to plot an array of lines in 3D. However, I would like to plot a 2D plot where the z data represents a color map on the x and y coordinates (something that looks like the output of a wavelet transform plot).

#dummy data
Y=[0]*10
for i in range(10):
    temp=np.array([[indx,np.cos(indx)] for indx in range(50)])
    Y[i]=np.array([i,temp.transpose()])
Y=np.array(Y,dtype=object)

#plotting in 3D
ax = plt.figure().add_subplot(projection='3d')
for indx in range(len(Y[:,0])):
    ax.plot(Y[indx,1][0], Y[indx,1][1], zs=Y[:,0][indx], zdir='x')

This gives me the 3D plot. How can I get the 2D plot as I described?

Solution

  • Are the y-coordinates going to be the same for every x-coordinate e.g. Y[0,1][0] = Y[1,1][0]? Also, will the y-coordinates be evenly spaced? The solution below assumes both of these to be true.

    Solution

    Create an array of z-values associated with each (x,y) location. Then call plt.imshow(...) using the new array as an argument. You can even include a colorbar!

    example output

    Note: The colormap is nothing impressive when using the dummy data, but it appropriately shows the plot you asked for.

    Code

    #dummy data
Y=[0]*10
for i in range(10):
    temp=np.array([[indx,np.cos(indx)] for indx in range(50)])
    Y[i]=np.array([i,temp.transpose()],dtype=object)

Y=np.array(Y,dtype=object)

xVals = Y.T[0] # x axis coordinates (not used, but a useful piece of information)
yVals = Y.T[1][1][0].astype(int) # y axis coordinates as an integer array (also not used)

# create the array of zVals
zVals = np.concatenate(Y.T[1])[1::2]

#plotting in 3D
fig = plt.figure()
ax3d = fig.add_subplot(211,projection='3d')

for indx in range(len(Y[:,0])):
    ax3d.plot(Y[indx,1][0], Y[indx,1][1], zs=Y[:,0][indx], zdir='x')

#plotting in 2D
ax = fig.add_subplot(212)
im = plt.imshow(zVals)
cb = plt.colorbar()