CODEWITHSUNDEEP
pythonperformancenumpyscipy
Eli Stevens
Eli Stevens

Reputation: 1447

Fastest way to fill numpy array with distances from a point

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.

Upvotes: 5

Views: 2343

Answers (2)

MSeifert
MSeifert

Reputation: 152870

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)

Upvotes: 7

heltonbiker
heltonbiker

Reputation: 27615

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.

Upvotes: 2

Related Questions