r - 爱因斯坦求和的 for 循环的更快替代方案？

Question

我有一段从 Python 移植到 R 的代码。原始 Python 版本使用 np.einsum。由于我在 R 中找不到等效的 np.einsum，并且我想确定我理解它，所以我直接使用 for 循环对其进行了编码。现在，我想知道是否有更快的选择。

示例代码：

n = 2 ; d = 3 ; nx = 4 ; v = 5

array4d <- array( runif(n*nx*v*d ,-1,0),
                  dim = c(n, nx, v, d) )

array3d <- array( runif(n*v*d ,-1,0),
                  dim = c(n, v, d) )

einsum_result <- array( rep(0, n*nx*d),
                        dim = c(n, nx, d))

# original Python version: np.einsum('ikl,ijkl->ijl', array3d, array4d, optimize=False)
# R version
for (i in 1: n) {
    for( j in 1: nx) {
        for( l in 1: d ) { 
            einsum_result[i, j, l] <- einsum_result[i, j, l] +
                sum( array3d[i, , l] * array4d[i, j, , l])
        }}}

我尝试使用矩阵乘法删除j循环（因为j/nx通常是最大的数字），但无法正确解决。任何建议表示赞赏！

标签： rmultidimensional-arraymatrix-multiplication

Answer 1

sweep好的，我在使用and消除 j 循环方面取得了一些进展rowSums，因此将其发布为答案。我增加了数组的维度，因为原始问题中的小数组的好处并不是很明显。

n = 5 ; d = 10; nx = 400; v = 15

array4d <- array( runif(n*nx*v*d ,-1,0),
                  dim = c(n, nx, v, d) )

array3d <- array( runif(n*v*d ,-1,0),
                  dim = c(n, v, d) )

einsum_result1 <- array( rep(0, n*nx*d),
                        dim = c(n, nx, d))
einsum_result2 <- einsum_result1
    
microbenchmark({
    for (i in 1: n) {
        for( j in 1: nx) {
            for( l in 1: d ) { 
                einsum_result1[i, j, l] <- einsum_result1[i, j, l] +
                    sum( array3d[i, , l] * array4d[i, j, , l])
            }}} },
    {
    for (i in 1: n) {
        #for( j in 1: nx) {
        for( l in 1: d ) {     
            einsum_result2[i, , l] <- einsum_result2[i, , l] +
                rowSums( sweep(array4d[i, , , l], MARGIN=2, array3d[i, , l], `*`) )
        }}})

     min       lq     mean   median       uq       max neval cld
 40.53148 42.64406 48.30703 47.65992 50.87022 179.93073   100   b
 10.86891 11.09062 12.02787 11.22610 11.56781  27.29123   100  a 
> identical(einsum_result1, einsum_result2)
[1] TRUE

有人看到删除任何剩余循环并用矢量化代码替换的方法吗？我已经把这个问题留了一段时间，但如果没有建议，我想我会接受我自己的答案