Repository: TheAlgorithms/Python
File: maths/numerical_analysis/weierstrass_method.py (Python)

Code before:

    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)

Review comment (suggestion):
Lets add some more tests:
```suggestion
        >>> def f(x): return x**3 - 6*x**2 + 11*x - 6
        >>> roots = weierstrass_method(f, 3)
        >>> np.allclose(np.sort(roots), np.array([1, 2, 3]))
        True
```

Code after:

    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 check(poly, degree, expected):
        ...     roots = weierstrass_method(poly, degree)
        ...     return np.allclose(np.sort(roots), np.sort(expected))

        >>> check(
        ...     lambda x: x**2 - 1,
        ...     2,
        ...     np.array([-1, 1]))
        True

        >>> check(
        ...     lambda x: x**3 - 4.5*x**2 + 5.75*x - 1.875,
        ...     3,
        ...     np.array([1.5, 0.5, 2.5])
        ... )
        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."
            )