python - 尝试使用 np.piecewise 进行多元函数时出现 ValueError

Question

我正在尝试使用 np.piecewise 定义一个多元分段函数，如下所示:

X = np.array([
    [1, 2],
    [3, 4], 
    [5, 6]
])

pw = np.piecewise(
    X,
    [
        np.abs(X[:, 0] - X[:, 1]) < 1,
        np.abs(X[:, 0] - X[:, 1]) >= 1
    ],
    [
        lambda X: 1 + 2 * X[:, 0] + 3 * X[:, 1],
        lambda X: 1.5 + 2.5 * X[:, 0] + 3.5 * X[:, 1]
    ]
)

运行此代码片段会出现以下错误:

ValueError: shape mismatch: value array of shape (3,) could not be broadcast to indexing result of shape (3,2)

对于上下文，我尝试在本例中表示一个映射 f: R^2 -> R，同时在 X 的每一行上对其进行评估。

有什么想法吗？我是否需要以不同的方式定义最终参数以便正确广播索引？

Answer 1

IMO np.piecewise 如果您有来自 np.meshgrid 的两个数组，则更适合，以便 np.piecewise 可以匹配条件的维度与您的数组维度。

在您的情况下，表示分段映射 $f:R^2\to R$ ，输入形状为 (n,2) 并逐行评估(每列代表一个变量)，生成矢量化代码的最简单方法就是使用 np.select:

def pw(X):
    return np.select([np.abs(X[:,0] - X[:,1]) < 1, np.abs(X[:,0] - X[:,1]) >= 1], 
            [1 + 2 * X[:, 0] + 3 * X[:, 1], 1.5 + 2.5 * X[:, 0] + 3.5 * X[:, 1]])

和pw(X)产生您想要的答案。