Efficiently apply function to spheric neighbourhood in numpy array

Question

I have a 3D numpy array of float values in Python. I need to retrieve all the elements in a sphere of radius r starting from a center point P(x, y, z). Then, I want to apply to the sphere points a function that updates their values and needs the distance to the center point to do this. I do these steps a lot of times and for large radius values, so I would like to have a solution that is as efficient as possible.

My current solution checks only the points in the bounding box of the sphere, as indicated here: Using a QuadTree to get all points within a bounding circle. A sketch of the code looks like this:

# P(x, y, z): center of the sphere
for k1 in range(x - r, x + r + 1):
    for k2 in range(y - r, y + r + 1):
        for k3 in range(z - r, z + r + 1):
            # Sphere center - current point distance
            dist = np.sum((np.array([k1, k2, k3]) - np.array([x, y, z])) ** 2)

            if (dist <= r * r):
                # computeUpdatedValue(distance, radius): function that computes the new value of the matrix in the current point 
                newValue = computeUpdatedValue(dist, r)

                # Update the matrix
                mat[k1, k2, k3] = newValue

However, I thought that applying a mask to retrive the points and, then, update them based on distance in a vectorized manner is more efficient. I have seen how to apply a circular kernel (How to apply a disc shaped mask to a numpy array?), but I do no know how to efficiently apply the function (depending on the indices) on each of the mask's elements.

jdehesa · Accepted Answer · 2018-12-06T14:50:00.923

EDIT: If your array is very big compared to the region you are updating, the solution below will take much more memory than necessary. You can apply the same idea but only to the region where the sphere may fall:

def updateSphereBetter(mat, center, radius):
    # Find beginning and end of region of interest
    center = np.asarray(center)
    start = np.minimum(np.maximum(center - radius, 0), mat.shape)
    end = np.minimum(np.maximum(center + radius + 1, 0), mat.shape)
    # Slice region of interest
    mat_sub = mat[tuple(slice(s, e) for s, e in zip(start, end))]
    # Center coordinates relative to the region of interest
    center_rel = center - start
    # Same as before but with mat_sub and center_rel
    ind = np.indices(mat_sub.shape)
    ind = np.moveaxis(ind, 0, -1)
    dist_squared = np.sum(np.square(ind - center_rel), axis=-1)
    mask = dist_squared <= radius * radius
    mat_sub[mask] = computeUpdatedValue(dist_squared[mask], radius)

Note that since mat_sub is a view of mat, updating it updates the original array, so this produces the same result as before, but with less resources.

Here is a little proof of concept. I defined computeUpdatedValue so that it shows the distance from the center, and then plotted a few "sections" of an example:

import numpy as np
import matplotlib.pyplot as plt

def updateSphere(mat, center, radius):
    # Make array of all index coordinates
    ind = np.indices(mat.shape)
    # Compute the squared distances to each point
    ind = np.moveaxis(ind, 0, -1)
    dist_squared = np.sum(np.square(ind - center), axis=-1)
    # Make a mask for squared distances within squared radius
    mask = dist_squared <= radius * radius
    # Update masked values
    mat[mask] = computeUpdatedValue(dist_squared[mask], radius)

def computeUpdatedValue(dist_squared, radius):
    # 1 at the center of the sphere and 0 at the surface
    return np.clip(1 - np.sqrt(dist_squared) / radius, 0, 1)

mat = np.zeros((100, 60, 80))
updateSphere(mat, [50, 20, 40], 20)

plt.subplot(131)
plt.imshow(mat[:, :, 30], vmin=0, vmax=1)
plt.subplot(132)
plt.imshow(mat[:, :, 40], vmin=0, vmax=1)
plt.subplot(133)
plt.imshow(mat[:, :, 55], vmin=0, vmax=1)

Output:

Efficiently apply function to spheric neighbourhood in numpy array

1 Answers1