How to color the point of a scatter plot separately?

Question

I would like to plot a three variable function in the following way. First I plot evenly distributed points in the 3d figure with scatter plot, than I want to use a method to color every point in the cube according the value of the function. The function is f(x,y,z)=(x^2+y^2-z^2)^2.

So far I managed to fill the plot with little balls evenly, but the coloring does not want to work. I tried to look up how to handle colormap, but the results was either the most basic or it was too hardcore.

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np
import math as mt

fig = plt.figure()
ax = fig.gca(projection='3d')

# Make data.
X = np.linspace(-5, 5, 6)
Y = np.linspace(-5, 5, 6)
Z = np.linspace(-5, 5, 6)
X, Y, Z = np.meshgrid(X, Y, Z)
# Plot the surface.
t = (X**2+Y**2-Z**2)**2
surf = ax.scatter(X, Y, Z, c=t, cmap=cm.viridis)


ax.set_zlim(-5, 5)
ax.zaxis.set_major_locator(LinearLocator(5))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')
ax.set_zlabel('Z axis')

# Add a color bar which maps values to colors.
#fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

I think what I tried wont work even if there is no error, since the pairing of colors with point is not handled. I thought about using plot instead of scatter, but again I did not find any useful material on colormaps.

Answer 1

I have three differences from what you have done

1. the function is simpler, so that we can check visually the correctness of the scatter plot
2. I have used ax.scatter3D in place of ax.scatter
3. the 3D matrix with the function values has to be flattened * t.flatten() )

and here the code and the plot

In [31]: from mpl_toolkits.mplot3d import Axes3D 
    ...: import matplotlib.pyplot as plt 
    ...: from matplotlib import cm 
    ...: from matplotlib.ticker import LinearLocator, FormatStrFormatter 
    ...: import numpy as np 
    ...: import math as mt 
    ...:  
    ...: fig = plt.figure() 
    ...: ax = fig.gca(projection='3d') 
    ...:  
    ...: # Make data. 
    ...: X = np.linspace(-5, 5, 6) 
    ...: Y = np.linspace(-5, 5, 6) 
    ...: Z = np.linspace(-5, 5, 6) 
    ...: X, Y, Z = np.meshgrid(X, Y, Z) 
    ...: # Plot the surface. 
    ...: t = Z 
    ...: ax.scatter3D(X, Y, Z, c=t.flatten())        

Note that, to put in evidence the layers, I have rotated the default view and also that Matplotlib draws the points using a sort of aerial perspective .