Weierstrass Method

maths/numerical_analysis/weierstrass_method.py

@@ -0,0 +1,85 @@
+from collections.abc import Callable
+
+import numpy as np
+
+
+def weierstrass_method(
+    polynomial: Callable[[np.ndarray], np.ndarray],
+    degree: int,
+    roots: np.ndarray | None = None,
+    max_iter: int = 100,
+) -> np.ndarray:
+    """
+    Approximates all complex roots of a polynomial using the
+    Weierstrass (Durand-Kerner) method.
+    Args:
+        polynomial: A function that takes a NumPy array of complex numbers and returns
+                    the polynomial values at those points.
+        degree: Degree of the polynomial (number of roots to find). Must be ≥ 1.
+        roots:  Optional initial guess as a NumPy array of complex numbers.
+                Must have length equal to 'degree'.
+                If None, perturbed complex roots of unity are used.
+        max_iter: Number of iterations to perform (default: 100).
+
+    Returns:
+        np.ndarray: Array of approximated complex roots.
+
+    Raises:
+        ValueError: If degree < 1, or if initial roots length doesn't match the degree.
+
+    Note:
+        - Root updates are clipped to prevent numerical overflow.
+
+    Example:
+        >>> import numpy as np
+        >>> def f(x): return x**2 - 1
+        >>> roots = weierstrass_method(f, 2)
+        >>> np.allclose(np.sort(roots), np.sort(np.array([1, -1])))
+        True
+
+    See Also:
+        https://en.wikipedia.org/wiki/Durand%E2%80%93Kerner_method
+    """
+
+    if degree < 1:
+        raise ValueError("Degree of the polynomial must be at least 1.")
+
+    if roots is None:
+        # Use perturbed complex roots of unity as initial guesses
+        rng = np.random.default_rng()
+        roots = np.array(
+            [
+                np.exp(2j * np.pi * i / degree) * (1 + 1e-3 * rng.random())
+                for i in range(degree)
+            ],
+            dtype=np.complex128,
+        )
+
+    else:
+        roots = np.asarray(roots, dtype=np.complex128)
+        if roots.shape[0] != degree:
+            raise ValueError(
+                "Length of initial roots must match the degree of the polynomial."
+            )
+
+    for _ in range(max_iter):
+        # Construct the product denominator for each root
+        denominator = np.array([root - roots for root in roots], dtype=np.complex128)
+        np.fill_diagonal(denominator, 1.0)  # Avoid zero in diagonal
+        denominator = np.prod(denominator, axis=1)
+
+        # Evaluate polynomial at each root
+        numerator = polynomial(roots).astype(np.complex128)
+
+        # Compute update and clip to prevent overflow
+        delta = numerator / denominator
+        delta = np.clip(delta, -1e10, 1e10)
+        roots -= delta
+
+    return roots
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()