如何在 numpy 中计算非平凡的外积和,而无需创建临时变量

51 阅读3分钟

给定一组 N 个单位向量和另一组 M 个向量,要求计算每个单位向量与其每个 M 个向量的点积的绝对值的平均值。实质上,这是计算两个矩阵的外积,然后求和并取平均值,其间插入了一个绝对值。当 N 和 M 不太大时,这不是很难,并且有很多方法可以进行(见下文)。问题在于当 N 和 M 很大的时候,创建的临时变量会非常大,并且对这种方法提出了实际限制。能否在不创建临时变量的情况下完成此计算?主な困难是由于存在绝对值。是否有用于“穿插”此类计算的通用技术?

作为一个例子,请考虑以下代码:

N = 7
M = 5

# 创建单位向量,只是为了我们有一些例子,这并不是为了优雅
phi = np.random.rand(N)*2*np.pi
ctheta = np.random.rand(N)*2 - 1
stheta = np.sqrt(1-ctheta**2)
nhat = np.array([stheta*np.cos(phi), stheta*np.sin(phi), ctheta]).T

# 创建其他向量
m = np.random.rand(M,3)

# 计算我们想要的数量,这里使用广播。
S = np.average(np.abs(np.sum(nhat*m[:,np.newaxis,:], axis=-1)), axis=0)

这很好,S 现在是一个长度为 N 的数组,其中包含所需的结果。不幸的是,在这个过程中我们已经创建了一些潜在的巨大数组。结果为:

np.sum(nhat*m[:,np.newaxis,:], axis=-1)

是一个 M X N 数组。最终结果当然只有 N 个大小。开始增加 N 和 M 的大小,我们很快就会遇到内存错误。如上所述,如果不需要绝对值,那么我们可以按照以下步骤进行,现在使用einsum():

T = np.einsum('ik,jk,j', nhat, m, np.ones(M)) / M

即使对于相当大的 N 和 M,这也可行并且快速工作。对于我需要包含 abs() 的特定问题,但更通用的解决方案(或许是一个更通用的ufunc)也将是有趣的。

2、解决方案

2.1 使用 Cython

根据一些评论,似乎使用 Cython 是最好的方法。我一直没有愚蠢地研究过使用 Cython。原来生成工作代码相对容易。在一些搜索之后,我整理了以下 Cython 代码。这不是最通用的代码,可能不是编写它的最佳方式,并且可能可以提高效率。即便如此,它只比原问题中的 einsum() 代码慢约 25%,所以并不会太差!它是为了明确地处理按原问题创建的数组而编写的(因此推断输入数组的模式)。

尽管存在注意事项,它确实为原始问题提供了一个相当有效的解决方案,并且可以在类似情况下作为起点。

import numpy as np
cimport numpy as np
import cython
DTYPE = np.float64
ctypedef np.float64_t DTYPE_t
cdef inline double d_abs (double a) : return a if a >= 0 else -a

@cython.boundscheck(False)
@cython.wraparound(False)
def process_vectors (np.ndarray[DTYPE_t, ndim=2, mode="fortran"] nhat not None,
                     np.ndarray[DTYPE_t, ndim=2, mode="c"] m not None) :
    if nhat.shape[1] != m.shape[1] :
        raise ValueError ("Arrays must contain vectors of the same dimension")
    cdef Py_ssize_t imax = nhat.shape[0]
    cdef Py_ssize_t jmax = m.shape[0]
    cdef Py_ssize_t kmax = nhat.shape[1] # same as m.shape[1]
    cdef np.ndarray[DTYPE_t, ndim=1] S = np.zeros(imax, dtype=DTYPE)
    cdef Py_ssize_t i, j, k
    cdef DTYPE_t val, tmp
    for i in range(imax) :
        val = 0
        for j in range(jmax) :
            tmp = 0
            for k in range(kmax) :
                tmp += nhat[i,k] * m[j,k]
            val += d_abs(tmp)
        S[i] = val / jmax
    return S

2.2 使用 root mean square

我不认为除了 Cython 之外,还有任何简单的方法可以加快您的确切操作。但您可能想要考虑您是否真的需要计算您正在计算的内容。因为如果您不是使用绝对值均值,而是使用均方根,您仍然可以在完全计算得到:

rms = np.sqrt(np.einsum('ij,il,kj,kl,k->i', nhat, nhat, m, m, np.ones(M)/M))

这与执行以下操作相同:

rms_2 = np.sqrt(np.average(np.einsum('ij,kj->ik', nhat, m)**2, axis=-1))

是的,这并不是您所要求的,但我担心这是您使用矢量化方法所能获得的。如果您决定走这条路,请查看 np.einsum 在 N 和 M 很大的情况下表现如何:当传入太多的参数和索引时,它往往会陷入困境。