форма вывода силы для функции NumPy

Question

У меня есть функция, которая принимает 2 входных массива в качестве входных данных и возвращает число (реальная функция намного сложнее, чем в примере ниже):

def the_func(x,y): 
    return np.sum(x)*np.sum(y)

>>> the_func([2,3,4],[10,11,12])
297

Пока все хорошо.

Но тогда

>>> the_func([[1,2,3],[4,5,6]],[[1,2,3],[4,5,6]])
441

вместо ожидаемого [36,225]

Поэтому я попытался vectorize исходная функция:

vfunc = np.vectorize(the_func)

Но

>>> vfunc([[1,2,3],[4,5,6]],[[1,2,3],[4,5,6]])
array([[ 1,  4,  9],
       [16, 25, 36]])

Как мне заставить эту функцию возвращать мой результат в форме, которую я хотел бы видеть (в последнем примере [36,225] = [(1+2+3)*(1+2+3),(4+5+6)*(4+5+6)])?

Answer 1

Вы должны использовать ключевое слово axis для этого, чтобы указать вместе с измерением, которое должно суммироваться. В вашем примере измените свою функцию следующим образом:

def the_func(x,y): 
    return np.sum(x, axis=1)*np.sum(y, axis=1)

Запуск его на вводе сэмпла дал результат [36 255].