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>< 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 != None: msg = "| current_root = ( %.6f, %.6f ) |" % (roots.real, roots.imag) counter += 1 elif roots == 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)]