基于索引函数的数组向量化操作答案

【问题标题】：Vectorized manipulation of an array, based on a function of indices基于索引函数的数组向量化操作
【发布时间】：2014-01-24 23:23:32
【问题描述】：

我有一个表示 3D 点之间的函数的数组。因此，作为索引，它得到 6 个元组。现在我想对这个数组的元素应用一个函数，但这个函数不仅取决于元素的值，还取决于它的索引。因此，如果A 是矩阵，m 和 n 是 A[m,n] 存储其值的 3D 点，k 是介于 0 和 3 之间的值，则 f(A,k)[m,n] 等于：

-m[k]**2 如果m==n

-m[k]**2-n[k]**2 否则

以下是我的代码：

import numpy as np
def func(a,k):
    b=np.empty(a.shape)
    for i in range(a.flatten().size):
        ind=np.unravel_index(i,a.shape)
        if ind[0:3]==ind[3:6]:
            b[ind]=a[ind]*ind[0:3][k]**2
        else:
            b[ind]=a[ind]*(ind[0:3][k]**2-ind[3:6][k]**2)
    return b
a=np.arange(729).reshape((3,3,3,3,3,3))
print func(a,2)

是否有对这段代码进行 vecotrizing 的方法？

附：这是我实际需要做的简化版本。

【问题讨论】：

标签： python arrays numpy vectorization

【解决方案1】：

使用numpy.indices()创建索引数组，然后就可以对计算进行vecotrizing：

import numpy as np
def func(a,k):
    b=np.empty(a.shape)
    for i in range(a.flatten().size):
        ind=np.unravel_index(i,a.shape)
        if ind[0:3]==ind[3:6]:
            b[ind]=a[ind]*ind[0:3][k]**2
        else:
            b[ind]=a[ind]*(ind[0:3][k]**2-ind[3:6][k]**2)
    return b

def func2(a,k):
    b = np.empty(a.shape)
    ind = np.indices(a.shape).reshape(6, -1)
    mask = np.all(ind[:3] == ind[3:6], axis=0)
    ar = a.ravel()
    br = b.ravel()
    br[mask] = ar[mask]*ind[k, mask]**2
    mask = ~mask
    br[mask] = ar[mask]*(ind[k, mask]**2 - ind[3+k, mask]**2)
    return b

a = np.arange(729).reshape((3,3,3,3,3,3))
b1 = func(a, 2)
b2 = func2(a, 2)
np.allclose(b1, b2)

这是 %timeit 结果：

%timeit func(a, 2)
%timeit func2(a, 2)

输出：

100 loops, best of 3: 16.4 ms per loop
1000 loops, best of 3: 579 µs per loop

您可以根据自己的情况对其进行一些优化：

def func3(a,k):
    b = np.empty(a.shape)
    ind = np.indices(a.shape).reshape(6, -1)
    mask = ~np.all(ind[:3] == ind[3:6], axis=0)
    ar = a.ravel()
    br = b.ravel()
    br[:] = ar*ind[k]**2
    br[mask] -= ar[mask]*ind[3+k, mask]**2
    return b

【讨论】：