小编Chr*_*ood的帖子

如何最好地编写一个函数,可以接受标量浮点数或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 …