Как правильно определить функцию, имеющую массивы numpy в качестве аргумента и тип возвращаемого значения?

Question

На мой взгляд, есть два варианта. (Здесь numpy импортируется как np)

1. Использование лямбда-выражения

Это можно использовать, когда мы имеем дело с простой функцией. Скажем, я хочу реализовать функцию f: (x, y) = (x ^ 2, x + y). Тогда я ожидал бы, что подобное сработает:

f = lambda x,y : x**2, x + y

f = lambda x,y : (x**2, x + y)

, когда я позвоню print(f(np.array([[1],[3]],float))) или print(f(np.array([1,3],float))). Но я не могу заставить работать любой вариант этих лямбда-выражений.

2. Определение «правильной» функции

Это, однако, работает:

def f(vec):
    return np.array([vec[0] ** 2, vec[0] + vec[1]])

Но является ли синтаксис оптимальным? И как бы я это сделал, если я не хочу, чтобы вышеуказанная «правильная» функция была векторизована. В целом, как мне определить numpy -значную функцию, принимающую numpy -значения массива?

Answer 1

На самом деле не имеет значения, используете ли вы lambda или def или делаете вычисления с локально определенными переменными. При использовании функций количество входных аргументов должно соответствовать определению.

Но давайте продемонстрируем на простых интерактивных примерах

Определение двух переменных:

In [121]: x, y = 1,3                                                                           
In [122]: x**2, x+y                                                                            
Out[122]: (1, 4)             # returns a tuple

Определение списка :

In [125]: vec = [1,3]                                                                          
In [126]: vec[0]**2, vec[0]+vec[1]                                                             
Out[126]: (1, 4)             # same tuple

или массив:

In [127]: arr = np.array([1,3])                                                                

[126] будет работать с этим тоже, но индексация массива будет медленнее.

Массивы лучше всего, когда мы делать вещи со всем массивом, а не только с элементами. Например, мы можем возвести в квадрат все члены массива или суммировать их.

In [128]: arr**2, arr.sum()                                                                    
Out[128]: (array([1, 9]), 4)

Эти вычисления будут работать с массивом любого размера, а не только с формой (2,).

Любой из вышеперечисленных кортежей может быть преобразован в массив:

In [133]: np.array([x**2, x+y])                                                                
Out[133]: array([1, 4])

Answer 2

Если вы знаете, что входные данные всегда будут массивом длины 2, вы можете использовать:

 f = lambda x : (x[0]**2, x[0] + x[1])

Тогда приведенные вами примеры также будут работать.