numpy.vectorize функции, которая возвращает массив

Question

Я пытаюсь векторизовать функцию, которая имеет 2 входа и выводит массив np.arpe, формы = (4,).Функция выглядит следующим образом:

def f(a, b):
    return np.array([a+b, a-b, a, b])

Мне удалось векторизовать функцию, используя параметр сигнатуры, однако она работает, только если я исключаю один из параметров, используя аргумент excluded np.vectorize:

Это работает:

vec = np.vectorize(f, signature='()->(n)', excluded=[1])
x = np.arange(5)
y = 3
vec(x, y)

>> output:
array([[ 3, -3,  0,  3],
       [ 4, -2,  1,  3],
       [ 5, -1,  2,  3],
       [ 6,  0,  3,  3],
       [ 7,  1,  4,  3]])

Однако, если я уберу аргумент excluded, все пойдет не так, как планировалось.

Это не работает:

vec = np.vectorize(f, signature='()->(n)')
x = np.arange(5)
y = 3
vec(x, y)

>> Error:
TypeError: wrong number of positional arguments: expected 1, got 2

Как сделать так, чтобы векторизованная функция могла получать массив / список входных значений для одного (или обоих) входных параметров?

Ожидаемый вывод будет функцией vec, которая позволит вызывать его с несколькими входами для любого из входных параметров.

Answer 1

In [237]: f1 = np.vectorize(f, signature='(),()->(n)')                          
In [238]: f1(np.arange(5),3)                                                    
Out[238]: 
array([[ 3, -3,  0,  3],
       [ 4, -2,  1,  3],
       [ 5, -1,  2,  3],
       [ 6,  0,  3,  3],
       [ 7,  1,  4,  3]])

In [241]: f1(np.arange(5),np.ones((4,5))).shape                                 
Out[241]: (4, 5, 4)
In [242]: f1(np.arange(5),np.ones((1,5))).shape                                 
Out[242]: (1, 5, 4)

---

frompyfunc возвращает массив dtype объекта:

In [336]: f2 = np.frompyfunc(f,2,1)                                             
In [337]: f2(np.arange(5), 3)                                                   
Out[337]: 
array([array([ 3, -3,  0,  3]), array([ 4, -2,  1,  3]),
       array([ 5, -1,  2,  3]), array([6, 0, 3, 3]), array([7, 1, 4, 3])],
      dtype=object)
In [338]: _.shape                                                               
Out[338]: (5,)

np.vectorize, без signature, использует frompyfunc, но добавляет свое собственное преобразование dtype.

---

In [340]: f1(np.arange(5), np.arange(3)) 
ValueError: shape mismatch: objects cannot be broadcast to a single shape

Сбой по той же причине, что и при добавлении следующего:

In [341]: np.arange(5)+np.arange(3)                                             
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-341-fb1c4f4372da> in <module>
----> 1 np.arange(5)+np.arange(3)

ValueError: operands could not be broadcast together with shapes (5,) (3,) 

Чтобы получить (5,3) результат, нам нужно создать первый аргумент (5,1):

In [342]: np.arange(5)[:,None]+np.arange(3)                                     
Out[342]: 
array([[0, 1, 2],
       [1, 2, 3],
       [2, 3, 4],
       [3, 4, 5],
       [4, 5, 6]])