Draw and/or get indices for a 2D plane in a 3D numpy array

Question

I have a 3D numpy array, e.g. of shape (200,200,200).

I also have a 2D plane, which I can express either in the form ax + by + cz + d = 0 or a (x,y,z) point and a direction vector in terms of the three x,y,z axes.

I want to "plot" this plane on the 3D numpy array, by essentially changing the value of all points that lie on this plane to a constant value.

Like this, but with only a single colour:

The most versatile solution to this (I think) is to get all the integer indices of points that lie along this line. I can then just use this to index the numpy array. However, if a library offered the ability to plot a plane directly, without providing the indices, then this would be fine.

I have seen problems on this site that seem superficially similar to my problem, such as approaches of finding the values that lie along an x,y,z vector using scipy.ndimage.map_coordinates, but of course I am dealing with a plane not a line (also this just returned the values, not the indices).

Another approach I considered, but seems difficult (and maybe slow) is something similar to this question where the reply shows how to draw a 2D triangle in a 3D space. I could instead draw a square in 3D space, but this shows that filling this square is not trivial, and also the square will only fill the entire plane if at least one axis is parallel to x, y or z (or a "corner" will remain unfilled).

Has anyone got any idea how I might achieve this?

Just to be clear, you want to set a plane of points within the array to some value, not plot a plane in matplotlib? — Mad Physicist, Jul 24 '21 at 05:13

aerobiomat · Accepted Answer · 2021-07-23T19:37:12.460

If the plane is near horizontal, you can give values to x and y and calculate z:

a, b, c, d = 1, 2, 3, -600                     # Plane parameters
v = np.arange(200)
x, y = np.meshgrid(v, v)                       # All xy combinations
z = np.rint((a*x + b*y + d) / -c).astype(int)  # Rounded
plane_voxels = np.dstack([x,y,z]).reshape(-1,3)
print(plane_voxels)

Results in:

[[  0   0 200]
 [  1   0 200]
 [  2   0 199]
 ...
 [197 199   2]
 [198 199   1]
 [199 199   1]]

For the general case you need to find the dimension with less variability, that will be the calculated variable, and give values to the other two:

a, b, c, d = 1, 2, 3, -600
v = np.arange(200)
dim = np.argmax(np.abs((a, b, c)))
if dim == 0:
    y, z = np.meshgrid(v, v)
    x = np.rint((b*y + c*z + d) / -a).astype(int)
elif dim == 1:
    x, z = np.meshgrid(v, v)
    y = np.rint((a*x + c*z + d) / -b).astype(int)
elif dim == 2:
    x, y = np.meshgrid(v, v)
    z = np.rint((a*x + b*y + d) / -c).astype(int)
plane_voxels = np.dstack([x,y,z]).reshape(-1,3)

Very nice. Two nitpicks: (1), you don't need a meshgrid. Just `np.arange(200)` and `np.arange(200)[:, None]` is sufficient. The third dimension will broadcast correctly. (2) OP wanted to set the elements of the cube, so you can do `cube[y, x, z] = 1` or whatever. — Mad Physicist, Jul 24 '21 at 05:22

score 0 · Answer 2 · answered Jul 23 '21 at 19:39

Here is an idea. First let's say you have a (3,3,3) of one. Then define your function of the plane equal to 0. create your coordinates with meshgrid and use it in your plane function. Finally in case you want some tolerance, use isclose to 0 as a mask of your array to change the values

n = 3 # change to 200 for your real case or arr.shape[0]

# just for the example
arr =  np.ones([n]*3)

# plane function = 0
f = lambda x,y,z: 2*x + 3*y**2 + 4*z - 3

# create the mask
mask = np.isclose(f(*np.meshgrid(*[np.arange(n)]*3)), 
                  0, # because your plane function is equal to 0
                  atol=1) # this is if you want some tolerance to catch nearby points

# change the value to whatever
arr[mask] = 5

print(arr)
[[[1. 5. 1.]
  [5. 1. 1.]
  [5. 1. 1.]]

 [[5. 1. 1.]
  [1. 1. 1.]
  [1. 1. 1.]]

 [[1. 1. 1.]
  [1. 1. 1.]
  [1. 1. 1.]]]

Draw and/or get indices for a 2D plane in a 3D numpy array

2 Answers2