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

Code before:
    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:

Review comment (suggestion):
Perhaps instead of defining `f` for every test case there could just be a `check` function (taking the function) which short-hands the `weierstrass_method` and `allclose`

Code after:
    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 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