/* SAS/IML functions to evaluate the upper and lower branches 
   of the Lambert W function. Program accompanies
   Rick Wicklin, "The Lambert W Function in SAS",
   published 31Aug2016 on The DO Loop blog:
   http://blogs.sas.com/content/iml/2016/08/31/lambert-w-function-sas.html ?

   LambertWp evaluates the upper (principal) branch by using Boyd's
      approximation followed by Halley iterations.
      Input: x >= -1/e 
      Optional input: HalleyIters controls the number of Halley iterations (default 2)
                      CheckBounds is boolean flag that specifies whether input might
                        be outside open interval (-1/e,infinity). If so, handle those values.
      Output: W(x) > -1. Output is same shape as x.

   LambertWm evaluates the lower branch by using Chapeau-Blondeau and Monir (2002)
      approximation followed by Halley itertions.
      Input: x where -1/e <= x < 0
      Optional input: HalleyIters specifies the number of Halley iterations (default 2)
                      CheckBounds is boolean flag that specifies whether input might
                        be outside open interval (-1/e,0). If so, handle those values.
      Output: W(x) < -1. Output is same shape as x.
*/

proc iml;

/************************************************/
/*  Evaluation of upper branch Wp(x)            */
/************************************************/
/* x is a solution to x#exp(x) = c  for x >= -1
   iff
   w=x+1 and y=1+exp(1)*c and w is soln of (w-1)#exp(w) = y-1   for w > 0

   Let g(x) = x#exp(x) and define W(z) implicitly as the value x for which g(x)=z.

   Goal: Evaluate W(z) for W the principal branch of the inverse function.
   1. Given z (where z> -1/e), define y = 1 + ez;
   2. Compute t = B(y) where B is Boyd's direct approximation to W.
   3. Use Newton or Halley iterations to improve approximation. 
      Define f(t) = (t-1)#exp(t) - y + 1
      Then  f'(t) = t#exp(t)
           f''(t) = (t+1)#exp(t)  ==> R(t) = f'' / f' = (t+1)/t
      and new approx is
      t_new = t - f(t) / ( f'(t) - 0.5*f(t)#R(t) )
   4. Let t_f be the final iterate. Then f(t_f)=0 or 
      (t_f-1)#exp(t_f) = y - 1
      Let x = t_f - 1. Then 
      x#exp(x)#e = (1 + ez) - 1 or
      x#exp(x) = z  ==> x solves g(x)=z ==> we have evaluated W(z)
*/
proc iml;
/* MAIN FUNCTION TO CALL FOR POSITIVE BRANCH */
/* positive branch of Lambert W function.
   If you know that all(x > -1/e), no need to check domain */
start LambertWp(x, HalleyIters=2, CheckBounds=1); 
   y = 1 + exp(1)*x;
   if ^CheckBounds then do;
      v = Wp_Boyd1(y);
      w = HalleyRefineWp(y, v, HalleyIters);
      return w-1;
   end;

   /* otherwise, compute y=0 and y>0 separately */
   w = j(nrow(y), ncol(y), .);
   idx = loc(y>0);               /* for y > 0 */
   if ncol(idx)>0 then do;
      v = Wp_Boyd1(y[idx]);
      w[idx] = HalleyRefineWp(y[idx], v, HalleyIters);
   end;
   if ncol(idx)=nrow(y)*ncol(y) then /* all found; return */
      return w-1;
      
   idx = loc(y=0);               /* if any y = 0 */
   if ncol(idx)>0 then 
      w[idx] = 0;
   return w-1;
finish;

/* HELPER FUNCTIONS FOR POSITIVE BRANCH */

/* Boyd's first approximation (Eqn 5) */
start Wp_Boyd0(y);
   lambda = 0.96299820407774; /* = sqrt(2) / (log(10)-log(log(10))) */
   Logy10 = log(y+10);
   w0 = (logy10 - log(logy10)) # tanh(lambda*sqrt(y));
   return w0;
finish;
/* Boyd's modified approximation (Eqn 7) */
start Wp_Boyd1(y);
   Logym14 = log(y)-1.4;
   w1 = Wp_Boyd0(y) # (1 + 0.1*logym14 # exp(-0.075*logym14##2));
   return w1;
finish;

/* Halley's method applied to f(x) = (x-1)#exp(x) - y + 1,  y>0 */
/* http://blogs.sas.com/content/iml/2016/08/24/halleys-method-finding-roots.html */
start HalleyRefineWp(y, z, numIters=2);
   x = z;                      /* initial guess */
   do i = 1 to numIters;
      fx = (x-1) # exp(x) - y + 1;     /* f(x)   */
      dfdx = x # exp(x);               /* f'(x)  */
      ratio = (x+1) / x;               /* f''(x) / f'(x) */
      x = x - fx / (dfdx - 0.5*fx#ratio);
   end;
   return x;
finish;

/************************************************/
/*  Evaluation of lower branch Wm(x)            */
/************************************************/

/* MAIN FUNCTION TO CALL FOR NEGATIVE BRANCH */
/* negative branch of Lambert W function. Domain: [-1/e, 0) */
start LambertWm(x, HalleyIters=1, CheckBounds=1); 
   if ^CheckBounds then do;
      v = Wm_CBM(x);
      return HalleyRefineWm(x, v, HalleyIters);
   end;

   /* otherwise, check domain. compute x=-1/e separately */
   w = j(nrow(x), ncol(x), .);
   idx = loc(x > -1/exp(1) & x < 0);   /* for -1/e < x < 0 */
   if ncol(idx)>0 then do;
      y = x[idx];
      v = Wm_CBM(y);
      w[idx] = HalleyRefineWm(y, v, HalleyIters);
   end;
   if ncol(idx)=nrow(x)*ncol(x) then   /* all found; return */
      return w;

   idx = loc(x = -1/exp(1));           /* if any x = -1/e */
   if ncol(idx)>0 then 
      w[idx] = -1;
   return w;
finish;

/* HELPER FUNCTIONS FOR NEGATIVE BRANCH */
/* http://blogs.sas.com/content/iml/2011/09/14/evaluate-polynomials-efficiently-by-using-horners-scheme.html */
start Polyval(coef, v);
/* Horner's evaluation of a polynomial of deg=n evaluated at x. 
   Input: coef: a column vector of length n+1 containing
                the coefficients in descending order.
          v: column vector of points at which to evaluate poly */
   c = colvec(coef); x = colvec(v); /* make column vectors */
   y = j(nrow(x), 1, c[1]);         /* initialize to c[1] */
   do j = 2 to nrow(c); 
      y = y # x + c[j];
   end;
   return(y);
finish;

/* Approximation to the negative branch of the Lambert W function. 
   Algorithm from p. 2163 of
   F. Chapeau-Blondeau and A. Monir (2002), "Numerical Evaluation of 
   the Lambert W Function and Application to Generation of Generalized
   Gaussian Noise With Exponent 1/2", IEEE Trans. on Signal Processing 

   If you know you are going to follow the approximation with Halley iterations,
   you can compute with less precision than CBM do. This implementation 
   simplifies the computation on R and L intervals.
*/
start Wm_CBM(x);
   /* handle function on Left, Middle, and Right disjoint domains */
   y = j(nrow(x), ncol(x), .);
   /* series expansion on L */
   Lidx = loc( -1/exp(1) <= x & x < -0.333 );
   if ncol(Lidx)>0 then do;
      p = -sqrt(2*(exp(1)#x[Lidx] + 1));
      coef = {-0.0259847148736038, 0.0445023148148148, -0.07962962962962, 
               0.15277777777777, -0.33333333333333, 1, -1};
      /* -221/8505 p^6 +769/17280 p^5 -43/540 p^4 +11/72 p^3 -1/3 p^2 +1 p -1 */
      *coef = { 0.0156356325323339,
        -0.0259847148736038, 0.0445023148148148, -0.07962962962962, 
         0.15277777777777, -0.33333333333333, 1, -1};
      y[Lidx] = Polyval(coef, p);
   end;
   /* rational function approximation M */
   Midx = loc(-0.333 <= x & x <= -0.033);
   if ncol(Midx)>0 then do;
      z = x[Midx];
      numCoef = {-5532.7760, -4877.6330, -47.4252, 391.0025, -8.0960};
      denomCoef= {1515.3060, 433.8688, -82.9423, 1};
      y[Midx] = Polyval(numCoef, z) / Polyval(denomCoef, z);
   end;
   /* Asymptotic expansion on R */
   Ridx = loc(-0.033 < x & x < 0);
   if ncol(Ridx)>0 then do;
      z = x[Ridx];
      L1 = log(-z);
      L2 = log( -log(-z) );
      L1_2 = L1#L1;
      L1_3 = L1#L1_2;
      y[Ridx] = L1 - L2 + L2/L1 + 
                (-2 + L2)#L2 / (2*L1_2) + 
                Polyval({2,-9,6}, L2) # L2 / (6*L1_3) 
                /* + 
                Polyval({3,-22,36,-12}, L2) # L2 / (12*L1_2#L1_2)+
                Polyval({12,-125,350,-300,60}, L2) # L2 / (60*L1_2#L1_3)*/;
   end;
   return y;
finish;

/* Halley's method applied to f(x) = x#exp(x) - y,  where x>-1 */
/* http://blogs.sas.com/content/iml/2016/08/24/halleys-method-finding-roots.html */
start HalleyRefineWm(y, z, numIters=2);
   x = z;
   do i = 1 to numIters;
      fx = x # exp(x) - y;     /* f(x)   */
      dfdx = (1+x) # exp(x);   /* f'(x)  */
      ratio = (2+x) / (1+x);   /* f''(x) / f'(x) */
      x = x - fx / (dfdx - 0.5*fx#ratio);
   end;
   return x;
finish;

store module=_all_;
quit;