3D距离向量化

3

我需要帮助将这段代码向量化。现在,当N=100时,它需要一分钟左右才能运行。我想加快速度。我已经为双重循环做了类似的事情,但从未使用过三维循环,并且遇到了困难。

import numpy as np
N = 100
n = 12
r = np.sqrt(2)

x = np.arange(-N,N+1)
y = np.arange(-N,N+1)
z = np.arange(-N,N+1)

C = 0

for i in x:
    for j in y:
        for k in z:
            if (i+j+k)%2==0 and (i*i+j*j+k*k!=0):
                p = np.sqrt(i*i+j*j+k*k)
                p = p/r
                q = (1/p)**n
                C += q

print '\n'
print C

1
@jonrsharpe:你看了问题吗? 代码是正确的,但很慢,所以OP希望将部分工作推给快速的numpy库,而不是在缓慢的Python中进行循环。 - Niklas B.
是的,代码是正确的,我只想将其向量化。我想学习如何使其尽可能高效。 - Wesley Bowman
1
ç­”و،ˆو¶‰هڈٹscipyçڑ„pdist,هڈ¯èƒ½è؟کو¶‰هڈٹbottleneckçڑ„ss(ه¹³و–¹ه’Œï¼Œن½ éœ€è¦په†چو¬،وڈگé«که®ƒçڑ„و½œهٹ›ï¼‰م€‚ - Dr. Jan-Philip Gehrcke
1
C 应该等于什么?我认为你可以轻松地使用 x,y,z = np.meshgrid(* [np.arange(-N,N + 1)] * 3) 来完成这个问题。 - wflynny
1
当 N=3 时,C 应该恰好为 12。 当 N=100 时,C 应该为 12.1318.... - Wesley Bowman
显示剩余2条评论
2个回答

2

感谢@Bill的帮助,我现在能够让它正常工作了。速度非常快。也许可以做得更好,特别是通过两个掩码来消除我最初为循环设置的两个条件。

    from __future__ import division
    import numpy as np

    N = 100
    n = 12
    r = np.sqrt(2)

    x, y, z = np.meshgrid(*[np.arange(-N, N+1)]*3)

    ind = np.where((x+y+z)%2==0)
    x = x[ind]
    y = y[ind]
    z = z[ind]
    ind = np.where((x*x+y*y+z*z)!=0)
    x = x[ind]
    y = y[ind]
    z = z[ind]

    p=np.sqrt(x*x+y*y+z*z)/r

    ans = (1/p)**n
    ans = np.sum(ans)
    print 'ans'
    print ans

我对这些条件(x+y+z)%2==0(x*x+y*y+z*z)!=0的来源很感兴趣。它们有任何物理意义吗?后者只是一个用于处理零值的过滤器吗?numpy可以处理这个问题。 - Dr. Jan-Philip Gehrcke
只是提醒一下,distances = distance.pdist(vectors, 'sqeuclidean') 可能会显著加快你的 np.sqrt(x*x+y*y+z*z),所以也许你想在过滤之前这样做。 - Dr. Jan-Philip Gehrcke
如果我把它放在过滤之前,那么我该如何进行过滤呢?此外,这些条件来自于我正在研究的原子晶格的物理意义。这是固态物理学,这些是我用于晶格位点的条件,其中晶格被设置为仅当坐标相加为偶数时才会有一个原子,并且在原点没有任何东西,因为它是参考原子。 - Wesley Bowman
我很乐意听取任何关于如何改进这个的建议,但我更希望您提交一个答案,我会很高兴地选择它作为最佳答案。 - Wesley Bowman
好的,现在我收到了你的代码。我正在研究它 :-) - Dr. Jan-Philip Gehrcke

2
网格/where/索引解决方案已经非常快速。我使它快了约65%。虽然这不算太多,但我还是会逐步解释一下:

对于我来说,最简单的方法是将网格中所有3D向量作为一个大的2D 3 x M 数组中的列来处理。 meshgrid 是创建所有组合的正确工具(请注意,需要numpy版本>= 1.7才能使用3D meshgrid),而vstack + reshape 则将数据转换为所需的形式。例如:

>>> np.vstack(np.meshgrid(*[np.arange(0, 2)]*3)).reshape(3,-1)
array([[0, 0, 1, 1, 0, 0, 1, 1],
       [0, 0, 0, 0, 1, 1, 1, 1],
       [0, 1, 0, 1, 0, 1, 0, 1]])

每一列都是一个三维向量。这八个向量中的每一个代表一个 1x1x1 的立方体顶点(在所有维度上步长和长度均为1的三维网格)。
我们将这个数组称为 vectors(其中包含表示网格中所有点的所有三维向量)。然后,准备一个布尔掩码以选择符合您的mod2条件的那些向量:
    mod2bool = np.sum(vectors, axis=0) % 2 == 0

np.sum(vectors, axis=0)会创建一个包含每个列向量元素总和的1 x M数组。因此,mod2bool是一个1 x M数组,其中每个列向量都有一个布尔值。现在使用这个布尔掩码:

    vectorsubset = vectors[:,mod2bool]

这个代码选择所有行(:),并使用布尔索引来过滤列,这两个操作在numpy中非常快。使用原生的numpy方法计算剩余向量的长度:
    lengths = np.sqrt(np.sum(vectorsubset**2, axis=0))

这很快 - 但是,scipy.stats.ssbottleneck.ss甚至可以比此更快地执行平方和计算。

按照您的说明转换长度:

    with np.errstate(divide='ignore'):
        p = (r/lengths)**n

这涉及到有限数除以零,导致输出数组中出现Inf。这是完全可以接受的。我们使用numpy的errstate上下文管理器来确保这些零除法不会引发异常或运行时警告。
现在将有限元素求和(忽略无穷大),并返回总和:
    return  np.sum(p[np.isfinite(p)])

我已经实现了以下两种方法。一种完全像刚才解释的那样,另一种涉及到瓶颈的ssnansum函数。我还添加了您的方法进行比较,以及跳过np.where((x*x+y*y+z*z)!=0)索引的修改版本,而是创建Inf,最后用isfinite方式求和。

import sys
import numpy as np
import bottleneck as bn

N = 100
n = 12
r = np.sqrt(2)


x,y,z = np.meshgrid(*[np.arange(-N, N+1)]*3)
gridvectors = np.vstack((x,y,z)).reshape(3, -1)


def measure_time(func):
    import time
    def modified_func(*args, **kwargs):
        t0 = time.time()
        result = func(*args, **kwargs)
        duration = time.time() - t0
        print("%s duration: %.3f s" % (func.__name__, duration))
        return result
    return modified_func


@measure_time
def method_columnvecs(vectors):
    mod2bool = np.sum(vectors, axis=0) % 2 == 0
    vectorsubset = vectors[:,mod2bool]
    lengths = np.sqrt(np.sum(vectorsubset**2, axis=0))
    with np.errstate(divide='ignore'):
        p = (r/lengths)**n
    return  np.sum(p[np.isfinite(p)])


@measure_time
def method_columnvecs_opt(vectors):
    # On my system, bn.nansum is even slightly faster than np.sum.
    mod2bool = bn.nansum(vectors, axis=0) % 2 == 0
    # Use ss from bottleneck or scipy.stats (axis=0 is default).
    lengths = np.sqrt(bn.ss(vectors[:,mod2bool]))
    with np.errstate(divide='ignore'):
        p = (r/lengths)**n
    return  bn.nansum(p[np.isfinite(p)])


@measure_time
def method_original(x,y,z):
    ind = np.where((x+y+z)%2==0)
    x = x[ind]
    y = y[ind]
    z = z[ind]
    ind = np.where((x*x+y*y+z*z)!=0)
    x = x[ind]
    y = y[ind]
    z = z[ind]
    p=np.sqrt(x*x+y*y+z*z)/r
    return np.sum((1/p)**n)


@measure_time
def method_original_finitesum(x,y,z):
    ind = np.where((x+y+z)%2==0)
    x = x[ind]
    y = y[ind]
    z = z[ind]
    lengths = np.sqrt(x*x+y*y+z*z)
    with np.errstate(divide='ignore'):
        p = (r/lengths)**n
    return  np.sum(p[np.isfinite(p)])


print method_columnvecs(gridvectors)
print method_columnvecs_opt(gridvectors)
print method_original(x,y,z)
print method_original_finitesum(x,y,z)

这是输出结果:

$ python test.py
method_columnvecs duration: 1.295 s
12.1318801965
method_columnvecs_opt duration: 1.162 s
12.1318801965
method_original duration: 1.936 s
12.1318801965
method_original_finitesum duration: 1.714 s
12.1318801965

所有方法产生的结果都是相同的。当使用isfinite风格求和时,您的方法会稍微快一些。我的方法更快,但我认为这是一种学术性质的练习,而不是重要的改进 :-)

我还有一个问题:您说对于N=3,计算应该产生12。即使您的方法也没有做到这一点。以上所有方法对于N=3都产生12.1317530867的结果。这是预期的吗?


我认为我原本想让N=1时给出12,这只是我用来测试是否正确的一个测试案例。非常感谢您的输入和帮助。我学到了很多以前没有想过的东西。谢谢! - Wesley Bowman

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接