Skip to content

Commit 3c720e5

Browse files
Add weierstrass_method for approximating complex roots
- Implements Durand-Kerner (Weierstrass) method for polynomial root finding - Accepts user-defined polynomial function and degree - Uses random perturbation of complex roots of unity for initial guesses - Handles validation, overflow clipping, and includes doctest
1 parent 7a0fee4 commit 3c720e5

File tree

1 file changed

+85
-0
lines changed

1 file changed

+85
-0
lines changed
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
from collections.abc import Callable
2+
3+
import numpy as np
4+
5+
6+
def weierstrass_method(
7+
polynomial: Callable[[np.ndarray], np.ndarray],
8+
degree: int,
9+
roots: np.ndarray | None = None,
10+
max_iter: int = 100,
11+
) -> np.ndarray:
12+
"""
13+
Approximates all complex roots of a polynomial using the
14+
Weierstrass (Durand-Kerner) method.
15+
Args:
16+
polynomial: A function that takes a NumPy array of complex numbers and returns
17+
the polynomial values at those points.
18+
degree: Degree of the polynomial (number of roots to find). Must be ≥ 1.
19+
roots: Optional initial guess as a NumPy array of complex numbers.
20+
Must have length equal to 'degree'.
21+
If None, perturbed complex roots of unity are used.
22+
max_iter: Number of iterations to perform (default: 100).
23+
24+
Returns:
25+
np.ndarray: Array of approximated complex roots.
26+
27+
Raises:
28+
ValueError: If degree < 1, or if initial roots length doesn't match the degree.
29+
30+
Note:
31+
- Root updates are clipped to prevent numerical overflow.
32+
33+
Example:
34+
>>> import numpy as np
35+
>>> def f(x): return x**2 - 1
36+
>>> roots = weierstrass_method(f, 2)
37+
>>> np.allclose(np.sort(roots), np.sort(np.array([1, -1])))
38+
True
39+
40+
See Also:
41+
https://en.wikipedia.org/wiki/Durand%E2%80%93Kerner_method
42+
"""
43+
44+
if degree < 1:
45+
raise ValueError("Degree of the polynomial must be at least 1.")
46+
47+
if roots is None:
48+
# Use perturbed complex roots of unity as initial guesses
49+
rng = np.random.default_rng()
50+
roots = np.array(
51+
[
52+
np.exp(2j * np.pi * i / degree) * (1 + 1e-3 * rng.random())
53+
for i in range(degree)
54+
],
55+
dtype=np.complex128,
56+
)
57+
58+
else:
59+
roots = np.asarray(roots, dtype=np.complex128)
60+
if roots.shape[0] != degree:
61+
raise ValueError(
62+
"Length of initial roots must match the degree of the polynomial."
63+
)
64+
65+
for _ in range(max_iter):
66+
# Construct the product denominator for each root
67+
denominator = np.array([root - roots for root in roots], dtype=np.complex128)
68+
np.fill_diagonal(denominator, 1.0) # Avoid zero in diagonal
69+
denominator = np.prod(denominator, axis=1)
70+
71+
# Evaluate polynomial at each root
72+
numerator = polynomial(roots).astype(np.complex128)
73+
74+
# Compute update and clip to prevent overflow
75+
delta = numerator / denominator
76+
delta = np.clip(delta, -1e10, 1e10)
77+
roots -= delta
78+
79+
return roots
80+
81+
82+
if __name__ == "__main__":
83+
import doctest
84+
85+
doctest.testmod()

0 commit comments

Comments
 (0)