How to color the surface base on color value meshgrid instead of Z-axis?

[lalit]

import numpy as np
from plotly.subplots import make_subplots

f=lambda x,y: x**2+y**2 # 2D parabolic function

x=np.linspace(-2,2,700) # define x intervals
y=np.linspace(-2,2,700) # define y intervals

X,Y=np.meshgrid(x,y) # creat x,y meshgrid

Z=f(X,Y) # calculate f(x,y) 
mask=Z>2 # create mask for circle
Z[mask]=np.nan # masking the circle 

colors=np.arctan2(Y,X) # it will calculate the radian for each of the point of meshgrid
colors[mask]=np.nan # masking the circle

fig=make_subplots(rows=1,cols=2,specs=[[{'type': 'surface'}, {'type': 'heatmap'}]])

# i want to color the surface based on the radian values not based on the Z values .. how to do it??
fig.add_surface(x=x, y=y ,z=Z, colorscale=px.colors.cyclical.HSV, row=1, col=1, showscale=False)

# i am coloring the heatmap based on the radian of each point
fig.add_heatmap(x=x, y=y, z=colors, colorscale=px.colors.cyclical.HSV, row=1,col=2,  showscale=False)

fig.update_layout(scene_aspectratio=dict(x=1, y=1,z=1))

fig.show()

So i want to plot the parabolic function f(X,Y)=X^2+X^2, …i plot it perfectly the problem is i want to color the surface not based on the Z-values but instead of Radian-values define as colors=np.arctan2(Y,X)

[lalit]

just got an property “surfacecolor”

import numpy as np
from plotly.subplots import make_subplots

f=lambda x,y: x**2+y**2 # 2D parabolic function

x=np.linspace(-2,2,700) # define x intervals
y=np.linspace(-2,2,700) # define y intervals

X,Y=np.meshgrid(x,y) # creat x,y meshgrid

Z=f(X,Y) # calculate f(x,y) 
mask=Z>2 # create mask for circle
Z[mask]=np.nan # masking the circle 

colors=np.arctan2(Y,X) # it will calculate the radian for each of the point of meshgrid
colors[mask]=np.nan # masking the circle

colorscale=["#A15AFF","#FF419E","#FF9D00","#A2FF29","#5DFF67","#50D0FF","#A15AFF"]

fig=make_subplots(rows=1,cols=2,specs=[[{'type': 'surface'}, {'type': 'heatmap'}]])

# i want to color the surface based on the radian values not based on the Z values .. how to do it??
fig.add_surface(x=x, y=y ,z=Z, surfacecolor=colors, colorscale=colorscale, row=1, col=1, showscale=False)

# i am coloring the heatmap based on the radian of each point
fig.add_heatmap(x=x, y=y, z=colors, colorscale=colorscale, row=1,col=2, showscale=False)

fig.update_layout(scene_aspectratio=dict(x=1, y=1,z=1))

fig.show()