处理标量或数组的Python函数

Question

如何最好地编写一个函数,可以接受标量浮点数或numpy向量(1-d数组),并返回标量,1-d数组或2-d数组,具体取决于输入？

该函数很昂贵并且经常被调用,我不想给调用者增加负担来对参数或返回值进行特殊强制转换.它只需要处理数字(不是列表或其他可迭代的东西).

np.vectorize可能很慢(将一个python函数广播到numpy数组)和其他答案(让Python函数干净地返回一个标量或列表,具体取决于参数的数量)和np.asarray(一个python函数,它接受为参数(标量或numpy数组)无助于获取输出数组所需的维度.

这种类型的代码可以在Matlab,Javascript和其他语言中使用:

import numpy as np

def func( xa, ya ):
    # naively, I thought I could do:
    xy = np.zeros( ( len(xa), len(ya) ) )
    for j in range(len( ya )):
        for i in range(len( xa )):
            # do something complicated
            xy[i,j] = x[i]+y[j]            
    return xy

适用于数组:

x = np.array([1., 2.])
y = np.array([2., 4.])
xy = func(x,y)
print xy

[[ 3.  5.]
 [ 4.  6.]]

但是对于标量浮点数不起作用:

x = 1.
y = 3.
xy = func(x,y)
print xy

<ipython-input-64-0f85ad330609> in func(xa, ya)
      4 def func( xa, ya ):
      5     # naively, I thought I could do:
----> 6     xy = np.zeros( ( len(xa), len(ya) ) )
      7     for j in range(len( ya )):
      8         for i in range(len( xa )):

TypeError: object of type 'float' has no len()

在类似的函数中使用np.asarray给出:

<ipython-input-63-9ae8e50196e1> in func(x, y)
      5     xa = np.asarray( x );
      6     ya = np.asarray( y );
----> 7     xy = np.zeros( ( len(xa), len(ya) ) )
      8     for j in range(len( ya )):
      9         for i in range(len( xa )):

TypeError: len() of unsized object

什么是快速,优雅和pythonic方法？

Answer 1

在整个numpy代码库中,您可以找到以下内容:

def func_for_scalars_or_vectors(x):
    x = np.asarray(x)
    scalar_input = False
    if x.ndim == 0:
        x = x[None]  # Makes x 1D
        scalar_input = True

    # The magic happens here

    if scalar_input:
        return np.squeeze(ret)
    return ret

Answer 2

根据意见，我更希望函数对输入类型灵活，但始终返回一致的类型；这将最终阻止调用者检查返回类型（既定目标）。

例如，允许标量和/或数组作为参数，但始终返回数组。

def func(x, y):
    # allow (x=1,y=2) OR (x=[1,2], y=[3,4]) OR (!) (x=1,y=[2,3])
    xn = np.asarray([x]) if np.isscalar(x) else np.asarray(x)
    yn = np.asarray([y]) if np.isscalar(y) else np.asarray(y)

    # calculations with numpy arrays xn and xy
    res = xn + yn  # ..etc...
    return res

（不过，可以通过设置标志“ scalar=True”，yada yada yada轻松修改该示例以返回标量。但您还必须处理一个参数是标量，另一个是数组等；看起来像对我来说很多 YAGNI。）

Answer 3

我会做以下事情：

def func( xa, ya ):
    xalen = xa if type(xa) is not list else len(xa)
    yalen = ya if type(ya) is not list else len(ya)
    xy = np.zeros( (xalen, yalen) )    
    for j in range(yalen): 
        for i in range(xalen):
            xy[i,j] = x[i]+y[j] 
    return xy