Applying colormaps to custom axis in Matplotlib 3D surface

Question

I have timeseries data which I've segmented into hundreds of chunks. I solved the autocorrelation for each segment and plotted them:

# plot superimposed
fig = plt.figure()
color = iter(plt.cm.Set2(np.linspace(0,1,num_segs)))
seg_iterator = df.iterrows() 
for index, seg in seg_iterator: # iterate over dataframe
    c=next(color)
    sns.plt.plot(seg, color=c)

Next, I plotted them as a 3D surface:

# plot as a surface
surfacefig = plt.figure()
surfaceax = surfacefig.gca(projection='3d')
X = np.arange(LAGS+1)
Y = np.arange(num_segs)
X, Y = np.meshgrid(X, Y)
surfaceax.plot_surface(X, Y, df, cmap=plt.cm.Set2)
plt.show()

How can I map colors to row index (rather than z-values)? I'd like to preserve the colors of the lines.

---

Update with result:

# updated lines. Make sure XX and YY are floats
surf = surfaceax.plot_surface(XX, YY, df, shade=False,
             facecolors=plt.cm.Set2((YY-YY.min()) / (YY.max()-YY.min())), 
             cstride=1, rstride=5, alpha=0.7)
plt.draw() # you need this to get the edge color
line = np.array(surf.get_edgecolor())
surf.set_edgecolor(line*np.array([0,0,0,0])+1)

Answer 1

You can try this:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

X = np.linspace(-np.pi, np.pi, 200, endpoint=True)
Y = np.linspace(-np.pi, np.pi, 200, endpoint=True)
XX, YY = np.meshgrid(X,Y)
Z = np.cos(XX)*np.cos(YY)

fig = plt.figure()
ax1 = plt.subplot2grid((1,2), (0,0), projection='3d')
ax2 = plt.subplot2grid((1,2), (0,1), projection='3d')
surf = ax1.plot_surface(XX, YY, Z,
                        cmap=plt.cm.Set2)
surf2 = ax2.plot_surface(XX, YY, Z, shade=False,
                         facecolors=plt.cm.Set2((XX-XX.min())/(XX.max()-XX.min()))
                         )

Where on the second plot, you set the facecolors as being function of XX, instead of Z by default. You need to rescale your XX values between 0 and 1 or the colormap will be saturated outside 0 and 1. You also need to remove the shade which is removed when yous use cmap (in the first plot).

However, for some unknown reasons, the lines disappear.

You can add them back with:

plt.draw() # you need this to get the edge color
lines = np.array(surf2.get_edgecolor())
surf2.set_edgecolor(lines*np.array([0,0,0,0])+1) # make lines white, and keep alpha==1. It's an array of colors like this: [r,g,b,alpha]

It gives:

HTH