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

Question

给定一个 5D 数组，目标是计算提取的两个数组之间的差异。为简单起见，假设测量第二个位置的差异，可以表示为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 方式。

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)]