Применить пользовательскую функцию на матрицей

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 с, есть некоторые требования к форме}

Answer 2

Вы можете использовать np.vectorize с ключевым словом signature:

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

# array([124.979052  , 892.46280834])

Answer 3

Вы можете попробовать либо map, либо список с помощью zip следующим образом.Обратите внимание, что я взял k=1, чтобы иметь рабочий код, поскольку вы не указали 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]