Find local minima and maxima of a 1D polynomial

main.py

import numpy as np

def find_local_min_max(poly_coefs: np.ndarray, atol: float | None = 1e-6):
    """
    Find the local minima and maxima of a polynomial function.

    Parameters
    ----------
    poly_coefs : np.ndarray
        The coefficients of the polynomial, ordered from lowest degree to highest.
    atol : float, optional
        The distance of the neighboring points of roots.
        Default is 1e-6.

    Returns
    -------
    tuple
        A tuple containing two arrays:
        - The real roots corresponding to the local minima of the polynomial.
        - The real roots corresponding to the local maxima of the polynomial.

    """
    # Create a polynomial object from the coefficients
    poly = np.polynomial.Polynomial(poly_coefs)

    # Find the real roots of the derivative polynomial
    real_roots = poly.deriv().roots().real

    # Find the local minima and maxima by checking if the polynomial values at the roots
    # are greater than or less than the values at the roots shifted by atol
    lower_bound = poly(real_roots) > poly(real_roots-atol)
    upper_bound = poly(real_roots) > poly(real_roots+atol)

    # Return the real roots corresponding to the local minima and maxima
    local_minima_args = np.logical_and(~lower_bound, ~upper_bound)
    local_maxima_args = np.logical_and(lower_bound, upper_bound)
    return real_roots[local_minima_args], real_roots[local_maxima_args]

if __name__ == "__main__":
    poly_coefs = np.array([0.911, 0.709, 0.341])
    print(find_local_min_max(poly_coefs))