Pyplot 3d, customize colors

Question

I'm plotting the spread using pyplot:

fig = plt.figure(figsize = (12, 12))
ax = fig.add_subplot(projection = '3d')

x = dataframe.AGE # X
y = dataframe.BMI # Y
z = dataframe.BP # Z

ax.scatter(x, y, z)

ax.set_xlabel("AGE")
ax.set_ylabel("BMI")
ax.set_zlabel("Bermuda plan (BP)")

plt.show()

How can I specify the color for these parameters? So x is red, y is blue, z is green

Answer 1

You can specify a matrix of n x 3 elements where n is the number of points, and each row represents an RGB color. Then, R, G, B are going to take values from zero (pure black) to one (red, or green or blue).

So, we need to scale the x, y, z data into the range [0, 1] and create the color matrix:

import numpy as np
import matplotlib.pyplot as plt

def shift_and_scale(t):
    t -= t.min()
    return t / t.max()

x, y, z = np.mgrid[-10:10:10j, -10:10:10j, -10:10:10j]
x, y, z = [t.flatten() for t in [x, y, z]]
colors = np.stack([shift_and_scale(t.copy()) for t in [x, y, z]]).T

fig = plt.figure()
ax = fig.add_subplot(projection="3d")
ax.scatter(x, y, z, c=colors)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")

As you can see, dots become more red when x increases. They become more green when y increases, and they become more blue when z increases.