Adding a 4th variable to a 3D plot in Python

Question

I have 3 different parameters X,Y and Z over a range of values, and for each combination of these a certain value of V. To make it clearer, the data would look something like this.

X  Y  Z  V
1  1  2  10
1  2  3  15
etc...

I'd like to visualize the data with a surface/contour plot, using V as a colour to see its value at that point, but I do not see how to add my custom colouring scheme into the mix using Python. Any idea on how to do this (or is this visualization outright silly)?

Thanks a lot!

Answer 1

It really depends on how you plan on plotting this data. I like to plot graphs with gnuplot: it's easy, free and intuitive. To plot your example with gnuplot you'd have to print those line into a file (with only those four columns) and plot using a code like the following

 reset
 set terminal png
 set output "out.png"
 splot "file.txt" using 1:2:3:4 with lines palette

Assuming that you save your data into the file file.txt. splot stands for surface plot. Of course, this is a minimum example.

Alternatively you can use matplotlib, but that is not, in my opinion, as intuitive. Although it has the advantage of centering all the processing in python.

Answer 2

Matplotlib allows one to pass the facecolors as an argument to e.g. ax.plot_surface.

That would imply then that you would have to perform 2D interpolation on your current array of colors, because you currently only have the colors in the corners of the rectangular faces (you did mention that you have a rectilinear grid).

You could use scipy.interpolate.interp2d for that, but as you see from the documentation, it is suggested to use scipy.interpolate.RectBivariateSpline.

To give you a simple example:

import numpy as np
y,x = np.mgrid[1:10:10j, 1:10:10j] # returns 2D arrays
# You have 1D arrays that would make a rectangular grid if properly reshaped.
y,x = y.ravel(), x.ravel()  # so let's convert to 1D arrays
z = x*(x-y)
colors = np.cos(x**2) - np.sin(y)**2

Now I have a similar dataset as you (one-dimensional arrays for x, y, z and colors). Remark that the colors are defined for each point (x,y). But when you want to plot with plot_surface, you'll generate rectangular patches, of which the corners are given by those points.

So, on to interpolation then:

from scipy.interpolate import RectBivariateSpline
# from scipy.interpolate import interp2d # could 've used this too, but docs suggest the faster RectBivariateSpline

# Define the points at the centers of the faces:
y_coords, x_coords = np.unique(y), np.unique(x)
y_centers, x_centers = [ arr[:-1] + np.diff(arr)/2 for arr in (y_coords, x_coords)]

# Convert back to a 2D grid, required for plot_surface:
Y = y.reshape(y_coords.size, -1)
X = x.reshape(-1, x_coords.size)
Z = z.reshape(X.shape)
C = colors.reshape(X.shape)
#Normalize the colors to fit in the range 0-1, ready for using in the colormap:
C -= C.min()
C /= C.max()

interp_func = RectBivariateSpline(x_coords, y_coords, C.T, kx=1, ky=1) # the kx, ky define the order of interpolation. Keep it simple, use linear interpolation.

In this last step, you could also have used interp2d (with kind='linear' replacing the kx=1, ky=1). But since the docs suggest to use the faster RectBivariateSpline...

Now you're ready to plot it:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.cm as cm

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
r = ax.plot_surface(X,Y,Z,
    facecolors=cm.hot(interp_func(x_centers, y_centers).T),
    rstride=1,  cstride=1) # only added because of this very limited dataset

As you can see, the colors on the faces have nothing to do anymore with the height of the dataset.

Note that you could have thought simply passing the 2D array C to facecolors would work, and matplotlib would not have complained. However, the result isn't accurate then, because matplotlib will use only a subset of C for the facecolors (it seems to ignore the last column and last row of C). It is equivalent to using only the color defined by one coordinate (e.g. the top-left) over the entire patch.

An easier method would have been to let matplotlib do the interpolation and obtain the facecolors and then pass those in to the real plot:

r = ax.plot_surface(X,Y,C, cmap='hot') # first plot the 2nd dataset, i.e. the colors
fc = r.get_facecolors()
ax.clear()
ax.plot_surface(X, Y, Z, facecolors=fc)

However, that won't work in releases <= 1.4.1 due to this recently submitted bug.