时间戳

首页 > 解决方案 > 在numpy中使用间接索引对多维数组求和的快速方法

python - 在numpy中使用间接索引对多维数组求和的快速方法

问题描述

计算以下表达式的最佳（最快）方法是什么：

\sum_{i \in I} \alpha_i \sum_{j \in J} \beta_j M[:, i, j] 对于给定的numpy数组：

在更一般的情况下，我需要为具有相同 M 数组的大量 alpha、beta、I、J 输入计算它。因此可以认为 alphas nad Is 具有形状 (N, 4)，betas 和 Js 具有形状 (N, 3)，我需要为范围 (N) 中的每个 n 计算此表达式。

先感谢您。

根据一些评论，为了使问题更清楚并添加一些上下文，这里是一种处理实际尺寸问题的幼稚方法：

M有形状(500, 200000, 20)

I有形状(10^6, 4)

J有形状(10^6, 3)

alpha有形状(10^6, 4)

beta有形状(10^6, 3) 

N = 10**6
M_new = np.zeros(M.shape[0], N)
for n in range(N):
    for i in range(4):
        for j in range(3):
            M_new[:, n] += alpha[n, i] * beta[n, j] * M[:, I[n, i], J[n, j]]

所以问题是如何尽可能快地计算 M_new。

解决方案

到目前为止，最快的解决方案是@jdehesa 使用 Numba 提出的解决方案。

@Han-KwangNienhuys 展示了与替代方法的速度比较。

标签： pythonnumpyscipy

解决方案

编辑：
我首先误解了这个问题。我在下面留下了原始答案，但它并没有按照问题的要求做。

冒着听起来很明显的风险，您总是可以求助于 Numba：

import numpy as np
import numba as nb

# Original loop implementation
def comb_loop(m, ii, jj, alpha, beta):
    n = ii.shape[0]
    m_new = np.zeros((m.shape[0], n))
    for col in range(n):
        for i in range(4):
            for j in range(3):
                m_new[:, col] += alpha[col, i] * beta[col, j] * m[:, ii[col, i], jj[col, j]]
    return m_new

# Numba implementation
@nb.njit(parallel=True)
def comb_nb(m, ii, jj, alpha, beta):
    n = ii.shape[0]
    m_new = np.empty((m.shape[0], n), m.dtype)
    for col in nb.prange(n):
        for row in range(m.shape[0]):
            val = 0
            for i in range(4):
                for j in range(3):
                    val += alpha[col, i] * beta[col, j] * m[row, ii[col, i], jj[col, j]]
            m_new[row, col] = val
    return m_new


# Test
np.random.seed(0)
N = 1_000  # Reduced for testing
m = np.random.rand(500, 200_000, 20)
ii = np.random.randint(m.shape[1], size=(N, 4))
jj = np.random.randint(m.shape[2], size=(N, 3))
alpha = np.random.rand(N, 4)
beta = np.random.rand(N, 3)

# Check results match
print(np.allclose(comb_loop(m, ii, jj, alpha, beta), comb_nb(m, ii, jj, alpha, beta)))
# True

# Timings
%timeit comb_loop(m, ii, jj, alpha, beta)
# 181 ms ± 1.77 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit comb_nb(m, ii, jj, alpha, beta)
# 31.1 ms ± 2.03 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

原来的错误答案

您可以使用np.einsum：

import numpy as np

def comb(alpha, beta, m):
    return np.einsum('i,j,nij->n', alpha, beta, m)

# Test
np.random.seed(0)
alpha = np.random.rand(10)
beta = np.random.rand(20)
m = np.random.rand(30, 10, 20)
result = comb(alpha, beta, m)
print(result.shape)
# (30,)

推荐阅读