Search code examples
pythonnumpymatplotlibslice

plot a slice of 3D data with pcolormesh

I have data in 2 numpy arrays: one a list of 3D positions, the other the values of a scalar AT each of those positions. The ordering of the position data is fairly 'weird' (see below).

The 3D position data is in an array:

pos = np.array([[1,1,1],[1,1,2],[1,1,3],[1,2,2],[1,2,1], ...])

pos.shape is (100000,3)

where the ordering is not intuitive (it is following a space filling curve).

and I also have the scalar values that I want to plot at each of those locations:

vel = np.array([1,2,1,3,4,...])

vel.shape = (1000000,1)

My question is, how do I plot an xy slice with pcolormesh from this data??? I can extract an xy plane with numpy:

xs = pos[:,:1][:,0]
ys = pos[:,1:2][:,0]

Now I have a bunch of essentially random x coords and y coords that no long map 1-1 to the vel data... :/. I don't know how to first map those initial positions to my vel data, so that I can generate a pcolormesh:

plt.pcolormesh(X, Y, V)

Can someone help me slice this data up, so everything stays mapped to the right position in xy (and z) space?

Solution

  • If I understand correctly, the following approach would work on your data:

    • find all the indices where the z is close to the desired z (np.where())
    • filter both pos and vel by these indices
    • find the order of the x and y values inside pos
    • apply this order to vel, and reshape to a 2D array
    • display the 2D array via imshow(); setting vmin and vmax will have the same colors for the same values for different subplots
    import matplotlib.pyplot as plt
import numpy as np

# first create some test data

# generate a test grid
N = 100  # size in one dimension
# create a grid of x y and z coordinates
xs, ys, zs = np.meshgrid(range(1, N + 1), range(1, N + 1), range(1, N + 1))
# rearrange them in the shape of the example
pos = np.vstack([xs.ravel(), ys.ravel(), zs.ravel()]).T
# randomly reorder the positions
np.random.shuffle(pos)

# generate a test function for `vel`, depending on pos
# calculate distance to the center of two spheres
center1 = [(N - 1) * 0.2, (N - 1) * 0.6, (N - 1) * 0.5]
radius1 = (N - 1) * 0.3
dist_to_center1 = np.sqrt((pos[:, 0] - center1[0]) ** 2 + (pos[:, 1] - center1[1]) ** 2 + (pos[:, 2] - center1[2]) ** 2)
center2 = [(N - 1) * 0.7, (N - 1) * 0.4, (N - 1) * 0.5]
radius2 = (N - 1) * 0.2
dist_to_center2 = np.sqrt((pos[:, 0] - center2[0]) ** 2 + (pos[:, 1] - center2[1]) ** 2 + (pos[:, 2] - center2[2]) ** 2)
# take union of the two spheres: vel is the minimum of signed distance to the spheres
vel = np.minimum(dist_to_center1 - radius1, dist_to_center2 - radius2)

    # display the layers for some z-values
xmin = pos[:, 0].min()
xmax = pos[:, 0].max()
ymin = pos[:, 1].min()
ymax = pos[:, 1].max()
fig, axs = plt.subplots(nrows=2, ncols=4, figsize=(15, 7))
for z_special, ax in zip(range(2, N - 2, N // 8 + 1), axs.flat):
    # indices where the z-value almost equals the desired z value
    filter = np.argwhere((pos[:, 2] > z_special - 0.001) & (pos[:, 2] < z_special + 0.001))[:, 0]
    xy_special = pos[:, :2][filter]
    vel_special = vel[filter]
    # find the order of xy_special on x and y
    order = np.lexsort(([xy_special[:, 0], xy_special[:, 1]]))
    vel_ordered = vel_special[order].reshape(N, -1)
    ax.imshow(vel_ordered, origin='lower', extent=[xmin, xmax, ymin, ymax], aspect='auto',
              cmap='turbo', vmin=-20, vmax=20)
    ax.set_xticks([]) # remove ticks to simplify the plot
    ax.set_yticks([])
    ax.set_title(f"z = {z_special:.3g}")

plt.tight_layout()
plt.show()

    cutting through a 3D volume with unordered x y z values

    Instead of images, you could also draw isolines for vel values. tricontourf() also works for unordered x, y pairs, even when not all pairs of a grid are present. To make the coloring consistent between subplots, you can explicitly set a list of levels.

    fig, axs = plt.subplots(nrows=2, ncols=4, figsize=(15, 7))
for z_special, ax in zip(range(2, N - 2, N // 8 + 1), axs.flat):
    # indices where the z-value almost equals the desired z value
    filter = np.argwhere((pos[:, 2] > z_special - 0.001) & (pos[:, 2] < z_special + 0.001))[:, 0]
    xy_special = pos[:, :2][filter]
    vel_special = vel[filter]
    ax.tricontourf(xy_special[:, 0], xy_special[:, 1], vel_special, cmap='turbo', levels=range(-20, 21, 5))
    ax.set_title(f"z = {z_special:.3g}")

    tricontourf over slices