81 lines
2.4 KiB
Python
81 lines
2.4 KiB
Python
|
"""This submodule contains tools for working with numpy.poly1d objects."""
|
||
|
|
||
|
# External Dependencies
|
||
|
from __future__ import division, absolute_import
|
||
|
from itertools import combinations
|
||
|
import numpy as np
|
||
|
|
||
|
# Internal Dependencies
|
||
|
from .misctools import isclose
|
||
|
|
||
|
|
||
|
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]
|
||
|
|
||
|
|
||
|
def polyroots01(p):
|
||
|
"""Returns the real roots between 0 and 1 of the polynomial with
|
||
|
coefficients given in p,
|
||
|
p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]
|
||
|
p can also be a np.poly1d object. See polyroots for more information."""
|
||
|
return polyroots(p, realroots=True, condition=lambda tval: 0 <= tval <= 1)
|
||
|
|
||
|
|
||
|
def rational_limit(f, g, t0):
|
||
|
"""Computes the limit of the rational function (f/g)(t)
|
||
|
as t approaches t0."""
|
||
|
assert isinstance(f, np.poly1d) and isinstance(g, np.poly1d)
|
||
|
assert g != np.poly1d([0])
|
||
|
if g(t0) != 0:
|
||
|
return f(t0)/g(t0)
|
||
|
elif f(t0) == 0:
|
||
|
return rational_limit(f.deriv(), g.deriv(), t0)
|
||
|
else:
|
||
|
raise ValueError("Limit does not exist.")
|
||
|
|
||
|
|
||
|
def real(z):
|
||
|
try:
|
||
|
return np.poly1d(z.coeffs.real)
|
||
|
except AttributeError:
|
||
|
return z.real
|
||
|
|
||
|
|
||
|
def imag(z):
|
||
|
try:
|
||
|
return np.poly1d(z.coeffs.imag)
|
||
|
except AttributeError:
|
||
|
return z.imag
|
||
|
|
||
|
|
||
|
def poly_real_part(poly):
|
||
|
"""Deprecated."""
|
||
|
return np.poly1d(poly.coeffs.real)
|
||
|
|
||
|
|
||
|
def poly_imag_part(poly):
|
||
|
"""Deprecated."""
|
||
|
return np.poly1d(poly.coeffs.imag)
|