如何根据NumPy数组元素值切换函数并绘制曲面图？

针对数组条件判断导致的"真值不明确"错误，无需手动循环，用numpy的向量化操作即可高效处理，以下是两种可行方案：

方案1：使用np.where实现元素级分支 #

np.where支持对数组每个元素做条件判断，分别执行对应计算，直接生成结果数组：

import numpy as np

def function(param, x_1, x_2):
    # 对每个x_2元素判断：>=0时取x_2，否则取x_2²
    term = np.where(x_2 >= 0, x_2, x_2 ** 2)
    est = param[0] + param[1] * (x_1 + term)
    return est

方案2：使用布尔掩码分别计算 #

通过布尔数组筛选符合条件的元素，分情况计算后合并结果：

import numpy as np

def function(param, x_1, x_2):
    # 初始化和输入形状一致的结果数组
    est = np.empty_like(x_1)
    # 生成布尔掩码：标记x_2 >= 0的位置
    mask_positive = x_2 >= 0
    # 对满足条件的位置计算对应值
    est[mask_positive] = param[0] + param[1] * (x_1[mask_positive] + x_2[mask_positive])
    # 对x_2 < 0的位置计算另一种情况
    est[~mask_positive] = param[0] + param[1] * (x_1[~mask_positive] + x_2[~mask_positive] ** 2)
    return est

错误原因说明 #

你之前的if x_2 > 0是针对单个标量的判断逻辑，但当x_2为numpy数组时，x_2 > 0会返回布尔数组，Python无法直接用整个布尔数组作为if的判断条件，因此抛出"真值不明确"错误。上述两种方法均为元素级操作，完全适配numpy数组输入。

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