CODEWITHSUNDEEP
numpy
durden2.0
durden2.0

Reputation: 9542

Best way to groups cells in 3D numpy array by proximity

I'm trying to find the best way to find/set all cells surrounding a 3D location. This is sort of a clustering problem. I would like to cluster 'groups' of cells around a few particular locations.

For example, given a cell distance of 2, a 3D array of (4,4,4), and a list of 3D locations:

(0, 0, 0)
(0, 2, 0)
(3, 0, 0)

Assume the given array is all zeros.

I'd like to return a 3D array that has 1 in all the cells that are within the cell distance of 2 from any location.

So, any cell that is 2 or less away from the above locations (including the location itself) should have a 1 in it.

I want to use numpy slicing instead of a Python for loops for performance.

Here's a sample input and output:

distance = 2
shape = (4, 4, 4)
locations = [(2, 0, 0)]

return should be:

array([[[1, 1, 1, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0],
        [0, 0, 0, 0]],

       [[1, 1, 1, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0],
        [0, 0, 0, 0]],

       [[1, 1, 1, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0],
        [0, 0, 0, 0]],

       [[1, 1, 1, 0],
        [1, 1, 1, 0],
        [1, 1, 1, 0],
        [0, 0, 0, 0]]], dtype=int32)

Essentially a cube of size 2x2x2 is formed around the 'central' point (2, 0, 0)

Upvotes: 1

Views: 210

Answers (1)

Eelco Hoogendoorn
Eelco Hoogendoorn

Reputation: 10759

Does this do what you want?

import numpy as np

distance = 2
shape = (4, 4, 4)
locations = np.array([(2, 0, 0)])

data = np.zeros(shape, np.int)

data[locations[:,0], locations[:,1], locations[:,2]] = 1

import scipy.ndimage
scipy.ndimage.binary_dilation(data, np.ones((3,3,3)), 2)

Upvotes: 1

Related Questions