[FIXED] Displaying 4th variable as colormap on 3D plot in Python

Issue

Task

I am trying to generate a plot that would represent my data in 4 dimensions. The option I am going for is a 3D plot where the 4th dimension is represented with a colormap. My data has some NaN values.

My attempt

Using an arbitrary set of data, my 4 variables are X, Y, Z, C. Where I want X, Y and Z to lay on their respective axes of the 3D plot and C to define the colormap. Both Z and C have some matching NaN values.

import numpy as np
import matplotlib.pyplot as plt

# generate some data
X, Y = np.meshgrid(np.linspace(1, 10, 20), np.linspace(1, 10, 20))  # 20 by 20 grid
Z = np.linspace(0, 10, 400).reshape((20, 20))  # data to plot on the vertical axis
Z.ravel()[np.random.choice(Z.size, 10, replace=False)] = np.nan  # some values are nan
C = np.cos(X) - np.sin(Y)  # 4th variable, to be represented in 4th dimension using colormap
C[np.isnan(Z)] = np.nan  # corresponding values are nan

I attempted to generate a 3D colormap of X, Y and C first, storing its colormap in variable fc. This code returns me the first colormap (unfortunately, StackOverflow is not letting me post images of the plots).

fig = plt.figure()
ax = fig.add_subplot(221, projection='3d')
# plot the surface for 4th dimension to get its colormap (facecolors)
surf = ax.plot_surface(X, Y, C, cmap='turbo')
fc = surf.get_facecolors()

Afterwards, I generated the desired 3D colormap of X, Y and Z, setting the colormap argument to fc (with some reshaping of fc on the way, as get_facecolors() returns a flat tuple of RGBA arrays), which returns me the second colormap.

# get_facecolors() returns a tuple of 1D arrays, but the plot_surface() facecolors argument requires a 3D array
fc = np.array(fc)
fc = fc.reshape((19, 19, 4))
ax = fig.add_subplot(222, projection='3d')
ax.plot_surface(X, Y, Z, facecolors=fc, cmap='turbo')

plt.show()

The problem

So as can be seen from the two plots, the colormap from the first plot gets mixed over before getting applied to the second plot. Reshaping in fc = fc.reshape((19, 19, 4)) is necessary as those are the required dimensions for the facecolors, but my guess it is this reshaping that causes mixing of colors.

I have attempted other solutions proposed in This question from a few years ago but no luck.

Any advice on how to get the colormap of one plot to translate to the other will be much appreciated! As well as any suggestions of easier/better options of representing 4 variables on a plot.

Thanks in advance!

Solution

You need to use Normalize and a color map to compute the facecolors, like this:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import Normalize

# generate some data
X, Y = np.meshgrid(np.linspace(1, 10, 20), np.linspace(1, 10, 20))  # 20 by 20 grid
Z = np.linspace(0, 10, 400).reshape((20, 20))  # data to plot on the vertical axis
Z.ravel()[np.random.choice(Z.size, 10, replace=False)] = np.nan  # some values are nan
C = np.cos(X) - np.sin(Y)  # 4th variable, to be represented in 4th dimension using colormap
C[np.isnan(Z)] = np.nan  # corresponding values are nan

cmap = cm.turbo
norm = Normalize(vmin=np.nanmin(C), vmax=np.nanmax(C))
fc = cmap(norm(C))

fig = plt.figure()
ax = fig.add_subplot(221, projection='3d')
# plot the surface for 4th dimension to get its colormap (facecolors)
surf = ax.plot_surface(X, Y, C, facecolors=fc)

ax = fig.add_subplot(222, projection='3d')
ax.plot_surface(X, Y, Z, facecolors=fc)

Answered By - Davide_sd