我现在想用Matplotlib绘制一个复杂函数的3D图，设定x是复数z的实部，y是虚部，z轴是f(z)的辐角（arg(f(z))），同时想用曲面的颜色来表示f(z)的模（mod(f(z))）。我写了下面这段代码，但颜色显示完全不对：

import matplotlib
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import matplotlib.pyplot as plt

x = np.linspace(-10, 10, 150)
y = np.linspace(-10, 10, 150)

X, Y = np.meshgrid(x, y)
arg = np.arctan2(fy(X, Y), fx(X, Y))
mod = np.sqrt(self.fx(X, Y)**2 + self.fy(X, Y)**2)
            
surface = ax.plot_surface(X, Y, arg, rstride=1, cstride=1, facecolors = plt.cm.hsv(mod), linewidth=0, antialiased=False)

这里的fx和fy分别是复函数的实部和虚部，ax是通过fig.add_subplot(...)创建的坐标轴对象（fig = plt.figure(...)）。

比如绘制f(z)=z²的时候，得到的图颜色看起来根本不代表模长：

我还做了个测试，把plt.cm.hsv(mod)换成plt.cm.hsv(X)，同时让fx(X,Y)=X，fy(X,Y)=Y，结果得到的图也完全不是我预期的——我本来以为会是颜色沿x方向渐变，等色线垂直于x轴的彩虹效果，但实际输出完全不对：

我实在搞不懂哪里出问题了，有没有大佬能指点一下？非常感谢！

备注：内容来源于stack exchange，提问作者Luke__