polytools.py 文件源码

def polyroots(p, realroots=False, condition=lambda r: True):
    """
    Returns the roots of a polynomial with coefficients given in p.
      p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]
    INPUT:
    p - Rank-1 array-like object of polynomial coefficients.
    realroots - a boolean.  If true, only real roots will be returned  and the
        condition function can be written assuming all roots are real.
    condition - a boolean-valued function.  Only roots satisfying this will be
        returned.  If realroots==True, these conditions should assume the roots
        are real.
    OUTPUT:
    A list containing the roots of the polynomial.
    NOTE:  This uses np.isclose and np.roots"""
    roots = np.roots(p)
    if realroots:
        roots = [r.real for r in roots if isclose(r.imag, 0)]
    roots = [r for r in roots if condition(r)]

    duplicates = []
    for idx, (r1, r2) in enumerate(combinations(roots, 2)):
        if isclose(r1, r2):
            duplicates.append(idx)
    return [r for idx, r in enumerate(roots) if idx not in duplicates]