#!/usr/bin/env python
"""Simple demonstration of polynomials and root finding."""
import numpy as N
import scipy as S
import pylab as P
def plot_poly(x,y,roots):
"""Simple utility to make a charting screen with x/y axes"""
P.figure()
P.axhline(0,color='g')
P.axvline(0,color='g')
P.grid()
P.plot(x,y,'b-')
P.scatter(roots.real,roots.imag,s=80,c='r')
P.xlabel('Re')
P.ylabel('Im')
# Create the coefficients for a polynomial with nroots_minus1 at x=-1, one
# root at -2 and one root at 1:
nroots_minus1 = 10
coefs = reduce(N.convolve,[[1,1]]*nroots_minus1+[[1,2],[1,-1]])
# Construct the polynomial and get its roots
pol = S.poly1d(coefs)
roots = pol.r
print 'Polynomial p(x):\n',pol,'\n'
print 'p(x) built with coefs:',coefs
print 'Roots of p(x):',roots
# Plot p(x)
x = N.linspace(-4,4,400)
y = pol(x)
# Show roots
plot_poly(x,y,roots)
P.xlim(-3,2)
P.ylim(-250,100)
P.title('All roots')
# Set limits and grid to make the plot clear
#P.ylim(-40,40)
plot_poly(x,y,roots)
eps = 0.05
P.xlim(-1-eps,-1+eps)
P.ylim(-eps,eps)
P.title('Root cluster around $x=-1$')
# Display on screen
P.show()
