python - 在 numpy 矩阵上应用自定义函数

Question

给定一个像 my_function(x,y) 这样的函数，它接受两个 ndarrays x 和 y 作为输入并输出一个标量:

    def my_function(x,y):
        perm = np.take(x, y)
        return np.sum((np.power(2, perm) - 1) / (np.log2(np.arange(3, k + 3))))

我想找到一种方法将其应用于两个矩阵 r 和 p

    r = np.asarray([[5,6,7],[8,9,10]])
    p = np.asarray([[2,1,0],[0,2,1]])

以这样的方式返回 ndarray 和值

    np.asarray([my_function([5,6,7],[2,1,0]), my_function([8,9,10],[0,2,1])

Answer 1

您可以稍微修改您的函数以使用 take_along_axis 而不是 take，这将使您能够适应 2D 解决方案。

<小时/>

def my_function_2d(x, y, k=1):
    t = np.take_along_axis(x, y, -1)
    u = np.power(2, t) - 1
    v = np.log2(np.arange(3, k+3))
    return (u / v).sum(-1)

my_function_2d(r, p, k=1)

array([ 139.43547554, 1128.73332914])

<小时/>

验证

In [96]: k = 1

In [97]: my_function([5,6,7],[2,1,0])
Out[97]: 139.4354755392921

In [98]: my_function([8,9,10],[0,2,1])
Out[98]: 1128.7333291393375

<小时/>

这也适用于 1D 情况:

In [145]: my_function_2d(r[0], p[0], k=1)
Out[145]: 139.4354755392921

<小时/>

此方法可推广到 N 维情况:

In [157]: r = np.random.randint(1, 5, (2, 2, 2, 2, 2, 3))

In [158]: p = np.random.randint(0, r.shape[-1], r.shape)

In [159]: my_function_2d(r, p, k=3)
Out[159]:
array([[[[[ 8.34718483, 14.25597598],
          [12.25597598, 19.97868221]],

         [[12.97868221,  4.68481893],
          [ 2.42295943,  1.56160631]]],


        [[[23.42409467,  9.82346582],
          [10.93124418, 16.42409467]],

         [[23.42409467,  1.56160631],
          [ 3.68481893, 10.68481893]]]],



       [[[[15.97868221, 10.93124418],
          [ 5.40752517, 14.93124418]],

         [[ 4.14566566,  6.34718483],
          [14.93124418,  3.68481893]]],


        [[[ 9.20853795, 13.39462286],
          [23.42409467,  3.82346582]],

         [[23.42409467,  9.85293763],
          [ 4.56160631, 10.93124418]]]]])

<小时/>

_{我假设您意识到您的方法并不适用于所有输入和 k，有一些形状要求}