"""Return the roots of a polynomial."""
from __future__ import annotations
import numpy
import numpoly
from ..baseclass import PolyLike
# from ..dispatch import implements
# @implements(numpy.roots)
[docs]def roots(poly: PolyLike) -> numpy.ndarray:
"""
Return the roots of a polynomial.
Assumes the polynomial has a single dimension.
Args:
poly:
Polynomial to take roots of, or if constant, the coefficients of
said polynomial. This to be compatible with :func:`numpy.roots`.
Return:
An array containing the roots of the polynomial.
Raise:
ValueError:
When `poly` cannot be converted to a rank-1 polynomial.
Note:
The algorithm relies on computing the eigenvalues of the companion
matrix [1]_.
.. [1] R. A. Horn & C. R. Johnson, *Matrix Analysis*. Cambridge, UK:
Cambridge University Press, 1999, pp. 146-7.
Example:
>>> q0 = numpoly.variable()
>>> poly = 3.2*q0**2+2*q0+1
>>> numpoly.roots(poly)
array([-0.3125+0.46351241j, -0.3125-0.46351241j])
>>> numpoly.roots([3.2, 2, 1])
array([-0.3125+0.46351241j, -0.3125-0.46351241j])
"""
# backwards compatibility
poly = numpoly.aspolynomial(poly)
if poly.isconstant():
return numpy.roots(poly.tonumpy())
# only rank-1
poly = numpoly.clean_attributes(poly, retain_names=False)
if len(poly.names) > 1:
raise ValueError("polynomial is not of rank 1.")
# align exponents to include all coefficients
filled_basis = poly.indeterminants ** numpy.arange(
numpoly.lead_exponent(poly).item(), dtype=int
)
_, poly = numpoly.align_exponents(filled_basis, poly)
# pass coefficients to numpy
return numpy.roots(poly.coefficients[::-1])