将自定义函数应用于numpy矩阵

Question

Given a function like my_function(x,y) that takes two ndarrays x and y as an input and outputs a scalar:

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

I want to find a way to apply it to two matrices r and p

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

in such a way that an ndarray is returned with the values

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

Answer 1

You can slightly modify your function to use take_along_axis instead of take , which will allow you to adapt to the 2D solution.

---

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

---

Validation

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

---

This will also still work on the 1D case:

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

---

This approach generalizes to the N -dimensional case:

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

---

_{I assume you realize your approach doesn't work for all inputs and k s, there are some shape requirements}

Answer 2

You can try either map or a list comprehension with zip as following. Please note that I took k=1 to have a running code as you did not specify k

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


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

result = np.asarray([my_function(i, j) for i, j in zip(r, p)])
print (result)
# [ 139.43547554 1128.73332914]

Answer 3

You can use np.vectorize with the signature keyword:

k = 3
np.vectorize(my_function, signature='(i),(i)->()')(r, p)

# array([124.979052  , 892.46280834])