Эффективная сумма гауссианов в 3D с NumPy с использованием больших массивов

Question

У меня есть массив трехмерных координат M x 3, координаты (M ~ 1000-10000), и я хотел бы вычислить сумму гауссиан с центром в этих координатах в трехмерном массиве сетки сетки. Трехмерный массив ячеистой сетки, как правило, имеет размер 64 x 64 x 64, но иногда может быть больше 256 x 256 x 256 и может увеличиваться. Для начала я следовал этому вопросу , преобразовав свой массив meshgrid в массив координат N x 3, xyz , где N равно 64 ^ 3 или 256 ^ 3 и т. Д. Однако для больших размеров массива требуется слишком много памяти, чтобы векторизовать весь расчет (понятно, так как он может приблизиться к элементам 1e11 и потреблять терабайт ОЗУ), поэтому я разбил его на цикл по М координатам. Однако это слишком медленно.

Мне интересно, есть ли способ ускорить это вообще без перегрузки памяти. Преобразуя сетку сетки в xyz, я чувствую, что потерял какое-либо преимущество равномерного расположения сетки, и что каким-то образом, может быть, с помощью scipy.ndimage, я смогу воспользоваться преимуществами равномерного расстояния, чтобы ускорить процесс.

Вот мой начальный старт:

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()

М = 1000, N = 64 (~ 5 секунд):

М = 1000, N = 256 (~ 10 минут):

Answer 1

Учитывая, что многие из ваших вычислений расстояния приведут к нулевому весу после экспоненты, вы, вероятно, можете сбросить большую часть ваших расстояний. Выполнение больших блоков вычислений расстояния при отбрасывании расстояний, превышающих пороговое значение, обычно быстрее с 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)

Теперь, даже если вы сохраните eps = None, это будет немного быстрее, поскольку вы все равно возвращаете около 10% своих расстояний, но с eps = 1e-6 или около того вы должны получить большое ускорение. В моей системе:

%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