Plotly - different color surfaces

Question

I'm trying to plot several surfaces, each of a different color, in Plotly for Python.

Specifically, a surface shows the predicted reward function for taking an action at different points in phase space. Since I have several possible actions at each point, each is a different surface. I'd like to color each surface uniquely, but independent of the x,y, or z coordinate.

I've tried to follow answer in R, but I can't figure out what I've done wrong. I always get the same blue color. Since I'm using PyPlot in other parts of my code, I'm choosing colors from the default matplotlib tableau.

Here's a basic example with toy data.

import matplotlib.pyplot as plt
import numpy as np
import plotly.graph_objs as go
import plotly.offline as off

off.init_notebook_mode()

make_int = np.vectorize(int)
cmap = plt.get_cmap("tab10")

saddle = np.array([[x**2-y**2 for x in np.arange(-10,11)] for y in np.arange(-10,11)])
paraboloid = np.array([[x**2 + y**2-100 for x in np.arange(-10,11)] for y in np.arange(-10,11)])

mycolors_a = make_int(256*np.array(cmap(1)[0:3])).reshape((1, 1,-1)).repeat(21, axis = 0).repeat(21, axis =1)
mycolors_b = make_int(256*np.array(cmap(2)[0:3])).reshape((1, 1,-1)).repeat(21, axis = 0).repeat(21, axis =1)
trace_a = go.Surface(z = saddle, surfacecolor = mycolors_a, opacity = .7, showscale = False, name = "Trace A")
trace_b = go.Surface(z = paraboloid, surfacecolor = mycolors_b, opacity = .7, showscale = False, name = "Trace B")

data = [trace_a, trace_b]
off.iplot(data)

Produces the following:

I should see a blue saddle and an orange paraboloid, but I don't. Note that even if I change the argument to cmap, I always get the same blue color. Thanks for your help!

Answer 1

The documentation is a bit cryptic here.

surfacecolor

(list, numpy array, or Pandas series of numbers, strings, or datetimes.)

Sets the surface color values, used for setting a color scale independent of z.

I never managed to put a list of strings, i.e. color values like 'rgb(0.3, 0.5, 0)', or RGB tuples in it.

But you can define your own color scale with the needed colors.

colorscale = [[0, 'rgb' + str(cmap(1)[0:3])], 
              [1, 'rgb' + str(cmap(2)[0:3])]]

and then provide a numeric array with the same dimensions as your plotted values.

colors_saddle = np.zeros(shape=saddle.shape)    

All values are set to 0 and will therefore map to the first color in your colorscale. The same for the next color.

In addition you need to set cmax and cmin manually.

Complete code

import numpy as np
import matplotlib.pyplot as plt
import plotly.graph_objs as go
import plotly.offline as off


off.init_notebook_mode()

make_int = np.vectorize(int)
cmap = plt.get_cmap("tab10")

saddle = np.array([[x**2-y**2 for x in np.arange(-10,11)] for y in np.arange(-10,11)])
paraboloid = np.array([[x**2 + y**2-100 for x in np.arange(-10,11)] for y in np.arange(-10,11)])

colors_saddle = np.zeros(shape=saddle.shape)    
colors_paraboloid = np.ones(shape=paraboloid.shape)    

colorscale = [[0, 'rgb' + str(cmap(1)[0:3])], 
              [1, 'rgb' + str(cmap(2)[0:3])]]

trace_a = go.Surface(z=saddle, 
                     surfacecolor=colors_saddle, 
                     opacity=.7, 
                     name="Trace A",
                     cmin=0,
                     cmax=1,
                     colorscale=colorscale)
trace_b = go.Surface(z=paraboloid, 
                     surfacecolor=colors_paraboloid, 
                     opacity=.7, 
                     name="Trace B", 
                     cmin=0,
                     cmax=1,
                     showscale=False,
                     colorscale=colorscale)

data = [trace_a, trace_b]
off.iplot(data)