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