制作標量 python function 接受列表和 arrays

Question

In python, is there an easy and efficient way to make function f(x: float) accept both lists and numpy arrays as arguments (in which case I would want to apply f element-wise and return the result in the same format as it已發送）？ 現在，我只需要一維 arrays。

作為一個例子（我的真實f更復雜），假設我有：

def f(x):
    return math.log(x) if x > 0 else 0.0

然后這個工作，但不是那么優雅 - 也可能不是那么有效，因為遞歸（我使用它，因為它允許我只有一個函數）：

def f(x):
    if np.isscalar(x):
        return math.log(x) if x > 0 else 0.0
    elif isinstance(x, np.ndarray):
        return np.array([f(i) for i in x], dtype=float)
    else:
        return [f(i) for i in x]

有沒有更好的辦法？

Answer 1

如果您需要的不僅僅是代碼庫中的單個 function 定義，那么使用裝飾器 function 將是一個不錯的選擇。 這確實意味着您必須包裝您的代碼，但我假設您所說的“只有一個函數”的意思是您確實希望 f(x) 中的 function 表達式在所有情況下都相同。

使用您現有的代碼，裝飾器 function 將如下所示：

def elementwise_possible(func):
    def wrapper(x):
        if np.isscalar(x):
            return func(x)
        elif isinstance(x, np.ndarray):
            return np.array([func(i) for i in x], dtype=float)
        else:
            return [func(i) for i in x]
    return wrapper 

你會這樣寫你的 function ：

@elementwise_possible
def f(x):
    return math.log(x) if x > 0 else 0.0

結果 output 變為

In[2]: A = f(2)
In[3]: A
Out[3]: 0.6931471805599453

In[4]: B = f(np.array([2,3,4]))
In[5]: B
Out[5]: array([0.69314718, 1.09861229, 1.38629436])

In[6]:C = f([5,6,7])
In[7]:C
Out[7]: [1.6094379124341003, 1.791759469228055, 1.9459101490553132]

我認為效率應該是一樣的。