example using newton's method on x2 -4

example_x2_newton.py

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 i in range(maxiter):
    x = x - f(x)/df(x)
    #print(f"i: {i}, x: {x}, f(x): {f(x)}")
    if abs(f(x)) < tol:
      return x
  print("Newton's method did not converge")
  return x


def main():
  # Poly coefficients:
  c= [1, 0, -4]
  f, df = get_f_df(c)

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

  x1 = newton(-4, f, df)
  x2 = newton(4, f, df)
  print(f"Root 1: {x1}")
  print(f"Root 2: {x2}")

  x = np.linspace(-3, 3, 100)
  y = f(x)
  plt.plot(x, y, label='f(x)')
  plt.plot(
    actual_roots, np.zeros(actual_roots.shape[0]), 'ro', label='Actual roots')
  plt.plot(x1, 0, 'gx', label='Newton: Root 1')
  plt.plot(x2, 0, 'gx', label='Newton: Root 2')
  plt.grid()
  plt.legend()
  plt.show()


if __name__ == '__main__':
  main()