(python) plot 3d surface with colormap as 4th dimension, function of x,y,z

This code is based on the trisurf demo http://matplotlib.org/examples/mplot3d/trisurf3d_demo.html

I added a function make_colormap() based on the SO Create own colormap using matplotlib and plot color scale

Also added a sequence w=tan(-x*y) that generates a colour map based on that function, in the gray scale.
You can play with the construction of the cdict to add more colors to it but I think gray scale makes a good proof of concept...

Sorry I couldn't work directly with your example, due to lack of minimal working code.

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

###################

def make_colormap(seq):
    """Return a LinearSegmentedColormap
    seq: a sequence of floats and RGB-tuples. The floats should be increasing
    and in the interval (0,1).
    """
    #%
    cdict = {'red': [], 'green': [], 'blue': []}

    # make a lin_space with the number of records from seq.     
    x = np.linspace(0,1, len(seq))
    #%
    for i in range(len(seq)):
        segment = x[i]
        tone = seq[i]
        cdict['red'].append([segment, tone, tone])
        cdict['green'].append([segment, tone, tone])
        cdict['blue'].append([segment, tone, tone])
    #%
    return mcolors.LinearSegmentedColormap('CustomMap', cdict)


#############################



n_angles = 36
n_radii = 8

# An array of radii
# Does not include radius r=0, this is to eliminate duplicate points
radii = np.linspace(0.125, 1.0, n_radii)

# An array of angles
angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)

# Repeat all angles for each radius
angles = np.repeat(angles[...,np.newaxis], n_radii, axis=1)

# Convert polar (radii, angles) coords to cartesian (x, y) coords
# (0, 0) is added here. There are no duplicate points in the (x, y) plane
x = np.append(0, (radii*np.cos(angles)).flatten())
y = np.append(0, (radii*np.sin(angles)).flatten())

# Pringle surface
z = np.sin(-x*y)



w = np.tan(-x*y)
colors = make_colormap(w)



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

ax.plot_trisurf(x, y, z, cmap=colors, linewidth=0.2)

plt.show()

This answer addresses the 4d surface plot problem. It uses matplotlib's plot_surface function instead of plot_trisurf.

Basically you want to reshape your x, y and z variables into 2d arrays of the same dimension. To add the fourth dimension as a colormap, you must supply another 2d array of the same dimension as your axes variables.

Below is example code for a 3d plot with the colormap corresponding to the x values. The facecolors argument is used to alter the colormap to your liking. Note that its value is acquired from the to_rgba() function in the matplotlib.cm.ScalarMappable class.

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

# domains
x = np.logspace(-1.,np.log10(5),50) # [0.1, 5]
y = np.linspace(6,9,50)             # [6, 9]
z = np.linspace(-1,1,50)            # [-1, 1]

# convert to 2d matrices
Z = np.outer(z.T, z)        # 50x50
X, Y = np.meshgrid(x, y)    # 50x50

# fourth dimention - colormap
# create colormap according to x-value (can use any 50x50 array)
color_dimension = X # change to desired fourth dimension
minn, maxx = color_dimension.min(), color_dimension.max()
norm = matplotlib.colors.Normalize(minn, maxx)
m = plt.cm.ScalarMappable(norm=norm, cmap='jet')
m.set_array([])
fcolors = m.to_rgba(color_dimension)

# plot
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot_surface(X,Y,Z, rstride=1, cstride=1, facecolors=fcolors, vmin=minn, vmax=maxx, shade=False)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
fig.canvas.show()

The answer I referenced (and others) mentions that you should normalize your fourth dimension data. It seems that this may be avoided by explicitly setting the limits of the colormap as I did in the code sample.

As of May 2022 the top three answers to this question each have various issues. I found the example provided in the matplotlib 3.5.0 documentation to be far simpler and actually work as expected to calculate facecolors with shading using the LightSource class.

Just override the specific z passed into ls.shade:

from matplotlib import cbook
from matplotlib import cm
from matplotlib.colors import LightSource
import matplotlib.pyplot as plt
import numpy as np

# Load and format data
dem = cbook.get_sample_data('jacksboro_fault_dem.npz', np_load=True)
z = dem['elevation']
nrows, ncols = z.shape
x = np.linspace(dem['xmin'], dem['xmax'], ncols)
y = np.linspace(dem['ymin'], dem['ymax'], nrows)
x, y = np.meshgrid(x, y)

region = np.s_[5:50, 5:50]
x, y, z = x[region], y[region], z[region]

# Set up plot
fig, ax = plt.subplots(subplot_kw=dict(projection='3d'))

ls = LightSource(270, 45)
# To use a custom hillshading mode, override the built-in shading and pass
# in the rgb colors of the shaded surface calculated from "shade".
rgb = ls.shade(z, cmap=cm.gist_earth, vert_exag=0.1, blend_mode='soft')
surf = ax.plot_surface(x, y, z, rstride=1, cstride=1, facecolors=rgb,
                       linewidth=0, antialiased=False, shade=False)

plt.show()

Many thanks to @Frik for his great answer, it helped me achieve a similar plot as requested by the OP.

However, I found that a few simplifications to the code may be done and could be of interest. Snippet and figure below.

import matplotlib.pyplot as plt
# This import registers the 3D projection, but is otherwise unused.
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401 unused import
from mpl_toolkits.mplot3d.axes3d import get_test_data
import numpy as np
fig, ax = plt.subplots(subplot_kw={'projection': '3d'})
X, Y, Z = get_test_data(0.05)
C = np.linspace(-5, 5, Z.size).reshape(Z.shape)
scamap = plt.cm.ScalarMappable(cmap='inferno')
fcolors = scamap.to_rgba(C)
ax.plot_surface(X, Y, Z, facecolors=fcolors, cmap='inferno')
fig.colorbar(scamap)
plt.show()

Finally, I also wanted to comment on what @Frik wrote:

The answer I referenced (and others) mentions that you should normalize your fourth dimension data. It seems that this may be avoided by explicitly setting the limits of the colormap as I did in the code sample.

I found this statement to be incorrect. Indeed, if one has a look at to_rgba, one can see that there is a norm keyword which is by default set to True. This is exactly where normalization occurs. The following statement is also included:

If norm is False, no normalization of the input data is performed, and it is assumed to be in the range (0-1).

You indeed want your data to lie in (0-1).