使用np.vectorize时数组元素访问报错，需适配输入返回对应类型

首先咱们先拆解你遇到的问题根源：@np.vectorize这个装饰器的工作逻辑是把输入数组拆成单个标量元素，逐个传入你的函数处理，再把所有结果重新拼成数组。这意味着在被装饰的recovery函数内部，你处理的始终是单个的float类型s，而不是数组。

你之前的错误操作（比如return ans[0]、len(ans)）之所以报错，是因为此时simp(y,x)返回的是一个标量（对一维数组做辛普森积分的结果就是单个数值），而不是数组——标量没有索引、长度这些属性，自然会触发各类错误。

好在你的核心需求（输入float返回float，输入np.ndarray返回数组），可以通过两种清晰的方式实现：

方案1：去掉np.vectorize，手动处理输入类型 #

这种方式更直观，也避免了vectorize带来的隐性性能问题（它本质是循环，效率并没有原生数组操作高）：

import numpy as np
from scipy.integrate import simpson  # 旧版scipy可能是simp，注意对应替换

def recovery(s, R0, n):
    # 处理单个float输入的情况
    if isinstance(s, float):
        x = np.linspace(s, 1, n)
        y = 1 / (x * R0(x))
        return simpson(y, x)
    # 处理numpy数组输入的情况
    elif isinstance(s, np.ndarray):
        results = []
        for s_val in s:
            x = np.linspace(s_val, 1, n)
            y = 1 / (x * R0(x))
            results.append(simpson(y, x))
        return np.array(results)
    # 处理不支持的输入类型
    else:
        raise TypeError("输入s必须是float或numpy.ndarray类型")

方案2：正确使用np.vectorize（如果坚持保留的话） #

如果你还是想保留vectorize，要记住：函数内部只需要处理单个标量输入的逻辑，不需要对结果做数组操作——vectorize会自动帮你把单个结果拼成数组。

你原来的代码（直接返回ans）其实已经满足需求了：

• 当输入是单个float时，函数返回float；
• 当输入是np.ndarray时，vectorize会遍历每个元素，把所有ans（标量）拼成数组返回。

如果确实需要对积分结果做额外处理（比如你的simp函数特殊，返回的是数组），可以在函数内部先判断结果类型再操作：

import numpy as np
from scipy.integrate import simpson

@np.vectorize
def recovery(s, R0, n):
    x = np.linspace(s, 1, n)
    y = 1 / (x * R0(x))
    ans = simpson(y, x)
    # 仅当ans是数组时取第一个元素，否则直接返回标量
    return ans[0] if isinstance(ans, np.ndarray) else ans

测试验证 #

用你的示例输入测试一下（假设R0是一个示例函数）：

s = np.array([0.2,0.3,0.4])
# 示例R0函数，可替换为你的实际函数
def R0(x):
    return x + 1

print(recovery(s, R0, 100))  # 返回数组
print(recovery(0.2, R0, 100))  # 返回float

这样就能完美匹配你的需求啦！

内容的提问来源于stack exchange，提问作者dontknowwhatimdoing