Welcome, guest | Sign In | My Account | Store | Cart

This package does simple polynomial manipulation: adding, multiplying, taking to powers, evaluating at a value, taking integrals and derivatives. Nothing as sophisticated as Mathematica, but useful all the same. I find that I make lots of dumb errors in multiplying out polynomials by hand (when I don't have Mathematica at my disposal), and this little script helps prevent those errors.

Python, 274 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
#!/usr/bin/env python
"""\
 Polynomial.py - A set of utilities to manipulate polynomials. This consists
 of a set of simple functions to convert polynomials to a Python list, and
 manipulate the resulting lists for multiplication, addition, and
 power functions. There is also a class that wraps these functions
 to make their use a little more natural.

 This can also evaluate the polynomials at a value, and take integrals
 and derivatives of the polynomials.

 Written by Rick Muller.
"""
import re

# Define some classes for some syntactic surgar:
class Polynomial:
    def __init__(self,val):
        if type(val) == type([]):
            self.plist = val
        elif type(val) == type(''):
            self.plist = parse_string(val)
        else:
            raise "Unknown argument to Polynomial: %s" % val
        return
    
    def __add__(self,other): return Polynomial(add(self.plist,plist(other)))
    def __radd__(self,other): return Polynomial(add(self.plist,plist(other)))
    def __sub__(self,other):  return Polynomial(sub(self.plist,plist(other)))
    def __rsub__(self,other): return -Polynomial(sub(self.plist,plist(other)))
    def __mul__(self,other): return Polynomial(multiply(self.plist,
                                                        plist(other)))
    def __rmul__(self,other): return Polynomial(multiply(self.plist,
                                                        plist(other)))
    def __neg__(self): return -1*self
    def __pow__(self,e): return Polynomial(power(self.plist,e))
    def __repr__(self): return tostring(self.plist)
    def __call__(self,x1,x2=None): return peval(self.plist,x1,x2)

    def integral(self): return Polynomial(integral(self.plist))
    def derivative(self): return Polynomial(derivative(self.plist))

# Define some simple utility functions. These manipulate "plists", polynomials
#  that are stored as python lists. Thus, 3x^2 + 2x + 1 would be stored
#  as [1,2,3] (lowest to highest power of x, even though polynomials
#  are typically written from highest to lowest power).

def plist(term):
    "Force term to have the form of a polynomial list"
    # First see if this is already a Polynomial object
    try:
        pl = term.plist
        return pl
    except:
        pass

    # It isn't. Try to force coercion from an integer or a float
    if type(term) == type(1.0) or type(term) == type(1):
        return [term]
    elif type(term) == type(''):
        return parse_string(term)
    # We ultimately want to be able to parse a string here
    else:
        raise "Unsupported term can't be corced into a plist: %s" % term
    return None

def peval(plist,x,x2=None):
    """\
    Eval the plist at value x. If two values are given, the
    difference between the second and the first is returned. This
    latter feature is included for the purpose of evaluating
    definite integrals.
    """
    val = 0
    if x2:
        for i in range(len(plist)): val += plist[i]*(pow(x2,i)-pow(x,i))
    else:
        for i in range(len(plist)): val += plist[i]*pow(x,i)
    return val

def integral(plist):
    """\
    Return a new plist corresponding to the integral of the input plist.
    This function uses zero as the constant term, which is okay when
    evaluating a definite integral, for example, but is otherwise
    ambiguous.

    The math forces the coefficients to be turned into floats.
    Consider importing __future__ division to simplify this.
    """
    if not plist: return []
    new = [0]
    for i in range(len(plist)):
        c = plist[i]/(i+1.)
        if c == int(c): c = int(c) # attempt to cast back to int
        new.append(c)
    return new

def derivative(plist):
    """\
    Return a new plist corresponding to the derivative of the input plist.
    """
    new = []
    if not plist: return new
    for i in range(1,len(plist)): new.append(i*plist[i])
    return new

def add(p1,p2):
    "Return a new plist corresponding to the sum of the two input plists."
    if len(p1) > len(p2):
        new = [i for i in p1]
        for i in range(len(p2)): new[i] += p2[i]
    else:
        new = [i for i in p2]
        for i in range(len(p1)): new[i] += p1[i]
    return new

def sub(p1,p2): return add(p1,mult_const(p2,-1))

def mult_const(p,c):
    "Return a new plist corresponding to the input plist multplied by a const"
    return [c*pi for pi in p]

def multiply(p1,p2):
    "Return a new plist corresponding to the product of the two input plists"
    if len(p1) > len(p2): short,long = p2,p1
    else: short,long = p1,p2
    new = []
    for i in range(len(short)): new = add(new,mult_one(long,short[i],i))
    return new

def mult_one(p,c,i):
    """\
    Return a new plist corresponding to the product of the input plist p
    with the single term c*x^i
    """
    new = [0]*i # increment the list with i zeros
    for pi in p: new.append(pi*c)
    return new

def power(p,e):
    "Return a new plist corresponding to the e-th power of the input plist p"
    assert int(e) == e, "Can only take integral power of a plist"
    new = [1]
    for i in range(e): new = multiply(new,p)
    return new

def parse_string(str=None):
    """\
    Do very, very primitive parsing of a string into a plist.
    'x' is the only term considered for the polynomial, and this
    routine can only handle terms of the form:
    7x^2 + 6x - 5
    and will choke on seemingly simple forms such as
    x^2*7 - 1
    or
    x**2 - 1
    """
    termpat = re.compile('([-+]?\s*\d*\.?\d*)(x?\^?\d?)')
    #print "Parsing string: ",str
    #print termpat.findall(str)
    res_dict = {}
    for n,p in termpat.findall(str):
        n,p = n.strip(),p.strip()
        if not n and not p: continue
        n,p = parse_n(n),parse_p(p)
        if res_dict.has_key(p): res_dict[p] += n
        else: res_dict[p] = n
    highest_order = max(res_dict.keys())
    res = [0]*(highest_order+1)
    for key,value in res_dict.items(): res[key] = value
    return res

def parse_n(str):
    "Parse the number part of a polynomial string term"
    if not str: return 1
    elif str == '-': return -1
    elif str == '+': return 1
    return eval(str)

def parse_p(str):
    "Parse the power part of a polynomial string term"
    pat = re.compile('x\^?(\d)?')
    if not str: return 0
    res = pat.findall(str)[0]
    if not res: return 1
    return int(res)

def strip_leading_zeros(p):
    "Remove the leading (in terms of high orders of x) zeros in the polynomial"
    # compute the highest nonzero element of the list
    for i in range(len(p)-1,-1,-1):
        if p[i]: break
    return p[:i+1]

def tostring(p):
    """\
    Convert a plist into a string. This looks overly complex at first,
    but most of the complexity is caused by special cases.
    """
    p = strip_leading_zeros(p)
    str = []
    for i in range(len(p)-1,-1,-1):
        if p[i]:
            if i < len(p)-1:
                if p[i] >= 0: str.append('+')
                else: str.append('-')
                str.append(tostring_term(abs(p[i]),i))
            else:
                str.append(tostring_term(p[i],i))
    return ' '.join(str)

def tostring_term(c,i):
    "Convert a single coefficient c and power e to a string cx^i"
    if i == 1:
        if c == 1: return 'x'
        elif c == -1: return '-x'
        return "%sx" % c
    elif i: 
        if c == 1: return "x^%d" % i
        elif c == -1: return "-x^%d" % i
        return "%sx^%d" % (c,i)
    return "%s" % c

def test():
    print tostring([1,2.,3]) # testing floats
    print tostring([1,2,-3]) # can we handle - signs
    print tostring([1,-2,3])
    print tostring([0,1,2])  # are we smart enough to exclude 0 terms?
    print tostring([0,1,2,0]) # testing leading zero stripping
    print tostring([0,1])
    print tostring([0,1.0]) # testing whether 1.0 == 1: risky
    print tostring(add([1,2,3],[1,-2,3]))  # test addition
    print tostring(multiply([1,1],[-1,1])) # test multiplication
    print tostring(power([1,1],2))         # test power

    # Some cases using the polynomial objects:
    print Polynomial([1,2,3]) + Polynomial([1,2]) # add
    print Polynomial([1,2,3]) + 1                 # add
    print Polynomial([1,2,3])-1                   # sub
    print 1-Polynomial([1,2,3])                   # rsub
    print 1+Polynomial([1,2,3])                   # radd
    print Polynomial([1,2,3])*-1                  # mul
    print -1*Polynomial([1,2,3])                  # rmul
    print -Polynomial([1,2,3])                    # neg
    print ''
    # Work out Niklasson's raising and lowering operators:
    #  tests putting constants into the polynomial.
    for m in range(1,4):
        print 'P^a_%d = ' % m,\
              1 - Polynomial([1,-1])**m*Polynomial([1,m])
        print 'P^b_%d = ' % m,\
              Polynomial([0,1])**m*Polynomial([1+m,-m])
    print ''

    # Test string parsing
    print parse_string('+3x - 4x^2 + 7x^5-x+1')
    print Polynomial(parse_string('+3x - 4x^2 + 7x^5-x+1'))
    print Polynomial('+3x - 4x^2 + 7x^5-x+1')
    print Polynomial('+3x -4x^2+7x^5-x+1') - 7
    print Polynomial('5x+x^2')

    # Test the integral and derivatives
    print integral([])
    print integral([1])
    print integral([0,1])
    print derivative([0,0,0.5])
    p = Polynomial('x')
    ip = p.integral()
    dp = p.derivative()
    print ip,dp
    print ip(0,1) # integral of y=x from (0,1)

if __name__ == '__main__': test()

I'm constantly making mistakes when I multiply out simple polynomials, and this program lets me evaluate, add, subtract, multiply, and take powers of simple polynomials. There are lots of limitations to the functionality here. All polynomials have to use 'x' as the dependent variable, and the syntax that you may use is fairly limited. Still, you can do some neat things like:

>>> from Polynomial import *
>>> Polynomial('x^2+5')**2-1
x^4 + 10x^2 + 24
>>> Polynomial('x^2+5')**2-'x'
x^4 + 10x^2 - x + 25

3 comments

Hunter Peress 16 years, 10 months ago  # | flag
Travis Oliphant 16 years, 9 months ago  # | flag

A similar Polynomial class is in scipy. There is a similar polynomial class in scipy (http://www.scipy.org): scipy.poly1d

Here are some examples:

>>> import scipy
>>> a = scipy.poly1d([1,2,3])
>>> print a
 2
x + 2 x + 3
>>> a = scipy.poly1d([10,20,30])
>>> print a
    2
10 x + 20 x + 30
>>> print a * a
     4       3        2
100 x + 400 x + 1000 x + 1200 x + 900

It has many of the same features and more (also can find the roots of the polynomial).

bruce wernick 12 years, 4 months ago  # | flag

Very nice.

I must say, I prefer this to scipy because it uses only primitive Python. Almost every attempt I have made to use code snippets that import scipy have failed for some or the other reason (other library needed, function not found...).

How would you extend this to solve for x

and then

read n polynomial strings to solve for n unknowns?

Created by Rick Muller on Thu, 6 Jan 2005 (PSF)
Python recipes (4591)
Rick Muller's recipes (10)

Required Modules

Other Information and Tasks