Coverage for /usr/lib/python3/dist-packages/sympy/polys/rootoftools.py: 21%
636 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1"""Implementation of RootOf class and related tools. """
4from sympy.core.basic import Basic
5from sympy.core import (S, Expr, Integer, Float, I, oo, Add, Lambda,
6 symbols, sympify, Rational, Dummy)
7from sympy.core.cache import cacheit
8from sympy.core.relational import is_le
9from sympy.core.sorting import ordered
10from sympy.polys.domains import QQ
11from sympy.polys.polyerrors import (
12 MultivariatePolynomialError,
13 GeneratorsNeeded,
14 PolynomialError,
15 DomainError)
16from sympy.polys.polyfuncs import symmetrize, viete
17from sympy.polys.polyroots import (
18 roots_linear, roots_quadratic, roots_binomial,
19 preprocess_roots, roots)
20from sympy.polys.polytools import Poly, PurePoly, factor
21from sympy.polys.rationaltools import together
22from sympy.polys.rootisolation import (
23 dup_isolate_complex_roots_sqf,
24 dup_isolate_real_roots_sqf)
25from sympy.utilities import lambdify, public, sift, numbered_symbols
27from mpmath import mpf, mpc, findroot, workprec
28from mpmath.libmp.libmpf import dps_to_prec, prec_to_dps
29from sympy.multipledispatch import dispatch
30from itertools import chain
33__all__ = ['CRootOf']
37class _pure_key_dict:
38 """A minimal dictionary that makes sure that the key is a
39 univariate PurePoly instance.
41 Examples
42 ========
44 Only the following actions are guaranteed:
46 >>> from sympy.polys.rootoftools import _pure_key_dict
47 >>> from sympy import PurePoly
48 >>> from sympy.abc import x, y
50 1) creation
52 >>> P = _pure_key_dict()
54 2) assignment for a PurePoly or univariate polynomial
56 >>> P[x] = 1
57 >>> P[PurePoly(x - y, x)] = 2
59 3) retrieval based on PurePoly key comparison (use this
60 instead of the get method)
62 >>> P[y]
63 1
65 4) KeyError when trying to retrieve a nonexisting key
67 >>> P[y + 1]
68 Traceback (most recent call last):
69 ...
70 KeyError: PurePoly(y + 1, y, domain='ZZ')
72 5) ability to query with ``in``
74 >>> x + 1 in P
75 False
77 NOTE: this is a *not* a dictionary. It is a very basic object
78 for internal use that makes sure to always address its cache
79 via PurePoly instances. It does not, for example, implement
80 ``get`` or ``setdefault``.
81 """
82 def __init__(self):
83 self._dict = {}
85 def __getitem__(self, k):
86 if not isinstance(k, PurePoly):
87 if not (isinstance(k, Expr) and len(k.free_symbols) == 1):
88 raise KeyError
89 k = PurePoly(k, expand=False)
90 return self._dict[k]
92 def __setitem__(self, k, v):
93 if not isinstance(k, PurePoly):
94 if not (isinstance(k, Expr) and len(k.free_symbols) == 1):
95 raise ValueError('expecting univariate expression')
96 k = PurePoly(k, expand=False)
97 self._dict[k] = v
99 def __contains__(self, k):
100 try:
101 self[k]
102 return True
103 except KeyError:
104 return False
106_reals_cache = _pure_key_dict()
107_complexes_cache = _pure_key_dict()
110def _pure_factors(poly):
111 _, factors = poly.factor_list()
112 return [(PurePoly(f, expand=False), m) for f, m in factors]
115def _imag_count_of_factor(f):
116 """Return the number of imaginary roots for irreducible
117 univariate polynomial ``f``.
118 """
119 terms = [(i, j) for (i,), j in f.terms()]
120 if any(i % 2 for i, j in terms):
121 return 0
122 # update signs
123 even = [(i, I**i*j) for i, j in terms]
124 even = Poly.from_dict(dict(even), Dummy('x'))
125 return int(even.count_roots(-oo, oo))
128@public
129def rootof(f, x, index=None, radicals=True, expand=True):
130 """An indexed root of a univariate polynomial.
132 Returns either a :obj:`ComplexRootOf` object or an explicit
133 expression involving radicals.
135 Parameters
136 ==========
138 f : Expr
139 Univariate polynomial.
140 x : Symbol, optional
141 Generator for ``f``.
142 index : int or Integer
143 radicals : bool
144 Return a radical expression if possible.
145 expand : bool
146 Expand ``f``.
147 """
148 return CRootOf(f, x, index=index, radicals=radicals, expand=expand)
151@public
152class RootOf(Expr):
153 """Represents a root of a univariate polynomial.
155 Base class for roots of different kinds of polynomials.
156 Only complex roots are currently supported.
157 """
159 __slots__ = ('poly',)
161 def __new__(cls, f, x, index=None, radicals=True, expand=True):
162 """Construct a new ``CRootOf`` object for ``k``-th root of ``f``."""
163 return rootof(f, x, index=index, radicals=radicals, expand=expand)
165@public
166class ComplexRootOf(RootOf):
167 """Represents an indexed complex root of a polynomial.
169 Roots of a univariate polynomial separated into disjoint
170 real or complex intervals and indexed in a fixed order:
172 * real roots come first and are sorted in increasing order;
173 * complex roots come next and are sorted primarily by increasing
174 real part, secondarily by increasing imaginary part.
176 Currently only rational coefficients are allowed.
177 Can be imported as ``CRootOf``. To avoid confusion, the
178 generator must be a Symbol.
181 Examples
182 ========
184 >>> from sympy import CRootOf, rootof
185 >>> from sympy.abc import x
187 CRootOf is a way to reference a particular root of a
188 polynomial. If there is a rational root, it will be returned:
190 >>> CRootOf.clear_cache() # for doctest reproducibility
191 >>> CRootOf(x**2 - 4, 0)
192 -2
194 Whether roots involving radicals are returned or not
195 depends on whether the ``radicals`` flag is true (which is
196 set to True with rootof):
198 >>> CRootOf(x**2 - 3, 0)
199 CRootOf(x**2 - 3, 0)
200 >>> CRootOf(x**2 - 3, 0, radicals=True)
201 -sqrt(3)
202 >>> rootof(x**2 - 3, 0)
203 -sqrt(3)
205 The following cannot be expressed in terms of radicals:
207 >>> r = rootof(4*x**5 + 16*x**3 + 12*x**2 + 7, 0); r
208 CRootOf(4*x**5 + 16*x**3 + 12*x**2 + 7, 0)
210 The root bounds can be seen, however, and they are used by the
211 evaluation methods to get numerical approximations for the root.
213 >>> interval = r._get_interval(); interval
214 (-1, 0)
215 >>> r.evalf(2)
216 -0.98
218 The evalf method refines the width of the root bounds until it
219 guarantees that any decimal approximation within those bounds
220 will satisfy the desired precision. It then stores the refined
221 interval so subsequent requests at or below the requested
222 precision will not have to recompute the root bounds and will
223 return very quickly.
225 Before evaluation above, the interval was
227 >>> interval
228 (-1, 0)
230 After evaluation it is now
232 >>> r._get_interval() # doctest: +SKIP
233 (-165/169, -206/211)
235 To reset all intervals for a given polynomial, the :meth:`_reset` method
236 can be called from any CRootOf instance of the polynomial:
238 >>> r._reset()
239 >>> r._get_interval()
240 (-1, 0)
242 The :meth:`eval_approx` method will also find the root to a given
243 precision but the interval is not modified unless the search
244 for the root fails to converge within the root bounds. And
245 the secant method is used to find the root. (The ``evalf``
246 method uses bisection and will always update the interval.)
248 >>> r.eval_approx(2)
249 -0.98
251 The interval needed to be slightly updated to find that root:
253 >>> r._get_interval()
254 (-1, -1/2)
256 The ``evalf_rational`` will compute a rational approximation
257 of the root to the desired accuracy or precision.
259 >>> r.eval_rational(n=2)
260 -69629/71318
262 >>> t = CRootOf(x**3 + 10*x + 1, 1)
263 >>> t.eval_rational(1e-1)
264 15/256 - 805*I/256
265 >>> t.eval_rational(1e-1, 1e-4)
266 3275/65536 - 414645*I/131072
267 >>> t.eval_rational(1e-4, 1e-4)
268 6545/131072 - 414645*I/131072
269 >>> t.eval_rational(n=2)
270 104755/2097152 - 6634255*I/2097152
272 Notes
273 =====
275 Although a PurePoly can be constructed from a non-symbol generator
276 RootOf instances of non-symbols are disallowed to avoid confusion
277 over what root is being represented.
279 >>> from sympy import exp, PurePoly
280 >>> PurePoly(x) == PurePoly(exp(x))
281 True
282 >>> CRootOf(x - 1, 0)
283 1
284 >>> CRootOf(exp(x) - 1, 0) # would correspond to x == 0
285 Traceback (most recent call last):
286 ...
287 sympy.polys.polyerrors.PolynomialError: generator must be a Symbol
289 See Also
290 ========
292 eval_approx
293 eval_rational
295 """
297 __slots__ = ('index',)
298 is_complex = True
299 is_number = True
300 is_finite = True
302 def __new__(cls, f, x, index=None, radicals=False, expand=True):
303 """ Construct an indexed complex root of a polynomial.
305 See ``rootof`` for the parameters.
307 The default value of ``radicals`` is ``False`` to satisfy
308 ``eval(srepr(expr) == expr``.
309 """
310 x = sympify(x)
312 if index is None and x.is_Integer:
313 x, index = None, x
314 else:
315 index = sympify(index)
317 if index is not None and index.is_Integer:
318 index = int(index)
319 else:
320 raise ValueError("expected an integer root index, got %s" % index)
322 poly = PurePoly(f, x, greedy=False, expand=expand)
324 if not poly.is_univariate:
325 raise PolynomialError("only univariate polynomials are allowed")
327 if not poly.gen.is_Symbol:
328 # PurePoly(sin(x) + 1) == PurePoly(x + 1) but the roots of
329 # x for each are not the same: issue 8617
330 raise PolynomialError("generator must be a Symbol")
332 degree = poly.degree()
334 if degree <= 0:
335 raise PolynomialError("Cannot construct CRootOf object for %s" % f)
337 if index < -degree or index >= degree:
338 raise IndexError("root index out of [%d, %d] range, got %d" %
339 (-degree, degree - 1, index))
340 elif index < 0:
341 index += degree
343 dom = poly.get_domain()
345 if not dom.is_Exact:
346 poly = poly.to_exact()
348 roots = cls._roots_trivial(poly, radicals)
350 if roots is not None:
351 return roots[index]
353 coeff, poly = preprocess_roots(poly)
354 dom = poly.get_domain()
356 if not dom.is_ZZ:
357 raise NotImplementedError("CRootOf is not supported over %s" % dom)
359 root = cls._indexed_root(poly, index, lazy=True)
360 return coeff * cls._postprocess_root(root, radicals)
362 @classmethod
363 def _new(cls, poly, index):
364 """Construct new ``CRootOf`` object from raw data. """
365 obj = Expr.__new__(cls)
367 obj.poly = PurePoly(poly)
368 obj.index = index
370 try:
371 _reals_cache[obj.poly] = _reals_cache[poly]
372 _complexes_cache[obj.poly] = _complexes_cache[poly]
373 except KeyError:
374 pass
376 return obj
378 def _hashable_content(self):
379 return (self.poly, self.index)
381 @property
382 def expr(self):
383 return self.poly.as_expr()
385 @property
386 def args(self):
387 return (self.expr, Integer(self.index))
389 @property
390 def free_symbols(self):
391 # CRootOf currently only works with univariate expressions
392 # whose poly attribute should be a PurePoly with no free
393 # symbols
394 return set()
396 def _eval_is_real(self):
397 """Return ``True`` if the root is real. """
398 self._ensure_reals_init()
399 return self.index < len(_reals_cache[self.poly])
401 def _eval_is_imaginary(self):
402 """Return ``True`` if the root is imaginary. """
403 self._ensure_reals_init()
404 if self.index >= len(_reals_cache[self.poly]):
405 ivl = self._get_interval()
406 return ivl.ax*ivl.bx <= 0 # all others are on one side or the other
407 return False # XXX is this necessary?
409 @classmethod
410 def real_roots(cls, poly, radicals=True):
411 """Get real roots of a polynomial. """
412 return cls._get_roots("_real_roots", poly, radicals)
414 @classmethod
415 def all_roots(cls, poly, radicals=True):
416 """Get real and complex roots of a polynomial. """
417 return cls._get_roots("_all_roots", poly, radicals)
419 @classmethod
420 def _get_reals_sqf(cls, currentfactor, use_cache=True):
421 """Get real root isolating intervals for a square-free factor."""
422 if use_cache and currentfactor in _reals_cache:
423 real_part = _reals_cache[currentfactor]
424 else:
425 _reals_cache[currentfactor] = real_part = \
426 dup_isolate_real_roots_sqf(
427 currentfactor.rep.rep, currentfactor.rep.dom, blackbox=True)
429 return real_part
431 @classmethod
432 def _get_complexes_sqf(cls, currentfactor, use_cache=True):
433 """Get complex root isolating intervals for a square-free factor."""
434 if use_cache and currentfactor in _complexes_cache:
435 complex_part = _complexes_cache[currentfactor]
436 else:
437 _complexes_cache[currentfactor] = complex_part = \
438 dup_isolate_complex_roots_sqf(
439 currentfactor.rep.rep, currentfactor.rep.dom, blackbox=True)
440 return complex_part
442 @classmethod
443 def _get_reals(cls, factors, use_cache=True):
444 """Compute real root isolating intervals for a list of factors. """
445 reals = []
447 for currentfactor, k in factors:
448 try:
449 if not use_cache:
450 raise KeyError
451 r = _reals_cache[currentfactor]
452 reals.extend([(i, currentfactor, k) for i in r])
453 except KeyError:
454 real_part = cls._get_reals_sqf(currentfactor, use_cache)
455 new = [(root, currentfactor, k) for root in real_part]
456 reals.extend(new)
458 reals = cls._reals_sorted(reals)
459 return reals
461 @classmethod
462 def _get_complexes(cls, factors, use_cache=True):
463 """Compute complex root isolating intervals for a list of factors. """
464 complexes = []
466 for currentfactor, k in ordered(factors):
467 try:
468 if not use_cache:
469 raise KeyError
470 c = _complexes_cache[currentfactor]
471 complexes.extend([(i, currentfactor, k) for i in c])
472 except KeyError:
473 complex_part = cls._get_complexes_sqf(currentfactor, use_cache)
474 new = [(root, currentfactor, k) for root in complex_part]
475 complexes.extend(new)
477 complexes = cls._complexes_sorted(complexes)
478 return complexes
480 @classmethod
481 def _reals_sorted(cls, reals):
482 """Make real isolating intervals disjoint and sort roots. """
483 cache = {}
485 for i, (u, f, k) in enumerate(reals):
486 for j, (v, g, m) in enumerate(reals[i + 1:]):
487 u, v = u.refine_disjoint(v)
488 reals[i + j + 1] = (v, g, m)
490 reals[i] = (u, f, k)
492 reals = sorted(reals, key=lambda r: r[0].a)
494 for root, currentfactor, _ in reals:
495 if currentfactor in cache:
496 cache[currentfactor].append(root)
497 else:
498 cache[currentfactor] = [root]
500 for currentfactor, root in cache.items():
501 _reals_cache[currentfactor] = root
503 return reals
505 @classmethod
506 def _refine_imaginary(cls, complexes):
507 sifted = sift(complexes, lambda c: c[1])
508 complexes = []
509 for f in ordered(sifted):
510 nimag = _imag_count_of_factor(f)
511 if nimag == 0:
512 # refine until xbounds are neg or pos
513 for u, f, k in sifted[f]:
514 while u.ax*u.bx <= 0:
515 u = u._inner_refine()
516 complexes.append((u, f, k))
517 else:
518 # refine until all but nimag xbounds are neg or pos
519 potential_imag = list(range(len(sifted[f])))
520 while True:
521 assert len(potential_imag) > 1
522 for i in list(potential_imag):
523 u, f, k = sifted[f][i]
524 if u.ax*u.bx > 0:
525 potential_imag.remove(i)
526 elif u.ax != u.bx:
527 u = u._inner_refine()
528 sifted[f][i] = u, f, k
529 if len(potential_imag) == nimag:
530 break
531 complexes.extend(sifted[f])
532 return complexes
534 @classmethod
535 def _refine_complexes(cls, complexes):
536 """return complexes such that no bounding rectangles of non-conjugate
537 roots would intersect. In addition, assure that neither ay nor by is
538 0 to guarantee that non-real roots are distinct from real roots in
539 terms of the y-bounds.
540 """
541 # get the intervals pairwise-disjoint.
542 # If rectangles were drawn around the coordinates of the bounding
543 # rectangles, no rectangles would intersect after this procedure.
544 for i, (u, f, k) in enumerate(complexes):
545 for j, (v, g, m) in enumerate(complexes[i + 1:]):
546 u, v = u.refine_disjoint(v)
547 complexes[i + j + 1] = (v, g, m)
549 complexes[i] = (u, f, k)
551 # refine until the x-bounds are unambiguously positive or negative
552 # for non-imaginary roots
553 complexes = cls._refine_imaginary(complexes)
555 # make sure that all y bounds are off the real axis
556 # and on the same side of the axis
557 for i, (u, f, k) in enumerate(complexes):
558 while u.ay*u.by <= 0:
559 u = u.refine()
560 complexes[i] = u, f, k
561 return complexes
563 @classmethod
564 def _complexes_sorted(cls, complexes):
565 """Make complex isolating intervals disjoint and sort roots. """
566 complexes = cls._refine_complexes(complexes)
567 # XXX don't sort until you are sure that it is compatible
568 # with the indexing method but assert that the desired state
569 # is not broken
570 C, F = 0, 1 # location of ComplexInterval and factor
571 fs = {i[F] for i in complexes}
572 for i in range(1, len(complexes)):
573 if complexes[i][F] != complexes[i - 1][F]:
574 # if this fails the factors of a root were not
575 # contiguous because a discontinuity should only
576 # happen once
577 fs.remove(complexes[i - 1][F])
578 for i, cmplx in enumerate(complexes):
579 # negative im part (conj=True) comes before
580 # positive im part (conj=False)
581 assert cmplx[C].conj is (i % 2 == 0)
583 # update cache
584 cache = {}
585 # -- collate
586 for root, currentfactor, _ in complexes:
587 cache.setdefault(currentfactor, []).append(root)
588 # -- store
589 for currentfactor, root in cache.items():
590 _complexes_cache[currentfactor] = root
592 return complexes
594 @classmethod
595 def _reals_index(cls, reals, index):
596 """
597 Map initial real root index to an index in a factor where
598 the root belongs.
599 """
600 i = 0
602 for j, (_, currentfactor, k) in enumerate(reals):
603 if index < i + k:
604 poly, index = currentfactor, 0
606 for _, currentfactor, _ in reals[:j]:
607 if currentfactor == poly:
608 index += 1
610 return poly, index
611 else:
612 i += k
614 @classmethod
615 def _complexes_index(cls, complexes, index):
616 """
617 Map initial complex root index to an index in a factor where
618 the root belongs.
619 """
620 i = 0
621 for j, (_, currentfactor, k) in enumerate(complexes):
622 if index < i + k:
623 poly, index = currentfactor, 0
625 for _, currentfactor, _ in complexes[:j]:
626 if currentfactor == poly:
627 index += 1
629 index += len(_reals_cache[poly])
631 return poly, index
632 else:
633 i += k
635 @classmethod
636 def _count_roots(cls, roots):
637 """Count the number of real or complex roots with multiplicities."""
638 return sum([k for _, _, k in roots])
640 @classmethod
641 def _indexed_root(cls, poly, index, lazy=False):
642 """Get a root of a composite polynomial by index. """
643 factors = _pure_factors(poly)
645 # If the given poly is already irreducible, then the index does not
646 # need to be adjusted, and we can postpone the heavy lifting of
647 # computing and refining isolating intervals until that is needed.
648 if lazy and len(factors) == 1 and factors[0][1] == 1:
649 return poly, index
651 reals = cls._get_reals(factors)
652 reals_count = cls._count_roots(reals)
654 if index < reals_count:
655 return cls._reals_index(reals, index)
656 else:
657 complexes = cls._get_complexes(factors)
658 return cls._complexes_index(complexes, index - reals_count)
660 def _ensure_reals_init(self):
661 """Ensure that our poly has entries in the reals cache. """
662 if self.poly not in _reals_cache:
663 self._indexed_root(self.poly, self.index)
665 def _ensure_complexes_init(self):
666 """Ensure that our poly has entries in the complexes cache. """
667 if self.poly not in _complexes_cache:
668 self._indexed_root(self.poly, self.index)
670 @classmethod
671 def _real_roots(cls, poly):
672 """Get real roots of a composite polynomial. """
673 factors = _pure_factors(poly)
675 reals = cls._get_reals(factors)
676 reals_count = cls._count_roots(reals)
678 roots = []
680 for index in range(0, reals_count):
681 roots.append(cls._reals_index(reals, index))
683 return roots
685 def _reset(self):
686 """
687 Reset all intervals
688 """
689 self._all_roots(self.poly, use_cache=False)
691 @classmethod
692 def _all_roots(cls, poly, use_cache=True):
693 """Get real and complex roots of a composite polynomial. """
694 factors = _pure_factors(poly)
696 reals = cls._get_reals(factors, use_cache=use_cache)
697 reals_count = cls._count_roots(reals)
699 roots = []
701 for index in range(0, reals_count):
702 roots.append(cls._reals_index(reals, index))
704 complexes = cls._get_complexes(factors, use_cache=use_cache)
705 complexes_count = cls._count_roots(complexes)
707 for index in range(0, complexes_count):
708 roots.append(cls._complexes_index(complexes, index))
710 return roots
712 @classmethod
713 @cacheit
714 def _roots_trivial(cls, poly, radicals):
715 """Compute roots in linear, quadratic and binomial cases. """
716 if poly.degree() == 1:
717 return roots_linear(poly)
719 if not radicals:
720 return None
722 if poly.degree() == 2:
723 return roots_quadratic(poly)
724 elif poly.length() == 2 and poly.TC():
725 return roots_binomial(poly)
726 else:
727 return None
729 @classmethod
730 def _preprocess_roots(cls, poly):
731 """Take heroic measures to make ``poly`` compatible with ``CRootOf``."""
732 dom = poly.get_domain()
734 if not dom.is_Exact:
735 poly = poly.to_exact()
737 coeff, poly = preprocess_roots(poly)
738 dom = poly.get_domain()
740 if not dom.is_ZZ:
741 raise NotImplementedError(
742 "sorted roots not supported over %s" % dom)
744 return coeff, poly
746 @classmethod
747 def _postprocess_root(cls, root, radicals):
748 """Return the root if it is trivial or a ``CRootOf`` object. """
749 poly, index = root
750 roots = cls._roots_trivial(poly, radicals)
752 if roots is not None:
753 return roots[index]
754 else:
755 return cls._new(poly, index)
757 @classmethod
758 def _get_roots(cls, method, poly, radicals):
759 """Return postprocessed roots of specified kind. """
760 if not poly.is_univariate:
761 raise PolynomialError("only univariate polynomials are allowed")
762 # get rid of gen and it's free symbol
763 d = Dummy()
764 poly = poly.subs(poly.gen, d)
765 x = symbols('x')
766 # see what others are left and select x or a numbered x
767 # that doesn't clash
768 free_names = {str(i) for i in poly.free_symbols}
769 for x in chain((symbols('x'),), numbered_symbols('x')):
770 if x.name not in free_names:
771 poly = poly.xreplace({d: x})
772 break
773 coeff, poly = cls._preprocess_roots(poly)
774 roots = []
776 for root in getattr(cls, method)(poly):
777 roots.append(coeff*cls._postprocess_root(root, radicals))
778 return roots
780 @classmethod
781 def clear_cache(cls):
782 """Reset cache for reals and complexes.
784 The intervals used to approximate a root instance are updated
785 as needed. When a request is made to see the intervals, the
786 most current values are shown. `clear_cache` will reset all
787 CRootOf instances back to their original state.
789 See Also
790 ========
792 _reset
793 """
794 global _reals_cache, _complexes_cache
795 _reals_cache = _pure_key_dict()
796 _complexes_cache = _pure_key_dict()
798 def _get_interval(self):
799 """Internal function for retrieving isolation interval from cache. """
800 self._ensure_reals_init()
801 if self.is_real:
802 return _reals_cache[self.poly][self.index]
803 else:
804 reals_count = len(_reals_cache[self.poly])
805 self._ensure_complexes_init()
806 return _complexes_cache[self.poly][self.index - reals_count]
808 def _set_interval(self, interval):
809 """Internal function for updating isolation interval in cache. """
810 self._ensure_reals_init()
811 if self.is_real:
812 _reals_cache[self.poly][self.index] = interval
813 else:
814 reals_count = len(_reals_cache[self.poly])
815 self._ensure_complexes_init()
816 _complexes_cache[self.poly][self.index - reals_count] = interval
818 def _eval_subs(self, old, new):
819 # don't allow subs to change anything
820 return self
822 def _eval_conjugate(self):
823 if self.is_real:
824 return self
825 expr, i = self.args
826 return self.func(expr, i + (1 if self._get_interval().conj else -1))
828 def eval_approx(self, n, return_mpmath=False):
829 """Evaluate this complex root to the given precision.
831 This uses secant method and root bounds are used to both
832 generate an initial guess and to check that the root
833 returned is valid. If ever the method converges outside the
834 root bounds, the bounds will be made smaller and updated.
835 """
836 prec = dps_to_prec(n)
837 with workprec(prec):
838 g = self.poly.gen
839 if not g.is_Symbol:
840 d = Dummy('x')
841 if self.is_imaginary:
842 d *= I
843 func = lambdify(d, self.expr.subs(g, d))
844 else:
845 expr = self.expr
846 if self.is_imaginary:
847 expr = self.expr.subs(g, I*g)
848 func = lambdify(g, expr)
850 interval = self._get_interval()
851 while True:
852 if self.is_real:
853 a = mpf(str(interval.a))
854 b = mpf(str(interval.b))
855 if a == b:
856 root = a
857 break
858 x0 = mpf(str(interval.center))
859 x1 = x0 + mpf(str(interval.dx))/4
860 elif self.is_imaginary:
861 a = mpf(str(interval.ay))
862 b = mpf(str(interval.by))
863 if a == b:
864 root = mpc(mpf('0'), a)
865 break
866 x0 = mpf(str(interval.center[1]))
867 x1 = x0 + mpf(str(interval.dy))/4
868 else:
869 ax = mpf(str(interval.ax))
870 bx = mpf(str(interval.bx))
871 ay = mpf(str(interval.ay))
872 by = mpf(str(interval.by))
873 if ax == bx and ay == by:
874 root = mpc(ax, ay)
875 break
876 x0 = mpc(*map(str, interval.center))
877 x1 = x0 + mpc(*map(str, (interval.dx, interval.dy)))/4
878 try:
879 # without a tolerance, this will return when (to within
880 # the given precision) x_i == x_{i-1}
881 root = findroot(func, (x0, x1))
882 # If the (real or complex) root is not in the 'interval',
883 # then keep refining the interval. This happens if findroot
884 # accidentally finds a different root outside of this
885 # interval because our initial estimate 'x0' was not close
886 # enough. It is also possible that the secant method will
887 # get trapped by a max/min in the interval; the root
888 # verification by findroot will raise a ValueError in this
889 # case and the interval will then be tightened -- and
890 # eventually the root will be found.
891 #
892 # It is also possible that findroot will not have any
893 # successful iterations to process (in which case it
894 # will fail to initialize a variable that is tested
895 # after the iterations and raise an UnboundLocalError).
896 if self.is_real or self.is_imaginary:
897 if not bool(root.imag) == self.is_real and (
898 a <= root <= b):
899 if self.is_imaginary:
900 root = mpc(mpf('0'), root.real)
901 break
902 elif (ax <= root.real <= bx and ay <= root.imag <= by):
903 break
904 except (UnboundLocalError, ValueError):
905 pass
906 interval = interval.refine()
908 # update the interval so we at least (for this precision or
909 # less) don't have much work to do to recompute the root
910 self._set_interval(interval)
911 if return_mpmath:
912 return root
913 return (Float._new(root.real._mpf_, prec) +
914 I*Float._new(root.imag._mpf_, prec))
916 def _eval_evalf(self, prec, **kwargs):
917 """Evaluate this complex root to the given precision."""
918 # all kwargs are ignored
919 return self.eval_rational(n=prec_to_dps(prec))._evalf(prec)
921 def eval_rational(self, dx=None, dy=None, n=15):
922 """
923 Return a Rational approximation of ``self`` that has real
924 and imaginary component approximations that are within ``dx``
925 and ``dy`` of the true values, respectively. Alternatively,
926 ``n`` digits of precision can be specified.
928 The interval is refined with bisection and is sure to
929 converge. The root bounds are updated when the refinement
930 is complete so recalculation at the same or lesser precision
931 will not have to repeat the refinement and should be much
932 faster.
934 The following example first obtains Rational approximation to
935 1e-8 accuracy for all roots of the 4-th order Legendre
936 polynomial. Since the roots are all less than 1, this will
937 ensure the decimal representation of the approximation will be
938 correct (including rounding) to 6 digits:
940 >>> from sympy import legendre_poly, Symbol
941 >>> x = Symbol("x")
942 >>> p = legendre_poly(4, x, polys=True)
943 >>> r = p.real_roots()[-1]
944 >>> r.eval_rational(10**-8).n(6)
945 0.861136
947 It is not necessary to a two-step calculation, however: the
948 decimal representation can be computed directly:
950 >>> r.evalf(17)
951 0.86113631159405258
953 """
954 dy = dy or dx
955 if dx:
956 rtol = None
957 dx = dx if isinstance(dx, Rational) else Rational(str(dx))
958 dy = dy if isinstance(dy, Rational) else Rational(str(dy))
959 else:
960 # 5 binary (or 2 decimal) digits are needed to ensure that
961 # a given digit is correctly rounded
962 # prec_to_dps(dps_to_prec(n) + 5) - n <= 2 (tested for
963 # n in range(1000000)
964 rtol = S(10)**-(n + 2) # +2 for guard digits
965 interval = self._get_interval()
966 while True:
967 if self.is_real:
968 if rtol:
969 dx = abs(interval.center*rtol)
970 interval = interval.refine_size(dx=dx)
971 c = interval.center
972 real = Rational(c)
973 imag = S.Zero
974 if not rtol or interval.dx < abs(c*rtol):
975 break
976 elif self.is_imaginary:
977 if rtol:
978 dy = abs(interval.center[1]*rtol)
979 dx = 1
980 interval = interval.refine_size(dx=dx, dy=dy)
981 c = interval.center[1]
982 imag = Rational(c)
983 real = S.Zero
984 if not rtol or interval.dy < abs(c*rtol):
985 break
986 else:
987 if rtol:
988 dx = abs(interval.center[0]*rtol)
989 dy = abs(interval.center[1]*rtol)
990 interval = interval.refine_size(dx, dy)
991 c = interval.center
992 real, imag = map(Rational, c)
993 if not rtol or (
994 interval.dx < abs(c[0]*rtol) and
995 interval.dy < abs(c[1]*rtol)):
996 break
998 # update the interval so we at least (for this precision or
999 # less) don't have much work to do to recompute the root
1000 self._set_interval(interval)
1001 return real + I*imag
1004CRootOf = ComplexRootOf
1007@dispatch(ComplexRootOf, ComplexRootOf)
1008def _eval_is_eq(lhs, rhs): # noqa:F811
1009 # if we use is_eq to check here, we get infinite recurion
1010 return lhs == rhs
1013@dispatch(ComplexRootOf, Basic) # type:ignore
1014def _eval_is_eq(lhs, rhs): # noqa:F811
1015 # CRootOf represents a Root, so if rhs is that root, it should set
1016 # the expression to zero *and* it should be in the interval of the
1017 # CRootOf instance. It must also be a number that agrees with the
1018 # is_real value of the CRootOf instance.
1019 if not rhs.is_number:
1020 return None
1021 if not rhs.is_finite:
1022 return False
1023 z = lhs.expr.subs(lhs.expr.free_symbols.pop(), rhs).is_zero
1024 if z is False: # all roots will make z True but we don't know
1025 # whether this is the right root if z is True
1026 return False
1027 o = rhs.is_real, rhs.is_imaginary
1028 s = lhs.is_real, lhs.is_imaginary
1029 assert None not in s # this is part of initial refinement
1030 if o != s and None not in o:
1031 return False
1032 re, im = rhs.as_real_imag()
1033 if lhs.is_real:
1034 if im:
1035 return False
1036 i = lhs._get_interval()
1037 a, b = [Rational(str(_)) for _ in (i.a, i.b)]
1038 return sympify(a <= rhs and rhs <= b)
1039 i = lhs._get_interval()
1040 r1, r2, i1, i2 = [Rational(str(j)) for j in (
1041 i.ax, i.bx, i.ay, i.by)]
1042 return is_le(r1, re) and is_le(re,r2) and is_le(i1,im) and is_le(im,i2)
1045@public
1046class RootSum(Expr):
1047 """Represents a sum of all roots of a univariate polynomial. """
1049 __slots__ = ('poly', 'fun', 'auto')
1051 def __new__(cls, expr, func=None, x=None, auto=True, quadratic=False):
1052 """Construct a new ``RootSum`` instance of roots of a polynomial."""
1053 coeff, poly = cls._transform(expr, x)
1055 if not poly.is_univariate:
1056 raise MultivariatePolynomialError(
1057 "only univariate polynomials are allowed")
1059 if func is None:
1060 func = Lambda(poly.gen, poly.gen)
1061 else:
1062 is_func = getattr(func, 'is_Function', False)
1064 if is_func and 1 in func.nargs:
1065 if not isinstance(func, Lambda):
1066 func = Lambda(poly.gen, func(poly.gen))
1067 else:
1068 raise ValueError(
1069 "expected a univariate function, got %s" % func)
1071 var, expr = func.variables[0], func.expr
1073 if coeff is not S.One:
1074 expr = expr.subs(var, coeff*var)
1076 deg = poly.degree()
1078 if not expr.has(var):
1079 return deg*expr
1081 if expr.is_Add:
1082 add_const, expr = expr.as_independent(var)
1083 else:
1084 add_const = S.Zero
1086 if expr.is_Mul:
1087 mul_const, expr = expr.as_independent(var)
1088 else:
1089 mul_const = S.One
1091 func = Lambda(var, expr)
1093 rational = cls._is_func_rational(poly, func)
1094 factors, terms = _pure_factors(poly), []
1096 for poly, k in factors:
1097 if poly.is_linear:
1098 term = func(roots_linear(poly)[0])
1099 elif quadratic and poly.is_quadratic:
1100 term = sum(map(func, roots_quadratic(poly)))
1101 else:
1102 if not rational or not auto:
1103 term = cls._new(poly, func, auto)
1104 else:
1105 term = cls._rational_case(poly, func)
1107 terms.append(k*term)
1109 return mul_const*Add(*terms) + deg*add_const
1111 @classmethod
1112 def _new(cls, poly, func, auto=True):
1113 """Construct new raw ``RootSum`` instance. """
1114 obj = Expr.__new__(cls)
1116 obj.poly = poly
1117 obj.fun = func
1118 obj.auto = auto
1120 return obj
1122 @classmethod
1123 def new(cls, poly, func, auto=True):
1124 """Construct new ``RootSum`` instance. """
1125 if not func.expr.has(*func.variables):
1126 return func.expr
1128 rational = cls._is_func_rational(poly, func)
1130 if not rational or not auto:
1131 return cls._new(poly, func, auto)
1132 else:
1133 return cls._rational_case(poly, func)
1135 @classmethod
1136 def _transform(cls, expr, x):
1137 """Transform an expression to a polynomial. """
1138 poly = PurePoly(expr, x, greedy=False)
1139 return preprocess_roots(poly)
1141 @classmethod
1142 def _is_func_rational(cls, poly, func):
1143 """Check if a lambda is a rational function. """
1144 var, expr = func.variables[0], func.expr
1145 return expr.is_rational_function(var)
1147 @classmethod
1148 def _rational_case(cls, poly, func):
1149 """Handle the rational function case. """
1150 roots = symbols('r:%d' % poly.degree())
1151 var, expr = func.variables[0], func.expr
1153 f = sum(expr.subs(var, r) for r in roots)
1154 p, q = together(f).as_numer_denom()
1156 domain = QQ[roots]
1158 p = p.expand()
1159 q = q.expand()
1161 try:
1162 p = Poly(p, domain=domain, expand=False)
1163 except GeneratorsNeeded:
1164 p, p_coeff = None, (p,)
1165 else:
1166 p_monom, p_coeff = zip(*p.terms())
1168 try:
1169 q = Poly(q, domain=domain, expand=False)
1170 except GeneratorsNeeded:
1171 q, q_coeff = None, (q,)
1172 else:
1173 q_monom, q_coeff = zip(*q.terms())
1175 coeffs, mapping = symmetrize(p_coeff + q_coeff, formal=True)
1176 formulas, values = viete(poly, roots), []
1178 for (sym, _), (_, val) in zip(mapping, formulas):
1179 values.append((sym, val))
1181 for i, (coeff, _) in enumerate(coeffs):
1182 coeffs[i] = coeff.subs(values)
1184 n = len(p_coeff)
1186 p_coeff = coeffs[:n]
1187 q_coeff = coeffs[n:]
1189 if p is not None:
1190 p = Poly(dict(zip(p_monom, p_coeff)), *p.gens).as_expr()
1191 else:
1192 (p,) = p_coeff
1194 if q is not None:
1195 q = Poly(dict(zip(q_monom, q_coeff)), *q.gens).as_expr()
1196 else:
1197 (q,) = q_coeff
1199 return factor(p/q)
1201 def _hashable_content(self):
1202 return (self.poly, self.fun)
1204 @property
1205 def expr(self):
1206 return self.poly.as_expr()
1208 @property
1209 def args(self):
1210 return (self.expr, self.fun, self.poly.gen)
1212 @property
1213 def free_symbols(self):
1214 return self.poly.free_symbols | self.fun.free_symbols
1216 @property
1217 def is_commutative(self):
1218 return True
1220 def doit(self, **hints):
1221 if not hints.get('roots', True):
1222 return self
1224 _roots = roots(self.poly, multiple=True)
1226 if len(_roots) < self.poly.degree():
1227 return self
1228 else:
1229 return Add(*[self.fun(r) for r in _roots])
1231 def _eval_evalf(self, prec):
1232 try:
1233 _roots = self.poly.nroots(n=prec_to_dps(prec))
1234 except (DomainError, PolynomialError):
1235 return self
1236 else:
1237 return Add(*[self.fun(r) for r in _roots])
1239 def _eval_derivative(self, x):
1240 var, expr = self.fun.args
1241 func = Lambda(var, expr.diff(x))
1242 return self.new(self.poly, func, self.auto)