simple python script to make a newton fractal

newton_fractal.py

# I ran on python 3.12
# requires `pip install matplotlib numpy`
import matplotlib.pyplot as plt
import numpy as np


def get_f_df(c: list[int]) -> tuple[callable, callable]:
  """ Get the function and its derivative for a given polynomial c"""

  def f(z):
    ans = 0
    for i, c_i in enumerate(c):
      ans += c_i * z**(len(c) - i - 1)
    return ans
  
  def df(z):
    ans = 0
    for i, c_i in enumerate(c):
      ans += (len(c) - i - 1) * c_i * z**(len(c) - i - 2)
    return ans
  
  return f, df



def newton(
    x0: complex, f: callable, df: callable,
    tol: float=1e-12, maxiter: int=100):
  """ Newton's method for solving f(x) = 0 """
  x = x0
  for _ in range(maxiter):
    x = x - f(x)/df(x)
    if abs(f(x)) < tol:
      return x
  print("Newton's method did not converge")
  return x

def newton_fractal_grid(
    f: callable, df: callable,
    xlim: tuple[float, float],
    ylim: tuple[float, float],
    nx: int =100, ny: int=100,
    maxiter: int=100):
  """ Mesh of which root Newton's method converges to for each starting point
  
  Args:
    f: function
    df: derivative of f
    xlim: x limits - x values are the real part of the complex number to
      start Newton's method with
    ylim: y limits - y values are the imaginary part of the complex number to
      start Newton's method with
    nx: number of points in x direction
    ny: number of points in y direction
  """
  x = np.linspace(*xlim, nx)
  y = np.linspace(*ylim, ny)
  X, Y = np.meshgrid(x, y)
  Z = X + 1j*Y
  W = np.zeros(Z.shape, dtype=complex)
  print(W.shape)
  for i in range(ny):
    for j in range(nx):
      W[i,j] = newton(Z[i,j], f, df, maxiter=maxiter)
  return X, Y, W

def map_root_to_index(W, actual_roots):
  """ Map the roots to integers so that we can color the roots """
  Z = np.zeros(W.shape)
  for i, root in enumerate(actual_roots):
    Z[np.isclose(W, root, 1e-12, 1e-2)] = i
  return Z


def main():
  # Poly coefficients:
  #   z^3 - 1 is cool
  #   z^3 - x is cool
  c= [1, 0, -1, 0]
  f, df = get_f_df(c)

  # Actual roots
  actual_roots = np.roots(c)
  print(f"Actual roots: {actual_roots}")

  X, Y, W = newton_fractal_grid(
    f, df, xlim=(-0.57, -0.45), ylim=(-0.05, 0.05), nx=1000, ny=1000)
  Z = map_root_to_index(W, actual_roots)
  plt.pcolormesh(X, Y, Z, cmap='brg', antialiased=True)
  plt.show()


if __name__ == '__main__':
  main()