Fastest way to fill numpy array with distances from a point

Question

I have a numpy array that represents a 3d grid over a volume of space, with each cell representing a non-cubic voxel (the scaling is arbitrary in all three dimensions). The array is O(500) voxels per dimension.

I would like to fill that array with the distance from a given XYZ point to the center of each voxel.

I can fill the array by using python for-loops, but that's slower than I would like. Is there a way to do so quickly using numpy/scipy?

The conversion to XYZ coordinates is done with two tuples, one giving the XYZ coordinates of the center of the 0,0,0 voxel, and the other giving the size of the voxels in XYZ units.

Answer 1

Create an ogrid with the distances in each dimension and then calculate the distance (using that the ogrid results broadcast correctly):

import numpy as np

x0, y0, z0 = 10, 10, 10

# assuming each dimension includes 500 points, from 0 to 500, step 1
x, y, z = np.ogrid[0:500, 0:500, 0:500]  
distances = np.sqrt((x-x0)**2+(y-y0)**2+(z-z0)**2)

If you need to include some scaling and offset for the grid:

x, y, z = np.ogrid[0:500, 0:500, 0:500]
x, y, z = (x * scale_x + offset_x, 
           y * scale_y + offset_y, 
           z * scale_z + offset_z)
distances = np.sqrt((x-x0)**2+(y-y0)**2+(z-z0)**2)

Answer 2

You can create three one-dimensional arrays representing the flattened X, Y and Z coordinates of your 3D array.

Then you perform the hypothenuse calculation for the whole arrays, using numpy methods:

D = numpy.sqrt(numpy.power(X - x_center, 2) + 
               numpy.power(Y - y_center, 2) + 
               numpy.power(Z - z_center, 2))

At the end, you reshape the array to the original shape.