Making a 4-D scatter plot using Matplotlib in Python

Question

Essentially, I'm trying to make a 4-D scatter plot with 4 columns of data (see sample below).

X (mm)  Y (mm)  Z (mm)  Diameter (mm)
11.096  11.0972 13.2401 124.279
14.6836 11.0389 8.37134 138.949
19.9543 11.1025 31.1912 138.949
15.4079 10.9505 31.1639 152.21
20.6372 14.5175 6.94501 152.211
20.47   11.225  31.3612 152.211
19.0432 11.3234 8.93819 152.213
29.4091 10.1331 26.6354 186.417
12.9391 10.6616 28.9523 186.418
29.9102 10.4828 25.1129 186.418
30.5483 12.163  15.9116 186.418
19.0631 10.5784 30.9791 186.418
9.65332 10.8563 12.975  186.419
8.4003  11.0417 17.0181 186.419
26.0134 10.6857 9.41572 186.419
13.7451 11.1495 28.7108 186.419

The first three columns of data (X, Y, Z) are the coordinate positions of the 4th column of data (Diameter) so I was able to generate a 3-D scatter plot of these positions. However, I'm trying to plot these Diameters with different color markers based on certain threshold values (ie. Diameters that are less than 100 mm are red, 101-200 mm are blue, 201-300 mm are green, etc.) Once the color of the markers are determined, it would plot these markers based on its X, Y, Z coordinates. I tried writing a simple for loop to do this, but for some reason it becomes very slow/laggy and will only plot one color too. Can anyone see if there's something wrong with my approach? Thanks!

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import pandas
import os

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

os.chdir(r'C:\Users\Me\Documents')
data = pandas.read_excel("Diameter Report", "Data")

X = data['X (mm)'].values.tolist()
Y = data['Y (mm)'].values.tolist()
Z = data['Z (mm)'].values.tolist()
dims = data['Diameter (mm)'].values.tolist()

for i in dims:
    if i < int(100):
        ax.plot(X, Y, Z, c='r', marker='o')
    elif i >= int(101) and i <200:
        ax.plot(X, Y, Z, c='b', marker='o')
    elif i >= int(201) and i <300:
        ax.plot(X, Y, Z, c='g', marker='o')

ax.set_xlabel('Center X (mm)')
ax.set_ylabel('Center Y (mm)')
ax.set_zlabel('Center Z (mm)')

plt.show()

Answer 1

It seems the thresholds for the values are equally spaced, so you can just divide by 100 and truncate further decimal places. This allows to plot a single scatter instead of hundreds of plots.

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

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

data = pandas.read_excel("Diameter Report", "Data")

X = data['X (mm)'].values
Y = data['Y (mm)'].values
Z = data['Z (mm)'].values
dims = data['Diameter (mm)'].values

ax.scatter(X,Y,Z, c=(dims/100).astype(int), marker="o", cmap="brg")

ax.set_xlabel('Center X (mm)')
ax.set_ylabel('Center Y (mm)')
ax.set_zlabel('Center Z (mm)')

plt.show()

The more general case of arbitrary boundaries would probably best be solved using a BoundaryNorm and a colormap with as many different colors as classifications.

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

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

d = np.random.rand(10,4)
d[:,3] = np.random.randint(1,300, 10)
data = pd.DataFrame(d, columns=["X (mm)","Y (mm)","Z (mm)","Diameter (mm)"])

X = data['X (mm)'].values
Y = data['Y (mm)'].values
Z = data['Z (mm)'].values
dims = data['Diameter (mm)'].values

bounds = [0,100,200,300]
colors = ["b", "r", "g"]
cmap = mcolors.ListedColormap(colors)
norm = mcolors.BoundaryNorm(bounds, len(colors))

sc = ax.scatter(X,Y,Z, c=dims, marker="o", cmap=cmap, norm=norm)

ax.set_xlabel('Center X (mm)')
ax.set_ylabel('Center Y (mm)')
ax.set_zlabel('Center Z (mm)')
fig.colorbar(sc)
plt.show()

Answer 2

Here is a slightly more general solution where you can explicitly specify the ranges you want regardless of the spacing. I did not have the complete data so I modified your limits from 100, 200, 300 to 140, 180, 200 based on the provided data.

A couple of things:

• You probably want to use scatter3d as you mentioned it in your question instead of plot.
• I am using NumPy to read in the data because this way you will have the data as NumPy arrays which make the masking and slicing easy.
• Here I am creating 3 conditional masks depending on the magnitude of dims.
• Next, you store these masks in a list and then iterate over it to use one mask at a time.

---

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

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

X, Y, Z, dims = np.loadtxt('sample.txt', unpack=True, skiprows=1)

mask1 = (dims<140)
mask2 = ((dims>=140) & (dims<180))
mask3 = ((dims>=180) & (dims<200))

masks = [mask1, mask2, mask3]
colors = ['r', 'b', 'g'] # color order as you specified in the question

for mask, color in zip(masks, colors): 
    ax.scatter3D(X[mask], Y[mask], Z[mask], c=color)

ax.set_xlabel('Center X (mm)')
ax.set_ylabel('Center Y (mm)')
ax.set_zlabel('Center Z (mm)')
plt.show()