Although there are many instances of the question: "What is the numpy alternative to nested for loops", I was unable to find a suitable answer for my case. Here it goes:
I have a 3D numpy array with "0" background and other integers as foreground. I would like to find and store the foreground voxels which fall within a predefined mask (a sphere defining a given distance from a reference node). I have successfully done the task using nested 'for' loops and a chain of 'if' conditions as shown below. I am looking for a more efficient and compact alternative to avoid the loops and long conditions for this neighborhood search algorithm.
sample input data:
import numpy as np
im = np.array([[[ 60, 54, 47, 52, 57, 53, 46, 48]
, [ 60, 57, 53, 53, 54, 53, 50, 55]
, [ 60, 63, 56, 58, 59, 57, 50, 50]
, [ 70, 70, 64, 69, 74, 72, 64, 47]
, [ 73, 76, 77, 80, 82, 76, 58, 37]
, [ 85, 85, 86, 86, 78, 62, 38, 20]
, [ 94, 94, 92, 78, 54, 33, 16, 255]
, [ 94, 90, 72, 51, 32, 19, 255, 255]
, [ 65, 53, 29, 18, 255, 255, 255, 255]
, [ 29, 22, 255, 255, 255, 255, 255, 0]]
, [[ 66, 67, 70, 69, 75, 73, 72, 63]
, [ 68, 70, 73, 74, 78, 80, 74, 53]
, [ 75, 87, 87, 83, 89, 86, 61, 33]
, [ 81, 89, 88, 98, 99, 77, 41, 18]
, [ 84, 94, 100, 100, 82, 49, 21, 255]
, [ 99, 101, 92, 75, 48, 25, 255, 255]
, [ 93, 77, 52, 32, 255, 255, 255, 255]
, [ 52, 40, 25, 255, 255, 255, 255, 255]
, [ 23, 16, 255, 255, 255, 255, 255, 0]
, [255, 255, 255, 255, 255, 255, 0, 0]]
, [[ 81, 83, 92, 101, 101, 83, 49, 19]
, [ 86, 96, 103, 103, 95, 64, 28, 255]
, [ 94, 103, 107, 98, 79, 41, 255, 255]
, [101, 103, 98, 79, 51, 28, 255, 255]
, [102, 97, 76, 49, 27, 255, 255, 255]
, [ 79, 62, 35, 21, 255, 255, 255, 255]
, [ 33, 23, 15, 255, 255, 255, 255, 255]
, [ 16, 255, 255, 255, 255, 255, 255, 0]
, [255, 255, 255, 255, 255, 255, 0, 0]
, [255, 255, 255, 255, 255, 0, 0, 0]]
, [[106, 107, 109, 94, 58, 26, 15, 255]
, [110, 104, 90, 66, 37, 19, 255, 255]
, [106, 89, 61, 35, 22, 255, 255, 255]
, [ 76, 56, 34, 19, 255, 255, 255, 255]
, [ 40, 27, 18, 255, 255, 255, 255, 255]
, [ 17, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 0]
, [255, 255, 255, 255, 255, 255, 0, 0]
, [255, 255, 255, 255, 255, 0, 0, 0]
, [255, 255, 255, 0, 0, 0, 0, 0]]
, [[ 68, 51, 33, 19, 255, 255, 255, 255]
, [ 45, 34, 20, 255, 255, 255, 255, 255]
, [ 28, 18, 255, 255, 255, 255, 255, 255]
, [ 17, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 0]
, [255, 255, 255, 255, 255, 255, 0, 0]
, [255, 255, 255, 255, 255, 0, 0, 0]
, [255, 255, 255, 0, 0, 0, 0, 0]
, [255, 0, 0, 0, 0, 0, 0, 0]]
, [[255, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 255]
, [255, 255, 255, 255, 255, 255, 255, 0]
, [255, 255, 255, 255, 255, 255, 0, 0]
, [255, 255, 255, 255, 255, 0, 0, 0]
, [255, 255, 255, 255, 0, 0, 0, 0]
, [255, 255, 255, 0, 0, 0, 0, 0]
, [255, 0, 0, 0, 0, 0, 0, 0]
, [ 0, 0, 0, 0, 0, 0, 0, 0]]
, [[255, 255, 255, 255, 255, 255, 255, 0]
, [255, 255, 255, 255, 255, 255, 255, 0]
, [255, 255, 255, 255, 255, 255, 0, 0]
, [255, 255, 255, 255, 255, 0, 0, 0]
, [255, 255, 255, 255, 0, 0, 0, 0]
, [255, 255, 255, 0, 0, 0, 0, 0]
, [255, 255, 0, 0, 0, 0, 0, 0]
, [ 0, 0, 0, 0, 0, 0, 0, 0]
, [ 0, 0, 0, 0, 0, 0, 0, 0]
, [ 0, 0, 0, 0, 0, 0, 0, 0]]
, [[255, 255, 255, 255, 255, 255, 0, 0]
, [255, 255, 255, 255, 255, 0, 0, 0]
, [255, 255, 255, 255, 0, 0, 0, 0]
, [255, 255, 255, 0, 0, 0, 0, 0]
, [255, 255, 0, 0, 0, 0, 0, 0]
, [255, 0, 0, 0, 0, 0, 0, 0]
, [ 0, 0, 0, 0, 0, 0, 0, 0]
, [ 0, 0, 0, 0, 0, 0, 0, 0]
, [ 0, 0, 0, 0, 0, 0, 0, 0]
, [ 0, 0, 0, 0, 0, 0, 0, 0]]])
The implemented method:
[Z,Y,X]=im.shape
RN = np.array([3,4,4])
################Loading Area search
rad = 3
a,b,c = RN
x,y,z = np.ogrid[-c:Z-c,-b:Y-b,-a:X-a]
neighborMask = x*x + y*y + z*z<= rad*rad
noNodeMask = im > 0
mask = np.logical_and(neighborMask, noNodeMask)
imtemp = im.copy()
imtemp[mask] = -1
for i in range (X):
for j in range (Y):
for k in range (Z):
if imtemp[i,j,k]==-1:
if i in (0, X-1) or j in (0, Y-1) or k in (0, Z-1):
imtemp[i,j,k]=-2
elif imtemp[i+1,j,k] == 0 or imtemp[i-1,j,k] == 0 or imtemp[i,j+1,k] == 0 or imtemp[i,j-1,k] == 0 or imtemp[i,j,k+1] == 0 or imtemp[i,j,k-1] == 0:
imtemp[i,j,k]=-2
LA = np.argwhere(imtemp==-2)
The resulting LA from the above sample code is:
In [90]:LA
Out[90]:
array([[4, 4, 0],
[4, 4, 6],
[4, 5, 5],
[4, 6, 4],
[4, 6, 5],
[4, 7, 3],
[5, 3, 5],
[5, 4, 4],
[5, 4, 5],
[5, 5, 3],
[5, 5, 4],
[5, 6, 2],
[5, 6, 3],
[6, 2, 4],
[6, 3, 3],
[6, 3, 4],
[6, 4, 2],
[6, 4, 3],
[6, 5, 1],
[6, 5, 2]])
And a slice in Z direction (an XY plane instance) which shows different untouched, masked (-1), and target (-2) nodes:
Problem Statement
You have a "solid" shape in a large array. You carve out a ball from that. Your goal is to find the indices of the surface of the solid within the ball. Surfaces are defined as any point neighboring the outside of the solid with 6-point connectivity. Edges of the array are considered to be surfaces too.
Faster Loop Solution
You already computed the mask that represents the intersection of the solid and the ball. You can compute the mask a little more elegantly and convert it to indices instead. I suggest keeping the order of your dimensions constant, instead of switching between different notations. The order of RN
is affected, for example, and you run the risk of mismatching your axis limits.
RN = np.array([4, 4, 3])
rad = 3
im = ...
cutout = ((np.indices(im.shape) - RN.reshape(-1, 1, 1, 1))**2).sum(axis=0) <= rad**2
solid = im > 0
mask = solid & cutout
indices = np.argwhere(mask)
You can also get the cutout without reshaping RN
by doing
cutout = ((np.rollaxis(np.indices(im.shape, sparse=False), 0, 4) - RN)**2).sum(axis=-1) <= rad**2
The nice thing about computing indices is that your loops don't need to be huge any more. By using argwhere
, you basically strip off the outer three loops, leaving only the if
statement to loop over. You can also vectorize the connectivity check. This has the nice side effect that you can define arbitrary connectivity for each pixel.
limit = np.array(im.shape) - 1 # Edge of `im`
connectivity = np.array([[ 1, 0, 0], # Add rows to determine connectivity
[-1, 0, 0],
[ 0, 1, 0],
[ 0, -1, 0],
[ 0, 0, 1],
[ 0, 0, -1]], dtype=indices.dtype)
index_mask = np.ones(len(indices), dtype=bool)
for n, ind in enumerate(indices):
if ind.all() and (ind < limit).all() and im[tuple((ind + connectivity).T)].all():
index_mask[n] = False
LA = indices[index_mask, :]
Notice that there is really no point to imtemp
at all. Even in your original loop, you could just manipulate mask
directly. Instead of setting elements to -2
when they pass your criterion, you could set elements to False
if they didn't.
I do something like that here. We check each of the indices that were actually selected, and determine if any of them are inside the solid. These indices are eliminated from the mask. The list of indices is then updated based on the mask.
The check ind.all() and (ind < limit).all() and im[tuple((ind + connectivity).T)].all()
is a shortcut for what you were doing with the or
conditions, but reversed (testing for non-surface rather than surface).
ind.all()
checks that none of the indices are zero: i.e., not on a top/front/left surface.(ind < limit).all()
checks that none of the indices are equal to the corresponding image size minus one.im[tuple((ind + connectivity).T)].all()
checks that none of the connected pixels are zero. (ind + connectivity).T
is a (3, 6)
array of the six points that we are connected to (currently defined as +/-1 in each axis by the (6, 3)
connectivity
array). When you turn it into a tuple, it just becomes a fancy index, as if you had done something like im[x + connectivity[:, 0], y + connectivity[:, 1], z + connectivity[:, 2]]
. The commas in the index just make it into a tuple. They way I show is better suited for arbitrary numbers of dimensions.Pixels that pass all three tests are inside the solid, and get removed. You can of course write the loop to check the other way, but then you would have to alter your mask:
index_mask = np.zeros(len(indices), dtype=bool)
for n, ind in enumerate(indices):
if (ind == 0).any() or (ind == limit).any() or (im[tuple((ind + connectivity).T)] == 0).any():
index_mask[n] = True
LA = indices[index_mask, :]
Looping manually is not ideal by any means. However, it shows you how to shorten the loop (probably by a couple of orders of magnitude), and how to define arbitrary connectivity using vectorization and broadcasting, without getting bogged down with hard-coding it.
Fully Vectorized Solution
The loops above can be fully vectorized using the magic of broadcasting. Instead of looping over each row in indices
, we can add connectivity
to it in bulk and filter the results in bulk. The trick is to add enough dimensions that you add all of connectivity
to each element of indices
.
You will still want to omit the pixels that are at the edges:
edges = (indices == 0).any(axis=-1) | (indices == limit).any(axis=-1)
conn_index = indices[~edges, None, :] + connectivity[None, ...]
index_mask = np.empty(len(indices), dtype=bool)
index_mask[edges] = True
index_mask[~edges] = (im[tuple(conn_index.T)] == 0).any(axis=0)
LA = indices[index_mask, :]
I expect that a properly written loop compiled with numba will be significantly faster than this solution, because it will avoid much of the overhead with pipelining the operations. It will not require large temporary buffers or special handling.
TL;DR
# Parameters
RN = np.array([4, 4, 3])
rad = 3
im = ...
# Find subset of interest
cutout = ((np.indices(im.shape) - RN.reshape(-1, 1, 1, 1))**2).sum(axis=0) <= rad**2
solid = im > 0
# Convert mask to indices
indices = np.argwhere(solid & cutout)
# Find image edges among indices
edges = (indices == 0).any(axis=-1) | (indices == limit).any(axis=-1)
# Connectivity elements for non-edge pixels
conn_index = indices[~edges, None, :] + connectivity[None, ...]
# Mask the valid surface pixels
index_mask = np.empty(len(indices), dtype=bool)
index_mask[edges] = True
index_mask[~edges] = (im[tuple(conn_index.T)] == 0).any(axis=0)
# Final result
LA = indices[index_mask, :]