python - 3D 距离矢量化

Question

我需要矢量化此代码的帮助。现在，N=100，运行需要一分钟左右。我想加快速度。我已经为双循环做过这样的事情，但从来没有用过 3D 循环，而且我遇到了困难。

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

Answer 1

meshgrid/where/indexing 解决方案已经非常快了。我让它快了大约 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]])

每一列都是一个 3D 向量。这八个向量中的每一个都代表一个 1x1x1 立方体(步长为 1，所有维度的长度为 1 的 3D 网格)的一个角。

让我们称这个数组为 vectors(它包含代表网格中所有点的所有 3D 矢量)。然后，准备一个 bool 掩码来选择那些满足您的 mod2 标准的向量:

    mod2bool = np.sum(vectors, axis=0) % 2 == 0

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

    vectorsubset = vectors[:,mod2bool]

这会选择所有行 (:) 并使用 bool 索引来过滤列，两者都是 numpy 中的快速操作。使用原生 numpy 方法计算剩余向量的长度:

    lengths = np.sqrt(np.sum(vectorsubset**2, axis=0))

这相当快——但是，scipy.stats.ss 和 bottleneck.ss 可以比这更快地执行平方和计算。

使用您的说明转换长度:

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

这涉及有限数除以零，从而在输出数组中产生 Inf。这完全没问题。我们使用 numpy 的 errstate 上下文管理器来确保这些零除法不会引发异常或运行时警告。

现在对有限元求和(忽略 infs)并返回总和:

    return  np.sum(p[np.isfinite(p)])

下面我已经实现了两次这个方法。一次与刚刚解释的完全一样，一次涉及 bottleneck 的 ss 和 nansum 函数。我还添加了您的比较方法，以及跳过 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。这是预期的吗？