frac.py

import numpy as np
import matplotlib.pyplot as plt

def newton_fractal(xmin, xmax, ymin, ymax, width, height, max_iter, tol, power):
    x = np.linspace(xmin, xmax, width)
    y = np.linspace(ymin, ymax, height)
    X, Y = np.meshgrid(x, y)
    Z = X + 1j * Y
    roots = np.exp(2j * np.pi * np.arange(power) / power)
    F = np.zeros(Z.shape, dtype=int)
    for i in range(max_iter):
        Z = Z - (Z**power - 1) / (power * Z**(power - 1))
        for k, r in enumerate(roots):
            m = (np.abs(Z - r) < tol) & (F == 0)
            F[m] = k + 1
    return F

xmin, xmax, ymin, ymax = -2, 2, -2, 2
width, height, max_iter, tol, power = 800, 800, 50, 1e-6, 5
F = newton_fractal(xmin, xmax, ymin, ymax, width, height, max_iter, tol, power)
plt.imshow(F, extent=[xmin, xmax, ymin, ymax], cmap='hsv')
plt.axis('off')
plt.show()