python - 在 Numpy 数组中投影坐标

Question

所以我在我的软件中遇到了一个相当大的瓶颈。我在 cords 中有一组坐标，其中每一行对应于 X,Y,Z 坐标。 cords 中的每个坐标在 atom_proj 中都有一个定义的区域。 atoms 变量对应于 cords 变量，并为 atom_proj 提供 key 。

我将坐标投影到 grid 数组上，然后旋转并重复直到满足旋转次数。我只投影 X 和 Z 坐标，忽略 Y。

下面是我的代码的简化版本。对于小坐标集和旋转次数，代码运行相对较快。但如果坐标集和旋转列表都很大，则可能需要很长时间。坐标的数量可以从几百到几万不等。我将 grid 上的区域投影到一定数量或旋转范围内以生成热图。下面还显示了坐标集的热图示例。

问题:

(i) - 如何减少坐标到矩阵的投影时间

(ii) - 是否有更 pythonic 的方式将坐标区域应用到 grid 而不是数组拼接？

import numpy as np
cords = np.array([[5,4,5],[5,4,3],[6,4,6]])
atoms = np.array([['C'],['H'],['C']])
atom_proj = {'H':np.array([[0,0,0,0,0],[0,0,1,0,0],[0,1,1,1,0],[0,0,1,0,0],[0,0,0,0,0]]),'C':np.array([[0,0,0,0,0,0,0],[0,0,0,0,0,0,0],[0,0,1,1,1,0,0],[0,0,1,1,1,0,0],[0,0,1,1,1,0,0],[0,0,0,0,0,0,0],[0,0,0,0,0,0,0]])}

grid = np.zeros((10,10))

for rot in xrange(1,10): 
    # This for loop would contain a list of list of rotations to apply which are calculated before hand.
    # apply rotation
    for values in zip(cords, atoms):
        atom_shape = np.shape(atom_proj[values[1][0]])
        rad = (atom_shape[0]-1)/2                                
        grid[values[0][2]-rad:values[0][2]+rad+1,values[0][0]-rad:values[0][0]+rad+1] += atom_proj[values[1][0]]    
print grid

热图:

Answer 1

像这样的东西应该适用于内部循环

extruded = np.zeros((N, 10,10))
extruded[range(N), cords[:,2], cords[:,0]] = 1

grid = np.zeros((10,10))
for atom, proj in atom_proj.iteritems():
    centers = extruded[atoms==atom].sum(0)
    projected = nd.convolve(centers, proj)
    grid += projected

一些注意事项:

• 还有一个循环，但它是通过长度为 2 的原子类型数组，而不是长度为 N 的单个原子数组。
• 我省略了 for rot in [] 循环，因为它在这里没有做任何事情，但它应该可以很好地适应。
• 它的工作原理是在每个原子的中心位置放置一个原子，放在一堆网格中。然后，对于每种原子类型，将这些原子类型全部添加。然后，正如@Joe 所建议的那样，原子投影与这些中心进行卷积。
• 为了测试，我的原子是一维的，你的是二维的。不确定这是不是故意的。
• 在下面，我还添加了第三个版本，这是您的算法，但具有我能够理解的变量，它称为 OP_simplified

这是完整的套件:

import numpy as np
import scipy.ndimage as nd

N = 1000
cords = np.random.randint(3, 7, (N, 3)) #np.array([[5,4,5],[5,4,3],[6,4,6]])
atoms = np.random.choice(list('HC'), N) #np.array([['C'],['H'],['C']])
atom_proj = {'H': np.array([[0,0,0,0,0],
                            [0,0,1,0,0],
                            [0,1,1,1,0],
                            [0,0,1,0,0],
                            [0,0,0,0,0]]),
             'C': np.array([[0,0,0,0,0,0,0],
                            [0,0,0,0,0,0,0],
                            [0,0,1,1,1,0,0],
                            [0,0,1,1,1,0,0],
                            [0,0,1,1,1,0,0],
                            [0,0,0,0,0,0,0],
                            [0,0,0,0,0,0,0]])}

def project_atom(cords, atoms, atom_proj):
    extruded = np.zeros((N, 10,10))
    extruded[range(N), cords[:,2], cords[:,0]] = 1
    grid = np.zeros((10,10))
    for atom, proj in atom_proj.iteritems():
        grid += nd.convolve(extruded[atoms.squeeze()==atom].sum(0), proj, mode='constant')
    return grid

def OP_simplified(cords, atoms, atom_proj):
    rads = {atom: (proj.shape[0] - 1)/2 for atom, proj in atom_proj.iteritems()}
    grid = np.zeros((10,10))
    for (x,y,z), atom in zip(cords, atoms):
        rad = rads[atom]
        grid[z-rad:z+rad+1, x-rad:x+rad+1] += atom_proj[atom]
    return grid

def OP(cords, atoms, atom_proj):
    grid = np.zeros((10,10))
    for values in zip(cords, atoms):
        atom_shape = np.shape(atom_proj[values[1][0]])
        rad = (atom_shape[0]-1)/2
        grid[values[0][2]-rad:values[0][2]+rad+1,values[0][0]-rad:values[0][0]+rad+1] += atom_proj[values[1][0]]
    return grid

有效!

In [957]: np.allclose(OP(cords, atoms, atom_proj), project_atom(cords, atoms, atom_proj))
Out[957]: True

和时间:

In [907]: N = 1000

In [910]: timeit OP(cords, atoms, atom_proj)
10 loops, best of 3: 30.7 ms per loop

In [911]: timeit project_atom(cords, atoms, atom_proj)
100 loops, best of 3: 2.97 ms per loop

In [913]: N = 10000

In [916]: timeit project_atom(cords, atoms, atom_proj)
10 loops, best of 3: 33.3 ms per loop

In [917]: timeit OP(cords, atoms, atom_proj)
1 loops, best of 3: 314 ms per loop