带有3D图的分段函数

Question

我在使用 np.piecewise 进行多维绘图时遇到了麻烦，因为出现了广播错误。

有没有人知道怎么解决这个问题？

这是我在一个简化的可执行脚本中写的内容：

import numpy as np
from pylab import *
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d, Axes3D 

num_steps = 100
x_arr = np.linspace(0,100, num_steps)
y_arr = np.linspace(0,20, num_steps)

def zfunc(x, y):
    return np.piecewise(x, [x>=500, x<500], [x, -x])

x,y = np.meshgrid(x_arr, y_arr)
z =zfunc(x,y)

fig=plt.figure()
ax=fig.subplot(1,1,1,projection='3d')
p = x.plot_surface(x,y,z,rstride=1,cstride=1,cmap=cm.coolwarm,linewidth=0,antialiased=False)
plt.show()

这段代码会出现以下错误：

 return np.piecewise(x, [x>=500, x<500], [x, -x])
  File "C:\Python27\lib\site-packages\numpy\lib\function_base.py", line 716, in piecewise
    y[condlist[k]] = item
ValueError: array is not broadcastable to correct shape

Answer 1

查看你正在使用的函数的文档说明通常是个好主意。我在那儿找到了这个解决方案。

np.piecewise(x, [x>=500, x<500], [lambda x: x, lambda x: -x])

funclist：可调用对象的列表，f(x,args,*kw)，或者是标量值。
每个函数会在条件为真的地方对x进行评估。它应该接受一个数组作为输入，并返回一个数组或一个标量值作为输出。如果提供的是标量值而不是可调用对象，那么就会假设是一个常量函数（lambda x: scalar）。