polynomials.py 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213
  1. from ..libmp.backend import xrange
  2. from .calculus import defun
  3. #----------------------------------------------------------------------------#
  4. # Polynomials #
  5. #----------------------------------------------------------------------------#
  6. # XXX: extra precision
  7. @defun
  8. def polyval(ctx, coeffs, x, derivative=False):
  9. r"""
  10. Given coefficients `[c_n, \ldots, c_2, c_1, c_0]` and a number `x`,
  11. :func:`~mpmath.polyval` evaluates the polynomial
  12. .. math ::
  13. P(x) = c_n x^n + \ldots + c_2 x^2 + c_1 x + c_0.
  14. If *derivative=True* is set, :func:`~mpmath.polyval` simultaneously
  15. evaluates `P(x)` with the derivative, `P'(x)`, and returns the
  16. tuple `(P(x), P'(x))`.
  17. >>> from mpmath import *
  18. >>> mp.pretty = True
  19. >>> polyval([3, 0, 2], 0.5)
  20. 2.75
  21. >>> polyval([3, 0, 2], 0.5, derivative=True)
  22. (2.75, 3.0)
  23. The coefficients and the evaluation point may be any combination
  24. of real or complex numbers.
  25. """
  26. if not coeffs:
  27. return ctx.zero
  28. p = ctx.convert(coeffs[0])
  29. q = ctx.zero
  30. for c in coeffs[1:]:
  31. if derivative:
  32. q = p + x*q
  33. p = c + x*p
  34. if derivative:
  35. return p, q
  36. else:
  37. return p
  38. @defun
  39. def polyroots(ctx, coeffs, maxsteps=50, cleanup=True, extraprec=10,
  40. error=False, roots_init=None):
  41. """
  42. Computes all roots (real or complex) of a given polynomial.
  43. The roots are returned as a sorted list, where real roots appear first
  44. followed by complex conjugate roots as adjacent elements. The polynomial
  45. should be given as a list of coefficients, in the format used by
  46. :func:`~mpmath.polyval`. The leading coefficient must be nonzero.
  47. With *error=True*, :func:`~mpmath.polyroots` returns a tuple *(roots, err)*
  48. where *err* is an estimate of the maximum error among the computed roots.
  49. **Examples**
  50. Finding the three real roots of `x^3 - x^2 - 14x + 24`::
  51. >>> from mpmath import *
  52. >>> mp.dps = 15; mp.pretty = True
  53. >>> nprint(polyroots([1,-1,-14,24]), 4)
  54. [-4.0, 2.0, 3.0]
  55. Finding the two complex conjugate roots of `4x^2 + 3x + 2`, with an
  56. error estimate::
  57. >>> roots, err = polyroots([4,3,2], error=True)
  58. >>> for r in roots:
  59. ... print(r)
  60. ...
  61. (-0.375 + 0.59947894041409j)
  62. (-0.375 - 0.59947894041409j)
  63. >>>
  64. >>> err
  65. 2.22044604925031e-16
  66. >>>
  67. >>> polyval([4,3,2], roots[0])
  68. (2.22044604925031e-16 + 0.0j)
  69. >>> polyval([4,3,2], roots[1])
  70. (2.22044604925031e-16 + 0.0j)
  71. The following example computes all the 5th roots of unity; that is,
  72. the roots of `x^5 - 1`::
  73. >>> mp.dps = 20
  74. >>> for r in polyroots([1, 0, 0, 0, 0, -1]):
  75. ... print(r)
  76. ...
  77. 1.0
  78. (-0.8090169943749474241 + 0.58778525229247312917j)
  79. (-0.8090169943749474241 - 0.58778525229247312917j)
  80. (0.3090169943749474241 + 0.95105651629515357212j)
  81. (0.3090169943749474241 - 0.95105651629515357212j)
  82. **Precision and conditioning**
  83. The roots are computed to the current working precision accuracy. If this
  84. accuracy cannot be achieved in ``maxsteps`` steps, then a
  85. ``NoConvergence`` exception is raised. The algorithm internally is using
  86. the current working precision extended by ``extraprec``. If
  87. ``NoConvergence`` was raised, that is caused either by not having enough
  88. extra precision to achieve convergence (in which case increasing
  89. ``extraprec`` should fix the problem) or too low ``maxsteps`` (in which
  90. case increasing ``maxsteps`` should fix the problem), or a combination of
  91. both.
  92. The user should always do a convergence study with regards to
  93. ``extraprec`` to ensure accurate results. It is possible to get
  94. convergence to a wrong answer with too low ``extraprec``.
  95. Provided there are no repeated roots, :func:`~mpmath.polyroots` can
  96. typically compute all roots of an arbitrary polynomial to high precision::
  97. >>> mp.dps = 60
  98. >>> for r in polyroots([1, 0, -10, 0, 1]):
  99. ... print(r)
  100. ...
  101. -3.14626436994197234232913506571557044551247712918732870123249
  102. -0.317837245195782244725757617296174288373133378433432554879127
  103. 0.317837245195782244725757617296174288373133378433432554879127
  104. 3.14626436994197234232913506571557044551247712918732870123249
  105. >>>
  106. >>> sqrt(3) + sqrt(2)
  107. 3.14626436994197234232913506571557044551247712918732870123249
  108. >>> sqrt(3) - sqrt(2)
  109. 0.317837245195782244725757617296174288373133378433432554879127
  110. **Algorithm**
  111. :func:`~mpmath.polyroots` implements the Durand-Kerner method [1], which
  112. uses complex arithmetic to locate all roots simultaneously.
  113. The Durand-Kerner method can be viewed as approximately performing
  114. simultaneous Newton iteration for all the roots. In particular,
  115. the convergence to simple roots is quadratic, just like Newton's
  116. method.
  117. Although all roots are internally calculated using complex arithmetic, any
  118. root found to have an imaginary part smaller than the estimated numerical
  119. error is truncated to a real number (small real parts are also chopped).
  120. Real roots are placed first in the returned list, sorted by value. The
  121. remaining complex roots are sorted by their real parts so that conjugate
  122. roots end up next to each other.
  123. **References**
  124. 1. http://en.wikipedia.org/wiki/Durand-Kerner_method
  125. """
  126. if len(coeffs) <= 1:
  127. if not coeffs or not coeffs[0]:
  128. raise ValueError("Input to polyroots must not be the zero polynomial")
  129. # Constant polynomial with no roots
  130. return []
  131. orig = ctx.prec
  132. tol = +ctx.eps
  133. with ctx.extraprec(extraprec):
  134. deg = len(coeffs) - 1
  135. # Must be monic
  136. lead = ctx.convert(coeffs[0])
  137. if lead == 1:
  138. coeffs = [ctx.convert(c) for c in coeffs]
  139. else:
  140. coeffs = [c/lead for c in coeffs]
  141. f = lambda x: ctx.polyval(coeffs, x)
  142. if roots_init is None:
  143. roots = [ctx.mpc((0.4+0.9j)**n) for n in xrange(deg)]
  144. else:
  145. roots = [None]*deg;
  146. deg_init = min(deg, len(roots_init))
  147. roots[:deg_init] = list(roots_init[:deg_init])
  148. roots[deg_init:] = [ctx.mpc((0.4+0.9j)**n) for n
  149. in xrange(deg_init,deg)]
  150. err = [ctx.one for n in xrange(deg)]
  151. # Durand-Kerner iteration until convergence
  152. for step in xrange(maxsteps):
  153. if abs(max(err)) < tol:
  154. break
  155. for i in xrange(deg):
  156. p = roots[i]
  157. x = f(p)
  158. for j in range(deg):
  159. if i != j:
  160. try:
  161. x /= (p-roots[j])
  162. except ZeroDivisionError:
  163. continue
  164. roots[i] = p - x
  165. err[i] = abs(x)
  166. if abs(max(err)) >= tol:
  167. raise ctx.NoConvergence("Didn't converge in maxsteps=%d steps." \
  168. % maxsteps)
  169. # Remove small real or imaginary parts
  170. if cleanup:
  171. for i in xrange(deg):
  172. if abs(roots[i]) < tol:
  173. roots[i] = ctx.zero
  174. elif abs(ctx._im(roots[i])) < tol:
  175. roots[i] = roots[i].real
  176. elif abs(ctx._re(roots[i])) < tol:
  177. roots[i] = roots[i].imag * 1j
  178. roots.sort(key=lambda x: (abs(ctx._im(x)), ctx._re(x)))
  179. if error:
  180. err = max(err)
  181. err = max(err, ctx.ldexp(1, -orig+1))
  182. return [+r for r in roots], +err
  183. else:
  184. return [+r for r in roots]