如何优化多维numpy数组计算？

Question

给定一个5维数组，objective是计算提取的两个数组之间的差异。为简单起见，假设测量第二个位置的差异，可以表示为 bt 和 lf。这两个数组的值可以提取如下：

arr[ep, bt, mt, bd, :] - arr[ep, lf, mt, bd, :]

请注意，在上面，第一 (ep)、第三 (mt) 和第四 (bd) 轴的索引对于两个数组都是相同的，只有第二个轴的位置索引不同（bt 和 lf）。

基于这个需求，提出了如下代码，并打包在函数nested_for_loop下：

import numpy as np
from joblib import Parallel, delayed
np.random.seed(0)

ub_lb_pair = np.tril_indices (5, -1)

arr = np.random.randn(3, 5, 4, 3, 2)
my_shape = arr.shape

def nested_for_loop():
    store_cval = np.full([my_shape[0], 10, my_shape[2], my_shape[3], my_shape[4]],
                         np.nan)  # preallocate
    for ep in range(0, my_shape[0]):
        for mt in range(0, my_shape[2]):
            for bd in range(0, my_shape[3]):
                for idx,(bt, lf) in enumerate(zip(ub_lb_pair[0], ub_lb_pair[1])):
                    store_cval[ep, idx, mt, bd, :] = arr[ep, bt, mt, bd, :] - \
                                                     arr[ep, lf, mt, bd, :]
    return store_cval


store_cval = nested_for_loop()

不过，如果可能的话，我想使代码更加紧凑和高效。

我能想到的一种方法是利用 joblib parallel 模块，它可以实现如下所示的功能 multi_prop.

def multi_prop(my_arr, ep):
    temp_ = np.full([10, my_shape[2], my_shape[3], my_shape[4]],
                    np.nan)
    for mt in range(0, my_shape[2]):
        for bd in range(0, my_shape[3]):
            for idx, (bt, lf) in enumerate(zip(ub_lb_pair[0], ub_lb_pair[1])):
                temp_[idx, mt, bd, :] = my_arr[ep, bt, mt, bd, :] - my_arr[ep, lf, mt, bd, :]
                x = 1
    return  temp_

dist_abs = Parallel(n_jobs=-1)(delayed(multi_prop)(arr, ep) for ep in range(0, my_shape[0]))

dist_abs = np.array(dist_abs)
bb = np.array_equal(store_cval, dist_abs)

但是，我想知道这是否是一种更 numpythonic 的方式来实现相同的 objective。

Answer 1

你根本不需要任何循环。想象一下这对奇特的指数：

bt, lf = np.tril_indices (5, -1)

您正在寻找

store_cval = arr[:, bt] - arr[:, lf]

请记住，store_cval[ep, idx, mt, bd, :] = arr[ep, bt, mt, bd, :] - arr[ep, lf, mt, bd, :] 是对最后一个索引的隐式循环。它们都是循环，你不需要它们中的任何一个。

更通用的解决方案：

def diffs(arr, axis):
    a, b = np.tril_indices(arr.shape[axis], -1)
    ind1 = [slice(None) for _ in range(arr.ndim)]
    ind2 = ind1.copy()
    ind1[axis] = a
    ind2[axis] = b
    return arr[tuple(ind1)] - arr[tuple(ind2)]

如何优化多维numpy数组计算？

How to optimise multi dimension numpy array calculation?

python

performance

numpy

coding-efficiency