allow empty return value for realpolyroots
[PyX.git] / pyx / helper.py
blobd500aa01b62af9526aa97133dd8951d08caf08b4
1 #!/usr/bin/env python
2 # -*- coding: ISO-8859-1 -*-
5 # Copyright (C) 2002-2004 Jörg Lehmann <joergl@users.sourceforge.net>
6 # Copyright (C) 2002-2005 André Wobst <wobsta@users.sourceforge.net>
7 # Copyright (C) 2005 Michael Schindler <m-schindler@users.sourceforge.net>
9 # This file is part of PyX (http://pyx.sourceforge.net/).
11 # PyX is free software; you can redistribute it and/or modify
12 # it under the terms of the GNU General Public License as published by
13 # the Free Software Foundation; either version 2 of the License, or
14 # (at your option) any later version.
16 # PyX is distributed in the hope that it will be useful,
17 # but WITHOUT ANY WARRANTY; without even the implied warranty of
18 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
19 # GNU General Public License for more details.
21 # You should have received a copy of the GNU General Public License
22 # along with PyX; if not, write to the Free Software
23 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
26 def sign(x):
27 return (x >= 0) and 1 or -1
29 # XXX fallback for Numeric (eigenvalue computation) to be implemented along
30 # know algorithms (like from numerical recipes)
32 import Numeric, LinearAlgebra
34 def realpolyroots(coeffs, epsilon=1e-5, polish=0):
36 """returns the roots of a polynom with given coefficients
38 This helper routine uses the package Numeric to find the roots
39 of the polynomial with coefficients given in coeffs:
40 0 = \sum_{i=0}^N x^i coeffs[i]
41 The solution is found via an equivalent eigenvalue problem
42 """
44 if not coeffs:
45 return []
47 try:
48 1.0 / coeffs[-1]
49 except ZeroDivisionError:
50 roots = realpolyroots(coeffs[:-1], epsilon=epsilon)
51 else:
52 # divide the coefficients by the maximal order
53 # this is to be done in the matrix, anyhow and
54 # makes checking more stable
55 coeffs = [coeff/coeffs[-1] for coeff in coeffs]
57 N = len(coeffs) - 1
58 # build the Matrix of the polynomial problem
59 mat = Numeric.zeros((N, N), Numeric.Float)
60 for i in range(N-1):
61 mat[i+1][i] = 1
62 for i in range(N):
63 mat[0][i] = -coeffs[N-1-i]
64 # find the eigenvalues of the matrix (== the roots of the polynomial)
65 roots = [complex(root) for root in LinearAlgebra.eigenvalues(mat)]
66 # take only the real roots
67 roots = [root.real for root in roots if -epsilon < root.imag < epsilon]
69 # polish the roots with Newton-Raphson
70 if polish:
71 def polish(root, epsilon):
72 polynom = 2*epsilon
73 while abs(polynom) > epsilon:
74 polynom = coeffs[N]*root + coeffs[N-1]
75 poprime = coeffs[N]*N
76 for i in range(N-2,-1,-1):
77 polynom = polynom*root + coeffs[i]
78 poprime = poprime*root + coeffs[i+1]*(i+1)
79 root -= polynom / poprime
80 return root
82 roots = [polish(root, epsilon) for root in roots]
84 # # check if the roots are really roots!
85 # for root in roots:
86 # polynom = coeffs[N]*root + coeffs[N-1]
87 # for i in range(N-2,-1,-1):
88 # polynom = polynom*root + coeffs[i]
89 # if abs(polynom) > epsilon:
90 # raise Exception("value %f instead of 0" % polynom)
92 return roots