Efficient sum of Gaussians in 3D with NumPy using large arrays

Question

I have an M x 3 array of 3D coordinates, coords (M ~1000-10000), and I would like to compute the sum of Gaussians centered at these coordinates over a mesh grid 3D array. The mesh grid 3D array is typically something like 64 x 64 x 64, but sometimes upwards of 256 x 256 x 256, and can go even larger. I’ve followed this question to get started, by converting my meshgrid array into an array of N x 3 coordinates, xyz, where N is 64^3 or 256^3, etc. However, for large array sizes it takes too much memory to vectorize the entire calculation (understandable since it could approach 1e11 elements and consume a terabyte of RAM) so I’ve broken it up into a loop over M coordinates. However, this is too slow.

I’m wondering if there is any way to speed this up at all without overloading memory. By converting the meshgrid to xyz, I feel like I’ve lost any advantage of the grid being equally spaced, and that somehow, maybe with scipy.ndimage, I should be able to take advantage of the even spacing to speed things up.

Here’s my initial start:

import numpy as np
from scipy import spatial

#create meshgrid
side = 100.
n = 64 #could be 256 or larger
x_ = np.linspace(-side/2,side/2,n)
x,y,z = np.meshgrid(x_,x_,x_,indexing='ij')

#convert meshgrid to list of coordinates
xyz = np.column_stack((x.ravel(),y.ravel(),z.ravel()))

#create some coordinates
coords = np.random.random(size=(1000,3))*side - side/2

def sumofgauss(coords,xyz,sigma):
    """Simple isotropic gaussian sum at coordinate locations."""
    n = int(round(xyz.shape[0]**(1/3.))) #get n samples for reshaping to 3D later
    #this version overloads memory
    #dist = spatial.distance.cdist(coords, xyz)
    #dist *= dist
    #values = 1./np.sqrt(2*np.pi*sigma**2) * np.exp(-dist/(2*sigma**2))
    #values = np.sum(values,axis=0)
    #run cdist in a loop over coords to avoid overloading memory
    values = np.zeros((xyz.shape[0]))
    for i in range(coords.shape[0]):
        dist = spatial.distance.cdist(coords[None,i], xyz)
        dist *= dist
        values += 1./np.sqrt(2*np.pi*sigma**2) * np.exp(-dist[0]/(2*sigma**2))
    return values.reshape(n,n,n)

image = sumofgauss(coords,xyz,1.0)

import matplotlib.pyplot as plt
plt.imshow(image[n/2]) #show a slice
plt.show()

M = 1000, N = 64 (~5 seconds):

M = 1000, N = 256 (~10 minutes):

Answer 1

Considering that many of your distance calculations will give zero weight after the exponential, you can probably drop a lot of your distances. Doing big chunks of distance calculations while dropping distances which are greater than a threshhold is usually faster with KDTree:

import numpy as np
from scipy.spatial import cKDTree # so we can get a `coo_matrix` output

def gaussgrid(coords, sigma = 1, n = 64, side = 100, eps = None):
    x_ = np.linspace(-side/2,side/2,n)
    x,y,z = np.meshgrid(x_,x_,x_,indexing='ij')
    xyz = np.column_stack((x.ravel(),y.ravel(),z.ravel()))
    if eps is None:
        eps = np.finfo('float64').eps
    thr = -np.log(eps) * 2 * sigma**2
    data_tree = cKDTree(coords)
    discr = 1000 # you can tweak this to get best results on your system
    values = np.empty(n**3)
    for i in range(n**3//discr + 1):
        slc = slice(i * discr, i * discr + discr)
        grid_tree = cKDTree(xyz[slc])
        dists = grid_tree.sparse_distance_matrix(data_tree, thr, output_type = 'coo_matrix')
        dists.data = 1./np.sqrt(2*np.pi*sigma**2) * np.exp(-dists.data/(2*sigma**2))
        values[slc] = dists.sum(1).squeeze()
    return values.reshape(n,n,n)

Now, even if you keep eps = None it'll be a bit faster as you're still returning about 10% your distances, but with eps = 1e-6 or so, you should get a big speedup. On my system:

%timeit out = sumofgauss(coords, xyz, 1.0)
1 loop, best of 3: 23.7 s per loop

%timeit out = gaussgrid(coords)
1 loop, best of 3: 2.12 s per loop

%timeit out = gaussgrid(coords, eps = 1e-6)
1 loop, best of 3: 382 ms per loop