Newton fractal

newton_fractal.py

import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

N = 2500
epsilon_x = 1
epsilon_y = 1
center_position = [0,0]

solutions = np.array([2.7693, 0.11535-0.58974J, 0.11535+0.58974J])
residual_criteria = 0.1
max_step = 100

def f(z):
    return z**3-3*z**2+z-1

def diff_f(z):
    return 3*z**2-6*z+1

def find_nearest_solution_idx_and_residual(z):
    nearest_solution_idx = abs(z-solutions).argmin()
    residual = abs(z - solutions[nearest_solution_idx])
    
    return nearest_solution_idx, residual 

def solve_f(inital_z):
    z = inital_z
    for n in range(max_step):
        z = z - f(z) / diff_f(z)
        nearest_solution_idx, residual = find_nearest_solution_idx_and_residual(z)
        if residual < residual_criteria: #convergence judgment
            return nearest_solution_idx + 2*n/max_step

    return len(solutions)

def main():
    grid = np.empty((N, N), dtype=np.float32)
    points_x = np.linspace(center_position[0]-epsilon_x, center_position[0]+epsilon_x, N)
    points_y = np.linspace(center_position[1]-epsilon_y, center_position[1]+epsilon_y, N)

    for n, ReZ in tqdm(enumerate(points_x)):
        for m, ImZ in enumerate(points_y):
            inital_Z = ReZ + ImZ*1J
            grid[m, n] = solve_f(inital_Z)

    show(grid)

def show(grid):
    plt.figure(figsize=(15, 15))
    plt.axis("off")
    plt.tick_params(labelbottom=False, labelleft=False, labelright=False, labeltop=False)
    plt.imshow(grid, interpolation="nearest", cmap="gnuplot_r", vmax=len(solutions))
    plt.savefig("save/img.png", bbox_inches="tight", pad_inches=0)
    plt.show()

if __name__ == "__main__":
    main()