python绘制3D图形

学习深度学习入门时，想绘制一下书上
$$f(x_0,x_1)=x_0^2+x_1^2$$
这个公式的图形

但是仅靠pyplot无法完成3D图的绘制

参考了python画二元函数的图像成功完成绘制

def function_2(x0,x1):
    return x0**2+x1**2


fig = plt.figure()
ax = Axes3D(fig)

x0 = np.arange(-10,10,0.01)
x1 = np.arange(-10,10,0.01)
X, Y = np.meshgrid(x0, x1)#网格的创建，这个是关键
Z = function_2(X,Y) #用X,Y做参数
ax.plot_surface(X,Y,Z,color='b')
plt.show()

图型如下：