如何将牛顿分形的复数值网格数据转为实值且不丢失信息？

我正在实现牛顿分形，目前的代码如下：

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)

参考Stack Overflow的回答修改了牛顿法实现，但pic包含复数值，plt.imshow不支持该类型数据。尝试用2*cos(np.angle(zz))处理后效果不理想，我希望得到目标图的效果。请问能否将数据转换为plt.imshow可接受的格式且不丢失信息？若不行，有哪些替代方法能达到目标？我希望保留向量化实现以保证效率。

方案一：映射到收敛根的索引（最常用的牛顿分形可视化方式） #

牛顿分形的核心价值就是展示每个初始点最终收敛到哪个根，把收敛点匹配到对应根的索引，生成整数矩阵即可直接用imshow渲染，还能保留迭代次数这类额外信息，且全程向量化不损失效率：

1. 先修正并完善你的newton函数（原函数会因np.any提前终止，无法处理批量点）：

import numpy as np
import matplotlib.pyplot as plt

# 注意：原注释写的x³-1，但系数是x⁶+x³-1，这里先修正为x³-1，若要原多项式可改回原系数
f = np.poly1d([1, 0, 0, -1])
fp = np.polyder(f)
rts = f.r  # 获取多项式的所有复根

def newton(z0, f, fp, MAX_IT=100):
    TOL = 1e-8
    z = z0.copy()
    # 记录每个点的迭代次数
    it_counts = np.zeros(z.shape, dtype=int)
    # 标记已收敛的点
    converged = np.zeros(z.shape, dtype=bool)
    
    for i in range(MAX_IT):
        # 只对未收敛的点进行迭代
        mask = ~converged
        if not np.any(mask):
            break
        dz = np.divide(f(z[mask]), fp(z[mask]))
        z[mask] -= dz
        # 更新收敛状态
        converged[mask] = np.abs(dz) < TOL
        it_counts[mask] = i + 1
    
    # 匹配每个收敛点到最近的根，取索引
    dists = np.abs(z[:, np.newaxis] - rts)
    root_idx = np.argmin(dists, axis=1)
    # 未收敛的点标记为-1
    root_idx[~converged] = -1
    
    return root_idx.reshape(npts, npts), it_counts.reshape(npts, npts)

2. 生成可视化数据并渲染：

npts = 1000
x = np.linspace(-1.5, 1.5, npts)
y = np.linspace(-1.5, 1.5, npts)
xx, yy = np.meshgrid(x, y)
z0 = np.ravel(xx + yy*1j)

root_idx, it_counts = newton(z0, f, fp)

# 基础版：按收敛根的索引着色
plt.figure(figsize=(10,10))
plt.imshow(root_idx, extent=[x.min(), x.max(), y.min(), y.max()], cmap='viridis')
plt.colorbar(ticks=[0,1,2], label='收敛到的根索引')
plt.title('牛顿分形（按收敛根着色）')
plt.show()

# 增强版：结合迭代次数（用颜色深浅表示迭代次数）
plt.figure(figsize=(10,10))
colored = plt.cm.viridis(root_idx / root_idx.max())
# 用迭代次数调整亮度
colored[..., 0] *= it_counts / it_counts.max()
plt.imshow(colored, extent=[x.min(), x.max(), y.min(), y.max()])
plt.title('牛顿分形（结合收敛根与迭代次数）')
plt.show()

这个方案完全保留了所有关键信息：每个点收敛到哪个根、迭代了多少次，且全程向量化，效率拉满。

方案二：复数属性映射到RGB通道（保留完整复数信息） #

如果需要完整保留收敛点的复数信息，可以把复数的实部、虚部、相位等属性分别映射到RGB三个通道：

# 假设已经得到收敛后的z矩阵（reshape为npts×npts）
z_conv = z.reshape(npts, npts)
# 归一化后映射到RGB
r = (np.real(z_conv) - np.real(z_conv).min()) / (np.real(z_conv).max() - np.real(z_conv).min())
g = (np.imag(z_conv) - np.imag(z_conv).min()) / (np.imag(z_conv).max() - np.imag(z_conv).min())
b = (np.angle(z_conv) + np.pi) / (2 * np.pi)
rgb = np.stack([r, g, b], axis=-1)

plt.figure(figsize=(10,10))
plt.imshow(rgb)
plt.title('牛顿分形（复数RGB映射）')
plt.show()

这种方式能完整保留复数的所有信息，但视觉直观性不如方案一，适合需要分析收敛点具体数值的场景。

关键修正说明 #

• 原newton函数逻辑错误：np.any(adz < TOL)会在有一个点收敛时就直接返回，无法处理批量初始点，必须改成逐点跟踪收敛状态。
• 原多项式注释与系数不符：注释写的是x^3 - 1，但系数对应x⁶+x³-1，使用时注意修正一致。

内容的提问来源于stack exchange，提问作者fractal_girl