/* LOESS.G - Local Regression and optional Robust Fitting and Symmetric Errors
** Simon Jackman, WWS, Princeton U. e-mail: sjackman@edith.princeton.edu
** Provided free, without guarantee, for public non-commercial use.
** Posted: Mar 96
**
** Usage: {xstar,ystar,res,yhat}=loess(x,y);
**
** Inputs:  x       - n by 1 vector
**          y       - n by 1 vector
**
** Outputs: xstar   - m by 1 vector of evenly spaced values along x axis
**          ystar   - m by 1 vector of fitted loess values: g(xstar)
**          res     - n by 1 vector of residuals from loess fit
**          yhat    - b by 1 vector of fitted loess : g(x)
**
** Globals: _span   - controls degree of smoothing (higher values tend to
**                    linear fit).  Default=2/3.
**          _degree - degree of fitted polynomial g(x).  1 = linear (default),
**                    2 = quadratic.
**          _family - 1 = Gaussian errors (default). 2 = Robust, symmetric.
**          _eval   - number of evenly spaced values to (linearly) interpolate
**                    fitted values.  Default = 50.
**          _Loutput - 0 = no results printed (default), anything else prints
**                     (sometimes lots...)
**
** Comments: Hacked (in part) from a SAS/IML macro I saw floating around the
**           net.  Took ages to port to GAUSS, since sortind(x) is a little
**           buggy (at least in 2.2). e.g., x={1,1,2,3,1,4,1,5,6};
**                                 sortind(x)= 7 2 1 5 3 4 6 8 9 - go figure...
**           Lots of fixes in here to deal with various problems
**           that crop up, but aren't talked about in the references.
**           Experiment with _span, and use the residuals for specification
**           diagnostics.  This runs very fast compared with Splus on a
**           shared UNIX box...
**           I plot outside of the procedure so as to get more problem-
**           specific control over axes, labelling, etc (especially when
**           passing a transformed var to the loess procedure).
**           I've tried to implement the Splus defaults from scatter.smooth
**           as much as possible.
**
**  References:
**  William S. Cleveland. 1979. Robust Locally Weighted Regression and
**     Smoothing Scatterplots.  _JASA_. 74:829-836.
**  William S. Cleveland and Susan J. Devlin. 1988. Locally Weighted
**                 Regression: An Approach to Regression Analysis by Local
**                 Fitting. _JASA_. 83:596-609. [mainly on multivariate loess]
**  William S. Cleveland, Eric Grosse, and William M. Shyu. 1991.
**                 Local Regression Models. in John M. Chambers and Trevor
**                 J. Hastie (eds). _Statistical Models in S_. Pacific Grove:
**                 Wadsworth.
**  Michael Friendly's routines in _SAS System for Statistical
**                 Graphics, First Edition_. Cary, NC: SAS Institute.
**  Trevor J. Hastie and R.J. Tibshirani. 1990. _Generalized
**                 Additive Models_. London: Chapman and Hall.
*/

/* use Splus option names and defaults as much as possible */
declare matrix _span ?= .67;         /* set default smoothing parameter */
declare matrix _degree ?= 1;         /* locally linear = 1, quadratic =2 */
declare matrix _eval ?= 50;          /* evaluate loess at 50 evenly spaced
                                        points over range of data */
declare matrix _family ?= 1;         /* gaussian weights = 1, symmetric = 2 */
declare matrix _Loutput ?=0;         /* print results? no=0, yes = o'wise */

#include indexcat.src;
proc(2)=wls(x,y,w,i);                /* Weighted least squares */
  local xpx,xpy,b,sing;
  sing=0;
  if _degree .eq 2;
     x = ones(rows(x),1)~x~x^2;      /* quadratic */
  else;
     x = ones(rows(x),1)~x;          /* linear    */
  endif;
  xpx = x'diagrv(eye(rows(x)),w)*x;
  xpy = x'diagrv(eye(rows(x)),w)*y;
  if abs(det(xpx)) .gt .0000001;
     b=inv(xpx)*xpy;
  else; if _degree .eq 2; b=0|0|0;
           else; b=0|0; endif;
  print "Matrix singular at observation " i; sing=1; endif;
  retp(b,sing); endp;

proc(1)=robust(e);  /* calculate robustness weights pp315-316 */
  local m,r;
    m=median(abs(e)); r=bisquare(e,6*m); retp(r); endp;

proc(1)=bisquare(u,b);  /* bi-square function, p316 */
  u=u./b; retp((abs(u) .lt 1) .* ((1-u^2)^2)); endp;

proc(4)=loess(x,y);
  local data,n,q,iter,i,res,wts,b,yhat,dist,s,near,nx,ny,d,delta,ystar,
        xstar,u,x0,x1,x2,y1,y2,loind,upind,singlist,slope,dx,step,fixx,
        sing,flag2,mask,fmt,unit,flag;

print; print;
"*********** Locally weighted robust scatterplot smoothing ***********";
  print;
  n=rows(x);
  unit=seqa(1,1,n);
  if floor(_span*n) .lt 2;               /* check for valid _span */
  print "Invalid value for _span - setting to default (2/3)"; print;
  _span=2/3; endif;
  if _family .eq 1; print "Gaussian weights";
  else;             print "Symmetric weights"; endif;
  if _degree .eq 1; print "Locally linear";
  else;             print "Locally quadratic"; endif;
        format 4,2; print "Span = " _span;
        format 3,0; print "Number of evaluation points " _eval;
  format 8,4;

  data=sortc(x~y,1);                   /* sort data first */
  x=data[.,1]; y=data[.,2];

  yhat=zeros(n,1);
  delta=ones(n,1);                     /* initialize robustness weights */
  res=y;                               /* initialize residuals */
  s=ones(n,1);                         /* initialize neighbor ranks */
  singlist=zeros(n,1);                 /* initialize singularity indicators */

  xstar=seqa(minc(x),(maxc(x)-minc(x))./(_eval-1),_eval);
  yhat=ones(n,1);

  iter=1;
  do until iter .gt _family;           /* iterations start
                                          2 iterations for symmetric,
                                          1 for Gaussian */
     format 1,0; print "Iteration " iter; format 8,4;
     i=1;
     do until i gt n;                  /* loop over x values */
        flag2=0;                       /* initialize singularity flag */
loop:   if _span .le 1; q=floor(_span.*n);
        else; q=n; endif;              /* number of nearest neighbors */
        x0=x[i];                       /* reference obs for this iteration */
        dist=abs(x-x0);                /* distance to other points */
        s=sortc(x~y~unit~dist,4);      /* identify sorted distances */
        s=s[1:q,.];
        nx=s[.,1]; ny=s[.,2];          /* q nearest neighbors of x0 */
        d=maxc(s[.,4]);                /* distance to furthest nearest */
        if _span .gt 1; d=d.*_span; endif;    /* p314, CGS */
        if d .gt 0;
           u=abs(nx-x0)./d;
           wts=(u .lt 1).*((1-u^3)^3);        /* tri-cube neighborhood wts */
           wts=delta[s[.,3]].*wts;            /* symmetric robustness weights */
           if sumc(wts[2:q]) .gt .0001;
              {b,sing}=wls(nx,ny,wts,i);
              if sing .eq 1;
                 print "Singularity obs " i " x = " x0 " y = " y[i];
                 q=n;
                 if flag2 .ne 1;
                 print "trying again..."; flag2=1; goto loop; endif;
              endif;
              singlist[i]=sing;
              if _degree .eq 2;
                 yhat[i] = (1~x0~(x0^2))*b;   /* smoothed value - quad */
              else; yhat[i] = (1~x0)*b;       /* smoothed value - linear*/
              endif;
           else; yhat[i] = y[i];
           endif;
        else; yhat[i] = meanc(ny);
        print "x(0) = 0.  obs  " i " x = " x0 " y = " y[i];
        singlist[i]=1;
        endif;
        i=i+1;
     endo;

     /* singularity fixup - interpolation */
     if sumc(singlist) .ne 0;
        print "Singularity fixup (interpolation) for " sumc(singlist) " obs";
        fixx=indexcat(singlist,1);
        i=1;
        do until i eq rows(fixx);
           if fixx[i] .eq 1 .or fixx[i] .eq n; yhat[i]=y[i];
              i=i+1; continue; endif;
           dx=x[fixx[i]]-x[fixx[i]-1];
           flag=0; step=0;
           do until flag .eq 1;
              step=step+1;
              if scalerr(indexcat(fixx,fixx[i]+step)) .eq 13; flag=1; endif;
           endo;
           slope=(yhat[fixx[i]+step]-yhat[fixx[i]-1]);
           slope=slope./(x[fixx[i]+step]-x[fixx[i]-1]);
           yhat[fixx[i]]=yhat[fixx[i]-1]+(slope.*dx);
        i=i+1;
        endo;
     endif;

     res=y-yhat;                              /* residuals */
     if _family .eq 2;
        delta=robust(res);                    /* get robustness weights */
     endif;
     iter=iter+1;
  endo;

  if _Loutput .ne 0;
    if _family .eq 2;
      print "Obs    X      Y     Fitted Y    Weights  Residuals";
      mask={ 1 1 1 1 1 1};
      let fmt[6,3]=
      "*.*lf " 3 0
      "*.*lf" 8 3 "*.*lf" 8 3 "*.*lf" 8 3 "*.*lf" 8 3 "*.*lf" 8 3 ;
      call printfm(seqa(1,1,n)~x~y~yhat~delta~res,mask,fmt);
    else;
      print "Obs    X      Y     Fitted Y    Residuals";
      mask={ 1 1 1 1 1 };
      let fmt[5,3]=
      "*.*lf " 3 0 "*.*lf" 8 3 "*.*lf" 8 3 "*.*lf" 8 3 "*.*lf" 8 3 ;
      call printfm(seqa(1,1,n)~x~y~yhat~res,mask,fmt);
    endif;
  endif;

/* I N T E R P O L A T I O N */
  format 3,0; print "Evaluating at " _eval " equally-spaced points";
  format 8,4;
  if _eval .eq n .or abs(_eval-rows(unique(x,1))) .lt 2;
     xstar=unique(x,1); ystar=yhat[uniqindx(x,1)]; goto done; endif;
  ystar=ones(_eval,1);
  ystar[1]=yhat[1]; ystar[_eval]=yhat[n];   /* fix end points */
  i = 2;
  do until i .eq _eval;
     d=x-xstar[i];
     upind=minindc(recode(d,d .lt 0,maxc(d)+1));
     loind=maxindc(recode(d,d .gt 0,minc(d)-1));
     x1=x[loind]; x2=x[upind]; y1=yhat[loind]; y2=yhat[upind];
     ystar[i]=y1+((y2-y1)./(x2-x1).*(xstar[i]-x1));
     i=i+1;
  endo;
done:
  retp(xstar,ystar,res,yhat);
endp;
