我正在研究牛顿分形的实现，这是我目前所做的：

import numpy as np
import matplotlib.pyplot as plt

f = np.poly1d([1,0,0,0,1,0,-1]) # x^3 - 1
fp = np.polyder(f)
rts = f.r
def newton(z0, f, fp, MAX_IT=1000):
    TOL = 1e-8
    z = z0
    for i in range(MAX_IT):
        dz = np.divide(f(z),fp(z))
        adz = np.absolute(dz)
        if np.any(adz < TOL) :
            return z
        z -= dz
    return False

npts = 1000
x = np.linspace(-1,1,npts)
y = np.linspace(-1,1,npts)
xx, yy = np.meshgrid(x, y)
pic = newton(np.ravel(xx+yy*1j),f,fp)

我从this答案中得到灵感，并修改了牛顿方法的实现。我的问题是pic有复杂的条目，而plt.imshow不接受这种类型的数据。
我尝试在pic上使用2*cos(np.angle(zz))，得到的结果是：

这是一种工作，但不是真的。这个就是我想要的：

所以我的问题是我可以将我的数据转换成plt.imshow可以接受的格式而不丢失信息吗？如果不能，我可以使用什么替代方法来获得相同的结果？我更希望保留方法的向量化，因为它更快，更有效（在我看来）。