# This stuff is based on algorithms presented in
# "Numerical Recipes - The Art of Scientific Computing", by Press et al.
#
# Finds roots (real and complex) of a polynomial.
# For usage see note preceeding definition of zroots().
#
# MWS, July 1992.

silent

#
# Given coeffs of poly of degree (m+1) in a, and estimate x of root,
# this function improves x to desired accuracy; if polish != 0, this is
# to machine's roundoff-limit, otherwise it is to eps.
# NB: assumes arraybase is zero, as set in zroots() which calls laguer()
#
LEPSS = 6.e-8
LMAXIT = 100	# limit on iterations for convergence
#
func laguer(@a,m,*x,eps,polish) = {
	local j,iter,err,dxold,cdx,abx,sq,h,gp,gm,g2,g,b,d,dx,f,x1

	dxold=abs(x)
	for (iter=1;iter<=LMAXIT;iter+=1) {
		b=a[m]
		err=abs(b)
		d=f=0
		abx=abs(x)
		for (j=m-1;j>=0;j-=1) {
			# compute value of poly and 1st & 2nd derivatives at x
			f=x*f+d
			d=x*d+b
			b=x*b+a[j]
			err=abs(b)+abx*err;
		}
		err *= LEPSS			# estimate roundoff error
		if (abs(b) <= err) return 1	# we are at root
		g=d/b				# use Laguerre's formula
		g2=g*g
		h=g2-2*f/b
		sq=sqrt((m-1)*(m*h-g2))
		gp=g+sq
		gm=g-sq
		if (abs(gp) < abs(gm)) gp=gm
		dx=m/gp
		x1=x-dx
		if (x == x1) return 1		# convergence achieved
		x=x1
		cdx=abs(dx);
		dxold=cdx;
		if (!polish)
			if (cdx <= eps*abs(x)) return 1
	}
	error(0)	# too many iterations
}


#
# This is the main interface. Call zroots with two arrays:
# 	a	contains the coefficients of poly a[0] + a[1]*x +...+ a[m]*x^m
#	roots	storage area for roots of poly
# The polish parameter determines whether highest accuracy is required. Set
# to 1 if this is desired, otherwise set it to zero.
#
ZEPS=2.0e-6		# desired accuracy
array zrwkspc[1]	# resize as required
#
func zroots(@a,@roots,polish) = {
	local m,jj,j,ii,x,b,c,finis,oldbase

	oldbase = arraybase(0)	# we want a[0] as first element
	m = sizeof(a)-1		# m is degree of poly
	resize(zrwkspc,m+1)	# make workspace big enough

	for (j=0;j<=m;j+=1) zrwkspc[j]=a[j]	# copy coefficients
	for (j=m;j>=1;j-=1) {			# loop over each root to be found
		x=0
		laguer(zrwkspc,j,x,ZEPS,0)	# find root
		if (abs(Im(x)) <= (2.0*ZEPS*abs(Re(x)))) x=Re(x)
		roots[j-1]=x
		b=zrwkspc[j]			# forward deflation

		for (jj=j-1;jj>=0;jj-=1) {
			c=zrwkspc[jj]
			zrwkspc[jj]=b
			b=x*b+c
		}
	}
	if (polish)		# polish roots
		for (j=0;j<m;j+=1)
			laguer(a,m,roots[j],ZEPS,1);

	arraybase(oldbase)
	1			# standard default return-value
}


#
# from poly.ic - evaluates poly in c at x
#
func epoly(@c, x) = {
	local n,j,ss,oldbase

	oldbase = arraybase(1)
	n = sizeof(c)

	ss = c[n]
	for (j = n-1; j >= 1; j -= 1)
		ss = ss*x+c[j]

	arraybase(oldbase)
	ss
}


#
#an example of usage
#

N=7	#degree
array a[N+1] = {1,2,3,4,5,6,7,8}	# polynomial
array r[N]				# where to store roots

echo "Routines defined. Performing example:"
echo "\tPoly is 1+2x+3x^2+...+8x^7 - be patient..."
zroots(a,r,1)
echo "done. Roots are:"

display(r)
echo "\nPoly evaluated at each root:"
for (j=1; j <= N; j+=1)
	print(epoly(a,r[j]))

verbose