Question

我在下面定义了相当简单的函数func1()，主要由两个for块组成。

它适用于较小的N值，但由于for块组合在一起，它会非常快速地增长到N>1000的分钟数。

如何使用numpy广播来改善此功能的性能？

import numpy as np
import time as t


def func1(A_triang, B_triang):
    aa = []
    for i, A_tr in enumerate(A_triang):
        for j, B_tr in enumerate(B_triang):
            # Absolute value of differences.
            abs_diff = abs(np.array(A_tr) - np.array(B_tr))
            # Store the sum of the differences and the indexes
            aa.append([sum(abs_diff), i, j])
    return aa


# Generate random data with the proper format
N = 500
A_triang = np.random.uniform(0., 20., (N, 3))
A_triang[:, 0] = np.ones(N)
B_triang = np.random.uniform(0., 20., (N, 3))
B_triang[:, 0] = np.ones(N)

# Call function.
s = t.clock()
aa = func1(A_triang, B_triang)
print(t.clock() - s)

Answer 1

这是一个NumPy broadcasting，利用indices_merged_arr_generic_using_cp的修改版本作为索引分配部分 -

import functools

# Based on https://stackoverflow.com/a/46135435/ @unutbu
def indices_merged_arr_generic_using_cp(arr):
    """
    Based on cartesian_product
    http://stackoverflow.com/a/11146645/190597 (senderle)
    """
    shape = arr.shape
    arrays = [np.arange(s, dtype='int') for s in shape]
    broadcastable = np.ix_(*arrays)
    broadcasted = np.broadcast_arrays(*broadcastable)
    rows, cols = functools.reduce(np.multiply, broadcasted[0].shape),\
                                                  len(broadcasted)+1
    out = np.empty(rows * cols, dtype=arr.dtype)
    start, end = rows, 2*rows
    for a in broadcasted:
        out[start:end] = a.reshape(-1)
        start, end = end, end + rows
    out[0:rows] = arr.flatten()
    return out.reshape(cols, rows).T

def func1_numpy_broadcasting(a,b):
    val = np.abs(a[:,None,:] - b).sum(-1)
    return indices_merged_arr_generic_using_cp(val)

如果输入的第一列始终为1s，那么我们将不需要计算它们的差异，因为它们的差异总和将为零。因此，为了获得val，我们可以简单地使用最后两列 -

val = np.abs(a[:,1,None] - b[:,1]) + np.abs(a[:,2,None] - b[:,2])

这会节省内存，因为我们不会以这种方式3D。

使用numexpr模块 -

import numexpr as ne
def func1_numexpr_broadcasting(a,b):
    a3D = a[:,None,:]
    val = ne.evaluate('sum(abs(a3D - b),2)')
    return indices_merged_arr_generic_using_cp(val)

利用第一列为1s的事实，我们会有 -

def func1_numexpr_broadcasting_v2(a,b):
    a1 = a[:,1,None]
    b1 = b[:,1]
    a2 = a[:,2,None]
    b2 = b[:,2]
    val = ne.evaluate('abs(a1-b1) + abs(a2-b2)')
    return indices_merged_arr_generic_using_cp(val)

提高numpy块组合的性能

1 个答案: