Source code for sympy.polys.rootoftools

"""Implementation of RootOf class and related tools. """

from __future__ import print_function, division

from sympy.core import (S, Expr, Integer, Float, I, oo, Add, Lambda,
symbols, sympify, Rational, Dummy)
from sympy.core.cache import cacheit
from sympy.core.function import AppliedUndef
from sympy.functions.elementary.miscellaneous import root as _root

from sympy.polys.polytools import Poly, PurePoly, factor
from sympy.polys.rationaltools import together
from sympy.polys.polyfuncs import symmetrize, viete

from sympy.polys.rootisolation import (
dup_isolate_complex_roots_sqf,
dup_isolate_real_roots_sqf)

from sympy.polys.polyroots import (
preprocess_roots, roots)

from sympy.polys.polyerrors import (
MultivariatePolynomialError,
GeneratorsNeeded,
PolynomialError,
DomainError)

from sympy.polys.domains import QQ

from mpmath import mpf, mpc, findroot, workprec
from mpmath.libmp.libmpf import dps_to_prec, prec_to_dps

from sympy.utilities import lambdify, public, sift

from sympy.core.compatibility import range, ordered

from math import log as mathlog

__all__ = ['CRootOf']

class _pure_key_dict(object):
"""A minimal dictionary that makes sure that the key is a
univariate PurePoly instance.

Examples
========

Only the following actions are guaranteed:

>>> from sympy.polys.rootoftools import _pure_key_dict
>>> from sympy import S, PurePoly
>>> from sympy.abc import x, y

1) creation

>>> P = _pure_key_dict()

2) assignment for a PurePoly or univariate polynomial

>>> P[x] = 1
>>> P[PurePoly(x - y, x)] = 2

3) retrieval based on PurePoly key comparison (use this
instead of the get method)

>>> P[y]
1

4) KeyError when trying to retrieve a nonexisting key

>>> P[y + 1]
Traceback (most recent call last):
...
KeyError: PurePoly(y + 1, y, domain='ZZ')

5) ability to query with in

>>> x + 1 in P
False

NOTE: this is a *not* a dictionary. It is a very basic object
for internal use that makes sure to always address its cache
via PurePoly instances. It does not, for example, implement
get or setdefault.
"""
def __init__(self):
self._dict = {}

def __getitem__(self, k):
if not isinstance(k, PurePoly):
if not (isinstance(k, Expr) and len(k.free_symbols) == 1):
raise KeyError
k = PurePoly(k, expand=False)
return self._dict[k]

def __setitem__(self, k, v):
if not isinstance(k, PurePoly):
if not (isinstance(k, Expr) and len(k.free_symbols) == 1):
raise ValueError('expecting univariate expression')
k = PurePoly(k, expand=False)
self._dict[k] = v

def __contains__(self, k):
try:
self[k]
return True
except KeyError:
return False

_reals_cache = _pure_key_dict()
_complexes_cache = _pure_key_dict()

def _pure_factors(poly):
_, factors = poly.factor_list()
return [(PurePoly(f, expand=False), m) for f, m in factors]

def _imag_count_of_factor(f):
"""Return the number of imaginary roots for irreducible
univariate polynomial f.
"""
terms = [(i, j) for (i,), j in f.terms()]
if any(i % 2 for i, j in terms):
return 0
# update signs
even = [(i, I**i*j) for i, j in terms]
even = Poly.from_dict(dict(even), Dummy('x'))
return int(even.count_roots(-oo, oo))

[docs]@public
def rootof(f, x, index=None, radicals=True, expand=True):
"""An indexed root of a univariate polynomial.

Returns either a ComplexRootOf object or an explicit

Parameters
==========

f : Expr
Univariate polynomial.
x : Symbol, optional
Generator for f.
index : int or Integer
Return a radical expression if possible.
expand : bool
Expand f.
"""

[docs]@public
class RootOf(Expr):
"""Represents a root of a univariate polynomial.

Base class for roots of different kinds of polynomials.
Only complex roots are currently supported.
"""

__slots__ = ['poly']

def __new__(cls, f, x, index=None, radicals=True, expand=True):
"""Construct a new CRootOf object for k-th root of f."""

[docs]@public
class ComplexRootOf(RootOf):
"""Represents an indexed complex root of a polynomial.

Roots of a univariate polynomial separated into disjoint
real or complex intervals and indexed in a fixed order.
Currently only rational coefficients are allowed.
Can be imported as CRootOf.

Examples
========

>>> from sympy import CRootOf, rootof
>>> from sympy.abc import x

CRootOf is a way to reference a particular root of a
polynomial. If there is a rational root, it will be returned:

>>> CRootOf.clear_cache()  # for doctest reproducibility
>>> CRootOf(x**2 - 4, 0)
-2

Whether roots involving radicals are returned or not
depends on whether the radicals flag is true (which is
set to True with rootof):

>>> CRootOf(x**2 - 3, 0)
CRootOf(x**2 - 3, 0)
>>> CRootOf(x**2 - 3, 0, radicals=True)
-sqrt(3)
>>> rootof(x**2 - 3, 0)
-sqrt(3)

The following cannot be expressed in terms of radicals:

>>> r = rootof(4*x**5 + 16*x**3 + 12*x**2 + 7, 0); r
CRootOf(4*x**5 + 16*x**3 + 12*x**2 + 7, 0)

The root bounds can be seen, however, and they are used by the
evaluation methods to get numerical approximations for the root.

>>> interval = r._get_interval(); interval
(-1, 0)
>>> r.evalf(2)
-0.98

The evalf method refines the width of the root bounds until it
guarantees that any decimal approximation within those bounds
will satisfy the desired precision. It then stores the refined
interval so subsequent requests at or below the requested
precision will not have to recompute the root bounds and will
return very quickly.

Before evaluation above, the interval was

>>> interval
(-1, 0)

After evaluation it is now

>>. r._get_interval()
(-165/169, -206/211)

To reset all intervals for a given polynomial, the _reset method
can be called from any CRootOf instance of the polynomial:

>>> r._reset()
>>> r._get_interval()
(-1, 0)

The eval_approx method will also find the root to a given
precision but the interval is not modified unless the search
for the root fails to converge within the root bounds. And
the secant method is used to find the root. (The evalf
method uses bisection and will always update the interval.)

>>> r.eval_approx(2)
-0.98

The interval needed to be slightly updated to find that root:

>>> r._get_interval()
(-1, -1/2)

The evalf_rational will compute a rational approximation
of the root to the desired accuracy or precision.

>>> r.eval_rational(n=2)
-69629/71318

>>> t = CRootOf(x**3 + 10*x + 1, 1)
>>> t.eval_rational(1e-1)
15/256 - 805*I/256
>>> t.eval_rational(1e-1, 1e-4)
3275/65536 - 414645*I/131072
>>> t.eval_rational(1e-4, 1e-4)
6545/131072 - 414645*I/131072
>>> t.eval_rational(n=2)
104755/2097152 - 6634255*I/2097152

========
eval_approx
eval_rational
_eval_evalf
"""

__slots__ = ['index']
is_complex = True
is_number = True

def __new__(cls, f, x, index=None, radicals=False, expand=True):
""" Construct an indexed complex root of a polynomial.

See rootof for the parameters.

The default value of radicals is False to satisfy
eval(srepr(expr) == expr.
"""
x = sympify(x)

if index is None and x.is_Integer:
x, index = None, x
else:
index = sympify(index)

if index is not None and index.is_Integer:
index = int(index)
else:
raise ValueError("expected an integer root index, got %s" % index)

poly = PurePoly(f, x, greedy=False, expand=expand)

if not poly.is_univariate:
raise PolynomialError("only univariate polynomials are allowed")

degree = poly.degree()

if degree <= 0:
raise PolynomialError("can't construct CRootOf object for %s" % f)

if index < -degree or index >= degree:
raise IndexError("root index out of [%d, %d] range, got %d" %
(-degree, degree - 1, index))
elif index < 0:
index += degree

dom = poly.get_domain()

if not dom.is_Exact:
poly = poly.to_exact()

roots = cls._roots_trivial(poly, radicals)

if roots is not None:
return roots[index]

coeff, poly = preprocess_roots(poly)
dom = poly.get_domain()

if not dom.is_ZZ:
raise NotImplementedError("CRootOf is not supported over %s" % dom)

root = cls._indexed_root(poly, index)
return coeff * cls._postprocess_root(root, radicals)

@classmethod
def _new(cls, poly, index):
"""Construct new CRootOf object from raw data. """
obj = Expr.__new__(cls)

obj.poly = PurePoly(poly)
obj.index = index

try:
_reals_cache[obj.poly] = _reals_cache[poly]
_complexes_cache[obj.poly] = _complexes_cache[poly]
except KeyError:
pass

return obj

def _hashable_content(self):
return (self.poly, self.index)

@property
def expr(self):
return self.poly.as_expr()

@property
def args(self):
return (self.expr, Integer(self.index))

@property
def free_symbols(self):
# CRootOf currently only works with univariate expressions
# whose poly attribute should be a PurePoly with no free
# symbols
return set()

def _eval_is_real(self):
"""Return True if the root is real. """
return self.index < len(_reals_cache[self.poly])

def _eval_is_imaginary(self):
"""Return True if the root is imaginary. """
if self.index >= len(_reals_cache[self.poly]):
ivl = self._get_interval()
return ivl.ax*ivl.bx <= 0  # all others are on one side or the other
return False  # XXX is this necessary?

@classmethod
def real_roots(cls, poly, radicals=True):
"""Get real roots of a polynomial. """
return cls._get_roots("_real_roots", poly, radicals)

@classmethod
def all_roots(cls, poly, radicals=True):
"""Get real and complex roots of a polynomial. """
return cls._get_roots("_all_roots", poly, radicals)

@classmethod
def _get_reals_sqf(cls, factor, use_cache=True):
"""Get real root isolating intervals for a square-free factor."""
if use_cache and factor in _reals_cache:
real_part = _reals_cache[factor]
else:
_reals_cache[factor] = real_part = \
dup_isolate_real_roots_sqf(
factor.rep.rep, factor.rep.dom, blackbox=True)

return real_part

@classmethod
def _get_complexes_sqf(cls, factor, use_cache=True):
"""Get complex root isolating intervals for a square-free factor."""
if use_cache and factor in _complexes_cache:
complex_part = _complexes_cache[factor]
else:
_complexes_cache[factor] = complex_part = \
dup_isolate_complex_roots_sqf(
factor.rep.rep, factor.rep.dom, blackbox=True)
return complex_part

@classmethod
def _get_reals(cls, factors, use_cache=True):
"""Compute real root isolating intervals for a list of factors. """
reals = []

for factor, k in factors:
try:
if not use_cache:
raise KeyError
r = _reals_cache[factor]
reals.extend([(i, factor, k) for i in r])
except KeyError:
real_part = cls._get_reals_sqf(factor, use_cache)
new = [(root, factor, k) for root in real_part]
reals.extend(new)

reals = cls._reals_sorted(reals)
return reals

@classmethod
def _get_complexes(cls, factors, use_cache=True):
"""Compute complex root isolating intervals for a list of factors. """
complexes = []

for factor, k in ordered(factors):
try:
if not use_cache:
raise KeyError
c = _complexes_cache[factor]
complexes.extend([(i, factor, k) for i in c])
except KeyError:
complex_part = cls._get_complexes_sqf(factor, use_cache)
new = [(root, factor, k) for root in complex_part]
complexes.extend(new)

complexes = cls._complexes_sorted(complexes)
return complexes

@classmethod
def _reals_sorted(cls, reals):
"""Make real isolating intervals disjoint and sort roots. """
cache = {}

for i, (u, f, k) in enumerate(reals):
for j, (v, g, m) in enumerate(reals[i + 1:]):
u, v = u.refine_disjoint(v)
reals[i + j + 1] = (v, g, m)

reals[i] = (u, f, k)

reals = sorted(reals, key=lambda r: r[0].a)

for root, factor, _ in reals:
if factor in cache:
cache[factor].append(root)
else:
cache[factor] = [root]

for factor, roots in cache.items():
_reals_cache[factor] = roots

return reals

@classmethod
def _refine_imaginary(cls, complexes):
sifted = sift(complexes, lambda c: c[1])
complexes = []
for f in ordered(sifted):
nimag = _imag_count_of_factor(f)
if nimag == 0:
# refine until xbounds are neg or pos
for u, f, k in sifted[f]:
while u.ax*u.bx <= 0:
u = u._inner_refine()
complexes.append((u, f, k))
else:
# refine until all but nimag xbounds are neg or pos
potential_imag = list(range(len(sifted[f])))
while True:
assert len(potential_imag) > 1
for i in list(potential_imag):
u, f, k = sifted[f][i]
if u.ax*u.bx > 0:
potential_imag.remove(i)
elif u.ax != u.bx:
u = u._inner_refine()
sifted[f][i] = u, f, k
if len(potential_imag) == nimag:
break
complexes.extend(sifted[f])
return complexes

@classmethod
def _refine_complexes(cls, complexes):
"""return complexes such that no bounding rectangles of non-conjugate
roots would intersect. In addition, assure that neither ay nor by is
0 to guarantee that non-real roots are distinct from real roots in
terms of the y-bounds.
"""
# get the intervals pairwise-disjoint.
# If rectangles were drawn around the coordinates of the bounding
# rectangles, no rectangles would intersect after this procedure.
for i, (u, f, k) in enumerate(complexes):
for j, (v, g, m) in enumerate(complexes[i + 1:]):
u, v = u.refine_disjoint(v)
complexes[i + j + 1] = (v, g, m)

complexes[i] = (u, f, k)

# refine until the x-bounds are unambiguously positive or negative
# for non-imaginary roots
complexes = cls._refine_imaginary(complexes)

# make sure that all y bounds are off the real axis
# and on the same side of the axis
for i, (u, f, k) in enumerate(complexes):
while u.ay*u.by <= 0:
u = u.refine()
complexes[i] = u, f, k
return complexes

@classmethod
def _complexes_sorted(cls, complexes):
"""Make complex isolating intervals disjoint and sort roots. """
complexes = cls._refine_complexes(complexes)
# XXX don't sort until you are sure that it is compatible
# with the indexing method but assert that the desired state
# is not broken
C, F = 0, 1  # location of ComplexInterval and factor
fs = set([i[F] for i in complexes])
for i in range(1, len(complexes)):
if complexes[i][F] != complexes[i - 1][F]:
# if this fails the factors of a root were not
# contiguous because a discontinuity should only
# happen once
fs.remove(complexes[i - 1][F])
for i in range(len(complexes)):
# negative im part (conj=True) comes before
# positive im part (conj=False)
assert complexes[i][C].conj is (i % 2 == 0)

# update cache
cache = {}
# -- collate
for root, factor, _ in complexes:
cache.setdefault(factor, []).append(root)
# -- store
for factor, roots in cache.items():
_complexes_cache[factor] = roots

return complexes

@classmethod
def _reals_index(cls, reals, index):
"""
Map initial real root index to an index in a factor where
the root belongs.
"""
i = 0

for j, (_, factor, k) in enumerate(reals):
if index < i + k:
poly, index = factor, 0

for _, factor, _ in reals[:j]:
if factor == poly:
index += 1

return poly, index
else:
i += k

@classmethod
def _complexes_index(cls, complexes, index):
"""
Map initial complex root index to an index in a factor where
the root belongs.
"""
i = 0
for j, (_, factor, k) in enumerate(complexes):
if index < i + k:
poly, index = factor, 0

for _, factor, _ in complexes[:j]:
if factor == poly:
index += 1

index += len(_reals_cache[poly])

return poly, index
else:
i += k

@classmethod
def _count_roots(cls, roots):
"""Count the number of real or complex roots with multiplicities."""
return sum([k for _, _, k in roots])

@classmethod
def _indexed_root(cls, poly, index):
"""Get a root of a composite polynomial by index. """
factors = _pure_factors(poly)

reals = cls._get_reals(factors)
reals_count = cls._count_roots(reals)

if index < reals_count:
return cls._reals_index(reals, index)
else:
complexes = cls._get_complexes(factors)
return cls._complexes_index(complexes, index - reals_count)

@classmethod
def _real_roots(cls, poly):
"""Get real roots of a composite polynomial. """
factors = _pure_factors(poly)

reals = cls._get_reals(factors)
reals_count = cls._count_roots(reals)

roots = []

for index in range(0, reals_count):
roots.append(cls._reals_index(reals, index))

return roots

def _reset(self):
self._all_roots(self.poly, use_cache=False)

@classmethod
def _all_roots(cls, poly, use_cache=True):
"""Get real and complex roots of a composite polynomial. """
factors = _pure_factors(poly)

reals = cls._get_reals(factors, use_cache=use_cache)
reals_count = cls._count_roots(reals)

roots = []

for index in range(0, reals_count):
roots.append(cls._reals_index(reals, index))

complexes = cls._get_complexes(factors, use_cache=use_cache)
complexes_count = cls._count_roots(complexes)

for index in range(0, complexes_count):
roots.append(cls._complexes_index(complexes, index))

return roots

@classmethod
@cacheit
def _roots_trivial(cls, poly, radicals):
"""Compute roots in linear, quadratic and binomial cases. """
if poly.degree() == 1:
return roots_linear(poly)

return None

if poly.degree() == 2:
elif poly.length() == 2 and poly.TC():
return roots_binomial(poly)
else:
return None

@classmethod
def _preprocess_roots(cls, poly):
"""Take heroic measures to make poly compatible with CRootOf."""
dom = poly.get_domain()

if not dom.is_Exact:
poly = poly.to_exact()

coeff, poly = preprocess_roots(poly)
dom = poly.get_domain()

if not dom.is_ZZ:
raise NotImplementedError(
"sorted roots not supported over %s" % dom)

return coeff, poly

@classmethod
def _postprocess_root(cls, root, radicals):
"""Return the root if it is trivial or a CRootOf object. """
poly, index = root
roots = cls._roots_trivial(poly, radicals)

if roots is not None:
return roots[index]
else:
return cls._new(poly, index)

@classmethod
def _get_roots(cls, method, poly, radicals):
"""Return postprocessed roots of specified kind. """
if not poly.is_univariate:
raise PolynomialError("only univariate polynomials are allowed")

coeff, poly = cls._preprocess_roots(poly)
roots = []

for root in getattr(cls, method)(poly):

return roots

@classmethod
def clear_cache(cls):
"""Reset cache for reals and complexes.

The intervals used to approximate a root instance are updated
as needed. When a request is made to see the intervals, the
most current values are shown. clear_cache will reset all
CRootOf instances back to their original state.

========
_reset
"""
global _reals_cache, _complexes_cache
_reals_cache = _pure_key_dict()
_complexes_cache = _pure_key_dict()

def _get_interval(self):
"""Internal function for retrieving isolation interval from cache. """
if self.is_real:
return _reals_cache[self.poly][self.index]
else:
reals_count = len(_reals_cache[self.poly])
return _complexes_cache[self.poly][self.index - reals_count]

def _set_interval(self, interval):
"""Internal function for updating isolation interval in cache. """
if self.is_real:
_reals_cache[self.poly][self.index] = interval
else:
reals_count = len(_reals_cache[self.poly])
_complexes_cache[self.poly][self.index - reals_count] = interval

def _eval_subs(self, old, new):
# don't allow subs to change anything
return self

def _eval_conjugate(self):
if self.is_real:
return self
expr, i = self.args
return self.func(expr, i + (1 if self._get_interval().conj else -1))

def eval_approx(self, n):
"""Evaluate this complex root to the given precision.

This uses secant method and root bounds are used to both
generate an initial guess and to check that the root
returned is valid. If ever the method converges outside the
root bounds, the bounds will be made smaller and updated.
"""
prec = dps_to_prec(n)
with workprec(prec):
g = self.poly.gen
if not g.is_Symbol:
d = Dummy('x')
if self.is_imaginary:
d *= I
func = lambdify(d, self.expr.subs(g, d))
else:
expr = self.expr
if self.is_imaginary:
expr = self.expr.subs(g, I*g)
func = lambdify(g, expr)

interval = self._get_interval()
while True:
if self.is_real:
a = mpf(str(interval.a))
b = mpf(str(interval.b))
if a == b:
root = a
break
x0 = mpf(str(interval.center))
x1 = x0 + mpf(str(interval.dx))/4
elif self.is_imaginary:
a = mpf(str(interval.ay))
b = mpf(str(interval.by))
if a == b:
root = mpc(mpf('0'), a)
break
x0 = mpf(str(interval.center[1]))
x1 = x0 + mpf(str(interval.dy))/4
else:
ax = mpf(str(interval.ax))
bx = mpf(str(interval.bx))
ay = mpf(str(interval.ay))
by = mpf(str(interval.by))
if ax == bx and ay == by:
root = mpc(ax, ay)
break
x0 = mpc(*map(str, interval.center))
x1 = x0 + mpc(*map(str, (interval.dx, interval.dy)))/4
try:
# without a tolerance, this will return when (to within
# the given precision) x_i == x_{i-1}
root = findroot(func, (x0, x1))
# If the (real or complex) root is not in the 'interval',
# then keep refining the interval. This happens if findroot
# accidentally finds a different root outside of this
# interval because our initial estimate 'x0' was not close
# enough. It is also possible that the secant method will
# get trapped by a max/min in the interval; the root
# verification by findroot will raise a ValueError in this
# case and the interval will then be tightened -- and
# eventually the root will be found.
#
# It is also possible that findroot will not have any
# successful iterations to process (in which case it
# will fail to initialize a variable that is tested
# after the iterations and raise an UnboundLocalError).
if self.is_real or self.is_imaginary:
if not bool(root.imag) == self.is_real and (
a <= root <= b):
if self.is_imaginary:
root = mpc(mpf('0'), root.real)
break
elif (ax <= root.real <= bx and ay <= root.imag <= by):
break
except (UnboundLocalError, ValueError):
pass
interval = interval.refine()

# update the interval so we at least (for this precision or
# less) don't have much work to do to recompute the root
self._set_interval(interval)
return (Float._new(root.real._mpf_, prec) +
I*Float._new(root.imag._mpf_, prec))

def _eval_evalf(self, prec, **kwargs):
"""Evaluate this complex root to the given precision."""
# all kwargs are ignored
return self.eval_rational(n=prec_to_dps(prec))._evalf(prec)

def eval_rational(self, dx=None, dy=None, n=15):
"""
Return a Rational approximation of self that has real
and imaginary component approximations that are within dx
and dy of the true values, respectively. Alternatively,
n digits of precision can be specified.

The interval is refined with bisection and is sure to
converge. The root bounds are updated when the refinement
is complete so recalculation at the same or lesser precision
will not have to repeat the refinement and should be much
faster.

The following example first obtains Rational approximation to
1e-8 accuracy for all roots of the 4-th order Legendre
polynomial. Since the roots are all less than 1, this will
ensure the decimal representation of the approximation will be
correct (including rounding) to 6 digits:

>>> from sympy import S, legendre_poly, Symbol
>>> x = Symbol("x")
>>> p = legendre_poly(4, x, polys=True)
>>> r = p.real_roots()[-1]
>>> r.eval_rational(10**-8).n(6)
0.861136

It is not necessary to a two-step calculation, however: the
decimal representation can be computed directly:

>>> r.evalf(17)
0.86113631159405258

"""
dy = dy or dx
if dx:
rtol = None
dx = dx if isinstance(dx, Rational) else Rational(str(dx))
dy = dy if isinstance(dy, Rational) else Rational(str(dy))
else:
# 5 binary (or 2 decimal) digits are needed to ensure that
# a given digit is correctly rounded
# prec_to_dps(dps_to_prec(n) + 5) - n <= 2 (tested for
# n in range(1000000)
rtol = S(10)**-(n + 2)  # +2 for guard digits
interval = self._get_interval()
while True:
if self.is_real:
if rtol:
dx = abs(interval.center*rtol)
interval = interval.refine_size(dx=dx)
c = interval.center
real = Rational(c)
imag = S.Zero
if not rtol or interval.dx < abs(c*rtol):
break
elif self.is_imaginary:
if rtol:
dy = abs(interval.center[1]*rtol)
dx = 1
interval = interval.refine_size(dx=dx, dy=dy)
c = interval.center[1]
imag = Rational(c)
real = S.Zero
if not rtol or interval.dy < abs(c*rtol):
break
else:
if rtol:
dx = abs(interval.center[0]*rtol)
dy = abs(interval.center[1]*rtol)
interval = interval.refine_size(dx, dy)
c = interval.center
real, imag = map(Rational, c)
if not rtol or (
interval.dx < abs(c[0]*rtol) and
interval.dy < abs(c[1]*rtol)):
break

# update the interval so we at least (for this precision or
# less) don't have much work to do to recompute the root
self._set_interval(interval)
return real + I*imag

def _eval_Eq(self, other):
# CRootOf represents a Root, so if other is that root, it should set
# the expression to zero *and* it should be in the interval of the
# CRootOf instance. It must also be a number that agrees with the
# is_real value of the CRootOf instance.
if type(self) == type(other):
return sympify(self == other)
if not (other.is_number and not other.has(AppliedUndef)):
return S.false
if not other.is_finite:
return S.false
z = self.expr.subs(self.expr.free_symbols.pop(), other).is_zero
if z is False:  # all roots will make z True but we don't know
# whether this is the right root if z is True
return S.false
o = other.is_real, other.is_imaginary
s = self.is_real, self.is_imaginary
assert None not in s  # this is part of initial refinement
if o != s and None not in o:
return S.false
re, im = other.as_real_imag()
if self.is_real:
if im:
return S.false
i = self._get_interval()
a, b = [Rational(str(_)) for _ in (i.a, i.b)]
return sympify(a <= other and other <= b)
i = self._get_interval()
r1, r2, i1, i2 = [Rational(str(j)) for j in (
i.ax, i.bx, i.ay, i.by)]
return sympify((
r1 <= re and re <= r2) and (
i1 <= im and im <= i2))

CRootOf = ComplexRootOf

[docs]@public
class RootSum(Expr):
"""Represents a sum of all roots of a univariate polynomial. """

__slots__ = ['poly', 'fun', 'auto']

def __new__(cls, expr, func=None, x=None, auto=True, quadratic=False):
"""Construct a new RootSum instance of roots of a polynomial."""
coeff, poly = cls._transform(expr, x)

if not poly.is_univariate:
raise MultivariatePolynomialError(
"only univariate polynomials are allowed")

if func is None:
func = Lambda(poly.gen, poly.gen)
else:
try:
is_func = func.is_Function
except AttributeError:
is_func = False

if is_func and 1 in func.nargs:
if not isinstance(func, Lambda):
func = Lambda(poly.gen, func(poly.gen))
else:
raise ValueError(
"expected a univariate function, got %s" % func)

var, expr = func.variables[0], func.expr

if coeff is not S.One:
expr = expr.subs(var, coeff*var)

deg = poly.degree()

if not expr.has(var):
return deg*expr

add_const, expr = expr.as_independent(var)
else:

if expr.is_Mul:
mul_const, expr = expr.as_independent(var)
else:
mul_const = S.One

func = Lambda(var, expr)

rational = cls._is_func_rational(poly, func)
factors, terms = _pure_factors(poly), []

for poly, k in factors:
if poly.is_linear:
term = func(roots_linear(poly)[0])
term = sum(map(func, roots_quadratic(poly)))
else:
if not rational or not auto:
term = cls._new(poly, func, auto)
else:
term = cls._rational_case(poly, func)

terms.append(k*term)

@classmethod
def _new(cls, poly, func, auto=True):
"""Construct new raw RootSum instance. """
obj = Expr.__new__(cls)

obj.poly = poly
obj.fun = func
obj.auto = auto

return obj

@classmethod
def new(cls, poly, func, auto=True):
"""Construct new RootSum instance. """
if not func.expr.has(*func.variables):
return func.expr

rational = cls._is_func_rational(poly, func)

if not rational or not auto:
return cls._new(poly, func, auto)
else:
return cls._rational_case(poly, func)

@classmethod
def _transform(cls, expr, x):
"""Transform an expression to a polynomial. """
poly = PurePoly(expr, x, greedy=False)
return preprocess_roots(poly)

@classmethod
def _is_func_rational(cls, poly, func):
"""Check if a lambda is a rational function. """
var, expr = func.variables[0], func.expr
return expr.is_rational_function(var)

@classmethod
def _rational_case(cls, poly, func):
"""Handle the rational function case. """
roots = symbols('r:%d' % poly.degree())
var, expr = func.variables[0], func.expr

f = sum(expr.subs(var, r) for r in roots)
p, q = together(f).as_numer_denom()

domain = QQ[roots]

p = p.expand()
q = q.expand()

try:
p = Poly(p, domain=domain, expand=False)
except GeneratorsNeeded:
p, p_coeff = None, (p,)
else:
p_monom, p_coeff = zip(*p.terms())

try:
q = Poly(q, domain=domain, expand=False)
except GeneratorsNeeded:
q, q_coeff = None, (q,)
else:
q_monom, q_coeff = zip(*q.terms())

coeffs, mapping = symmetrize(p_coeff + q_coeff, formal=True)
formulas, values = viete(poly, roots), []

for (sym, _), (_, val) in zip(mapping, formulas):
values.append((sym, val))

for i, (coeff, _) in enumerate(coeffs):
coeffs[i] = coeff.subs(values)

n = len(p_coeff)

p_coeff = coeffs[:n]
q_coeff = coeffs[n:]

if p is not None:
p = Poly(dict(zip(p_monom, p_coeff)), *p.gens).as_expr()
else:
(p,) = p_coeff

if q is not None:
q = Poly(dict(zip(q_monom, q_coeff)), *q.gens).as_expr()
else:
(q,) = q_coeff

return factor(p/q)

def _hashable_content(self):
return (self.poly, self.fun)

@property
def expr(self):
return self.poly.as_expr()

@property
def args(self):
return (self.expr, self.fun, self.poly.gen)

@property
def free_symbols(self):
return self.poly.free_symbols | self.fun.free_symbols

@property
def is_commutative(self):
return True

def doit(self, **hints):
if not hints.get('roots', True):
return self

_roots = roots(self.poly, multiple=True)

if len(_roots) < self.poly.degree():
return self
else:
return Add(*[self.fun(r) for r in _roots])

def _eval_evalf(self, prec):
try:
_roots = self.poly.nroots(n=prec_to_dps(prec))
except (DomainError, PolynomialError):
return self
else:
return Add(*[self.fun(r) for r in _roots])

def _eval_derivative(self, x):
var, expr = self.fun.args
func = Lambda(var, expr.diff(x))
return self.new(self.poly, func, self.auto)