pythonperformancenumpynested-loopsneighbours

Efficient neighbourhood search in numpy ndarray instead of nested conditional for loops


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: a sample slice of the original and masked matrices


Solution

  • 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).

    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, :]