Finding (all) numeric roots of a nth degree polynomial – with python

This is a code I originally wrote in Matlab for my Computational Physics class. After that, I decided to rewrite it in python just to start programming in this awesomely easy to learn programming language.

This program uses the Newton-Rapson based algorithm described in
this paper by Sasan A. Muheadden.

The main function will receive just one parameter:

An array of float or integer numbers that represent the coefficients of our nth degree polynomial. See the examples below.

This is the polynomial we want to solve: x^5+x^3-1.42*x^2+1
This is the input array for the main function: [1,0,1,-1.42,0,1]

Notice that the input array has n+1 elements.

First we get our necessary variables for the algorithm

import cmath as m
import random as rnd

relative_tolerance = 0.000000001
interval_magnitude = 3
nmax = 1000

I created an equality function for complex numbers to have a more concise syntax in the following defined functions.


def eq(a,b,rel_tol=relative_tolerance):
    if(isinstance(a,complex) and isinstance(b,complex)):
        return m.isclose(a,b,rel_tol=rel_tol)
    else:
        return False

The nice things about polynomials is that their derivatives are dead easy to program since after differentiating we get another polynomial of degree n-1, so we get an easily recognizable recursive relation. The following function calculates takes as parameter an array with the coefficients of the original function and returns an array with the coefficients of the resulting polynomial.


def derivative(poly):
    deriv = []
    deg = len(poly)
    n = deg-1
    for i in range(0,deg):
        C = poly[i]
        deriv.append(C*n)
        n -= 1
    deriv.pop()
    return deriv

The function “eval”, as you may anticipate, evaluates a given polynomials at a given x. It takes as parameter an array with the coefficients of the polynomial to be evaluated and the value x to evaluate it at.


def eval(poly,x):
    eval = 0
    deg = len(poly)
    n = deg-1
    for i in range(0,deg):
        C = poly[i]
        eval += C*x**n
        n -= 1
    return eval

find_root will take as a parameter an initial estimation of a root of the polynomial represented by the variable “poly” (which, again, is an array with the coefficients of a polynomial).


def find_root(z0,poly,rel_tol=relative_tolerance,verbose=False):
    k = 0
    x = z0
    if(verbose == True):
        print('........................................')
    while True:
        f = eval(poly,x)
        f_prime = eval(derivative(poly),x)
        x = x - f/f_prime
        k += 1
        if(k == nmax):
            x = None
            break
        cur = x
        cur_f = eval(poly,x)
        if(verbose==True):
            msg = "| zr= ( %.6f, %.6f ), f(zr)= ( %.6f, %.6f ) |" % (cur.real, cur.imag, cur_f.real, cur_f.imag)
            print(msg)
        if abs(cur_f) <span id="mce_SELREST_end" style="overflow:hidden;line-height:0;"></span>&lt; rel_tol:
            break;
    return x

Aaand finally, the function "find_roots" is the main function where all the previously defined functions will be encapsuled. If you want to see the steps in finding the roots pass on the parameter verbose=True.

This function takes as parameter the poly array, which if you read everything above must know by now.

def find_roots(poly,rel_tol=relative_tolerance,int_mag=interval_magnitude,verbose = True):
found = False
counter = 1
deg = len(poly) – 1

if(verbose == True):
print('| polynomial_deg = ' + str(deg) + ' | \n')

C0 = poly[-1]
z0 = complex(abs(C0)+1,rnd.randint(-int_mag,int_mag)/rnd.randint(1,int_mag))
if(verbose == True):
print('| trial ' + str(counter) + ' |')
roots = [find_root(z0,poly,verbose=verbose)]
# prepare msg letting know if a root was found
if(verbose == True):
if roots[0] != None:
msg = "| current_root = ( %.6f, %.6f ) |" % (roots[0].real, roots[0].imag)
counter += 1
elif roots[0] == None:
msg = "| algorithm failed to find a root, try a different z0 |"
print('………………………………….')
print(msg)
print("\n")

while True:
# we are dividing to get a floating point number
# denominator interval is [1,int_mag] to avoid zeroes
a = rnd.randint(-int_mag,int_mag)/rnd.randint(1,int_mag)
b = rnd.randint(-int_mag,int_mag)/rnd.randint(1,int_mag)
z0 = complex(a,b)
if(verbose == True):
print('| trial ' + str(counter) + ' |')

current_root = find_root(z0,poly,verbose=verbose)

# prepare msg letting know if a root was found
if(verbose == True):
if current_root != None:
msg = "| current_root = ( %.6f, %.6f ) |" % (current_root.real, current_root.imag)
elif current_root == None:
msg = "| algorithm failed to find a root, try a different z0 |"
print('………………………………….')
print(msg)
print("\n")

if(counter == nmax):
print('WARNING: could only find ' + str(len(roots)) + ' out of ' + str(deg) + ' roots')
break

# about to make a possibly innecesary for loop
# return and modify this algorithm when
# you've found a better way to do it
for root in roots:
if(eq(root,current_root)):
found = True

if(found == False and current_root != None):
roots.append(current_root)
elif(found == True):
found = False

# deg + 1 accounts for the number of elements in polynomial
# (deg + 1) + 1 accounts for the 'starter' string inside roots
if(len(roots) == deg):
break

counter += 1
return roots

Finally, we run test the main function.

poly = [1,0,-1,0,-3,0,2,-6.4,0,-5,-2.2]
find_roots(poly,verbose=False)

These are the n roots.

[(1.622862875657454+0j),
(0.01189799744808773+0.8942986469642933j),
(0.011897997448073974-0.8942986469642864j),
(-0.369536384298153+4.895865258206433e-16j),
(-0.1539012130022848+1.2799938044260128j),
(0.8964383223119088+0.7405519518754988j),
(0.896438322311946-0.7405519518761811j),
(-1.381098352437382+0.36525414989632693j),
(-0.15390121300228396-1.2799938044260104j),
(-1.3810983524373823-0.365254149896327j)]

Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s