автоматическое вычисление массива в python (numpy)

Question

У меня есть функция, которая зависит от нескольких переменных, скажем y=f(x1,x2,x3,x4). Если каждая из переменных является простым числом, то результат должен быть простым числом. Если одна из переменных является массивом, мне нужен результат, чтобы быть также массивом. И так далее: если две переменные являются массивами, мне нужен результат, чтобы быть двумерным массивом. Etc.

Пример:

def f(x1,x2,x3,x4): 
    y=x1*x2/(x3+x4) 
    return y 

x1=1.0 
x2=2.0 
x3=3.0 
x4=4.0 
f(x1,x2,x3,x4) 
# should give 2.0/7.0 = 0.2857... 

x3=array([1.0,2.0,3.0,4.0,5.0]) 
f(x1,x2,x3,x4) 
# should give a one-dimensional array with shape (5,) 

x4=array([10.0,20.0,30.0,40.0,50.0,60.0,70.0]) 
f(x1,x2,x3,x4) 
# should give a two-dimensional array with shape (5,7) 

Как это сделать? (Чтобы быть как можно более ясным для читателя, не использующего Python моей программы?)

Answer 1

Правильный способ сделать это - передать данные надлежащей формы. Если вы хотите получить результат 2d, вы должны передать в 2d массивы. Это может быть выполнено np.newaxis.

import numpy as np 

def f(x1,x2,x3,x4): 
    y = x1*x2/(x3+x4) 
    return y 

x1 = 1.0 
x2 = 2.0 
x3 = 3.0 
x4 = 4.0 
print f(x1,x2,x3,x4) 

x3 = np.array([1.0,2.0,3.0,4.0,5.0]) 
print f(x1,x2,x3,x4) 

x3 = x3[:, np.newaxis] 
x4 = np.array([10.0,20.0,30.0,40.0,50.0,60.0,70.0]) 
x4 = x4[np.newaxis, :] 
print f(x1,x2,x3,x4) 

Конечно, путь ваш вопрос поставлен, это немного неоднозначным, почему вы должны ожидать, чтобы получить массив в форме (5, 7), а не массив в форме (7, 5).

Answer 2

Вы можете посмотреть на использование @when, чтобы определить функцию соответствующим образом.

Например:

# Case for single values 
@when(int, int, int, int) 
def f(x1,x2,x3,x4): 
    y = x1*x2/(x3+x4) 
    return y 

# Case for 1D arrays 
@when(list, list, list, list) 
def f(x1,x2,x3,x4): 
    y = [] 
    for i in range(len(x1)): 
     y.append(x1[i]*x2[i]/(x3[i]+x4[i])) 
    return y