使用 numpy 和 matplotlib，函数允许数字(float 或 int)或 numpy 数组作为参数似乎很常见，如下所示:
import numpy as np
print np.sin(0)
# 0
x = np.arange(0,4,0.1)
y = np.sin(x)
在这个例子中，我用一个整数参数调用 np.sin 一次，然后用一个 numpy 数组 x 调用一次。我现在想写一个允许类似处理的函数，但我不知道怎么做。例如:
def fun(foo, n):
a = np.zeros(n)
for i in range(n):
a[i] = foo
return a
将允许我像 fun(1, 5) 那样调用 fun，但不能像 fun(x, 5) 那样调用。当然，我的实际计算要复杂得多。
如何初始化 a 以便我可以将简单数字或整个数字数组作为元素？
非常感谢您的帮助!
请您参考如下方法:
内置 numpy 函数通常以
def foo(a, ...):
a = np.asarray(a)
...
也就是说，它们将输入参数转换为数组(如果它已经是数组，则不复制)。这允许他们使用标量和列表。
一旦参数是一个数组，它就有一个形状并且可以针对其他参数进行广播。
在您的示例中，不清楚当 foo 是数组时应该发生什么
def fun(foo, n):
a = np.zeros(n)
for i in range(n):
a[i] = foo
return a
a 被初始化为 dtype float 数组。这意味着 a[i]=foo 仅在 foo 是单个元素编号(标量，可能是单个元素数组)时才有效。如果 foo 是一个具有多个值的数组，您可能会收到有关尝试使用序列设置元素的错误。
a[i] 是 a[i,...] 的缩写。那就是它在第一维上建立索引。因此，如果 a 的初始值正确，它可以接受数组作为输入(遵守广播规则)。
如果 a 初始化为 np.zeros(n, dtype=object)，则 a[i]=foo 将与任何东西，因为它 a 只包含指向 Python 对象的指针。
np.frompyfunc 是一种从函数生成数组的方法。但是它返回一个 dtype=object 的数组。 np.vectorize 使用它，但可以让您更好地控制输出类型。但两者都适用于标量。数组(如果作为参数给出)将逐个元素传递给函数。
