3D scatter of points with colors as a 4th var

Question

I'm trying to scatter a dataset in 3D and for each point using a scalar variable as a color, to visualize better some phenomena. In order to keep things simple, suppose i want that each point plotted has a color that in some way represents another 3D array (density,temperature, energy,etc..).

How can I do this? I would like something like the following:

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

#some dummy data
x = np.linspace(-1,1,3, endpoint = True )
y  = np.linspace(-1,1,3, endpoint = True )
z = np.linspace(-1,1,3, endpoint = True )
X,Y,Z = np.meshgrid(x,y,z,indexing = 'ij')
E = X**2+Y**2+Z**2
fig = plt.figure()
ax = fig.add_subplot(111,projection = '3d')
ax.scatter(X,Y,Z, color = VV)
plt.show()

here VV would be an array that in some way is proportional to E and that gives to each (x,y,z) a color according to E[x,y,z]. The final part would be to put a colorbar that helps to read the approximate color value. I searched already on the forum but the answers I found were not adaptable to 3D datasets and i'm currently stuck.

Answer 1

If I have understood the problem, you can just set c=E in your call to scatter and Matplotlib will automatically normalise it to fit the colour scale. scatter returns a PathCollection which is a type of mappable so can be passed to colorbar.

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

#some dummy data
x = np.linspace(-1,1,3, endpoint = True )
y  = np.linspace(-1,1,3, endpoint = True )
z = np.linspace(-1,1,3, endpoint = True )
X,Y,Z = np.meshgrid(x,y,z,indexing = 'ij')
E = X**2+Y**2+Z**2

fig = plt.figure()
ax = fig.add_subplot(111, projection = '3d')
pc = ax.scatter(X, Y, Z, c=E)
fig.colorbar(pc)
plt.show()