/* Program to accompany 
   Wicklin, Rick (2015),
   "The sensitivity of Newton's method to an initial guess,"
   The DO Loop blog, published 24JUN2015.
   http://blogs.sas.com/content/iml/2015/06/24/sensitivity-newtons-method.html
   
   Define a fifth-degree polynomial with five roots. For each point x, 
   apply Newton's method and count how many iterations it takes to 
   converge to a root. For each value of x, record the number of 
   iterations and the root to which the point converges.
   Use SGPLOT to plot Y=number of iterations versus X=initial condition
   and color-code by the root to which Newton's method converges.

   For the examples in the blog post, 
   f(x) =  1*x##5 -2*x##4 -10*x##3 +20*x##2 +9*x -18
   which has roots at {-3 -1 1 2 3}.
*/
proc iml;
/* evaluate polynomial with coefficients in vector p */
start Func(x) global(p);
   y = j(nrow(x), ncol(x), p[1]); /* initialize to coef of largest deg */
   do j = 2 to ncol(p);           /* Horner's method of evaluation */
      y = y # x + p[j];
   end;
   return(y);
finish;

/* evaluate derivative of polynomial with coefficients in vector p */
start Deriv(x) global(p);
   d = ncol(p)-1;
   p0 = p;
   p = (d:1) # p0[,1:d];  /* coefficients of derivative polynomial */
   y = Func(x);           /* temporarily overwrite p; evaluate derivative */
   p = p0;                /* restore p */
   return( y );
finish;

/* Implementation of Newton's method with default arguments */
/* By default, maximum iterations (maxiter) is 25           */
/*             convergence criterion (converge) is 1e-6     */
/* Modified from 
   http://blogs.sas.com/content/iml/2011/08/05/using-newtons-method-to-find-the-zero-of-a-function.html
*/
start NewtonMethod(x, iter, x0, maxIter=25, converge=1e-6);
   x = x0;
   f = Func(x);       /* evaluate function at starting values */
   do iter = 1 to maxiter           /* iterate until maxiter */
      while(max(abs(f))>converge);         /* or convergence */
      J = Deriv(x);                  /* evaluate derivatives */
      delta = -solve(J, f);   /* solve for correction vector */
      x = x + delta;                /* the new approximation */
      f = func(x);                   /* evaluate the function */
   end;
   /* return missing if no convergence */
   if iter > maxIter then
      x = j(nrow(x0),ncol(x0),.);
finish NewtonMethod;

/* Iterate Newton's method from x. Return number of iterations and the root 
   number to which the method converged.  The roots are in the vector g_roots */
start NewtonConvergence(numIters, rootNumber, x) global(g_roots);
   numIters = j(1,ncol(x),.);
   rootNumber = j(1,ncol(x),.);
   epsilon = 0.01;

   do i = 1 to ncol(x);
      x0 = x[i];
      call NewtonMethod(root, iter, x0, 50);
      idx = loc(root>g_roots-epsilon  & root<g_roots+epsilon);
      numIters[i] = iter;
      if ^IsEmpty(idx) then rootNumber[i] = idx;
   end;
finish;

/* Given the vectors x, numIters, and rootNumber, create a needle
   plot of numIters vs x, colored by rootNumber */
start PlotNewton(x, numIters, rootNumber);
   create NewtonConverge var {"x" "numIters" "rootNumber"};
   append;
   close;

   submit; 
      proc sort data=NewtonConverge;
      by rootNumber x;
      run;

      proc sgplot data=NewtonConverge;
      label x="Initial Guess";
      needle x=x y=numIters / group=rootNumber lineattrs=(thickness=2);
      yaxis grid;
      xaxis grid minor minorcount=1;
      run;
   endsubmit;
finish;

/***************************************************************
  Main driver: Create polynomial and call NewtonConvergence
  and PlotNewton on different domains
 ***************************************************************/

/* c5*x##5 + c4*x##4 + ... + c0 */
p = {1 -2 -10 20 9 -18};       /* Coefficients: p[1]=c5, p[2]=c4, etc */
g_roots = {-3 -1 1 2 3};       /* roots for this problem */
n = 250;                       /* evaluate at n points */

/* apply Newton's method to initial guesses in this domain */
ab = {-4 4};
x = ab[1] + range(ab)/n * (1:n);         
run NewtonConvergence(numIters, rootNumber, x);

title "Newton's Method on Fifth-Degree Polynomial";
title2 "Initial Guess in [-4, 4]";
ods graphics / width=582px;
run PlotNewton(x, numIters, rootNumber);

/* apply Newton's method to initial guesses in this domain */
ab = {-0.5 0.5};
x = ab[1] + range(ab)/n * (1:n);         
run NewtonConvergence(numIters, rootNumber, x);

title "Newton's Method on Fifth-Degree Polynomial";
title2 "Initial Guess in [-0.5, 0.5]";
ods graphics / width=580px;
run PlotNewton(x, numIters, rootNumber);

/* apply Newton's method to initial guesses in this domain */
ab = {0.05 0.12};
x = ab[1] + range(ab)/n * (1:n);         
run NewtonConvergence(numIters, rootNumber, x);

title "Newton's Method on Fifth-Degree Polynomial";
title2 "Initial Guess in [0.05, 0.12]";
ods graphics / width=587px;
run PlotNewton(x, numIters, rootNumber);