Lambert W function in Maxima

(defun $lambert_w (z)
  (simplify (list '(%lambert_w) (resimplify z))))

;;; Set properties to give full support to the parser and display
(defprop $lambert_w %lambert_w alias)
(defprop $lambert_w %lambert_w verb)
(defprop %lambert_w $lambert_w reversealias)
(defprop %lambert_w $lambert_w noun)

;;; lambert_w is a simplifying function
(defprop %lambert_w simp-lambertw operators)

;;; Derivative of lambert_w
(defprop %lambert_w
  ((x)
   ((mtimes)
    ((mexpt) $%e ((mtimes ) -1 ((%lambert_w) x)))
    ((mexpt) ((mplus) 1 ((%lambert_w) x)) -1)))
  grad)

;;; Integral of lambert_w
;;; integrate(W(x),x) := x*(W(x)^2-W(x)+1)/W(x)
(defprop %lambert_w
  ((x)
   ((mtimes)
    x
    ((mplus)
     ((mexpt) ((%lambert_w) x) 2)
     ((mtimes) -1 ((%lambert_w) x))
     1)
    ((mexpt) ((%lambert_w) x) -1)))
  integral)

(defun simp-lambertw (x yy z)
  (declare (ignore yy))
  (oneargcheck x)
  (setq x (simpcheck (cadr x) z))
  (cond ((equal x 0) 0)
        ((equal x 0.0) 0.0)
        ((zerop1 x) ($bfloat 0))        ;bfloat case
        ((alike1 x '$%e)
         ;; W(%e) = 1
         1)
        ((alike1 x '((mtimes simp) ((rat simp) -1 2) ((%log simp) 2)))
         ;; W(-log(2)/2) = -log(2)
         '((mtimes simp) -1 ((%log simp) 2)))
        ((alike1 x '((mtimes simp) -1 ((mexpt simp) $%e -1)))
         ;; W(-1/e) = -1
         -1)
        ((alike1 x '((mtimes) ((rat) -1 2) $%pi))
         ;; W(-%pi/2) = %i*%pi/2
         '((mtimes simp) ((rat simp) 1 2) $%i $%pi))
        ;; numerical evaluation
        ((complex-float-numerical-eval-p x)
          ;; x may be an integer.  eg "lambert_w(1),numer;"
          (if (integerp x)
            (to (bigfloat::lambert-w-k 0 (bigfloat:to ($float x))))
            (to (bigfloat::lambert-w-k 0 (bigfloat:to x)))))
        ((complex-bigfloat-numerical-eval-p x)
         (to (bigfloat::lambert-w-k 0 (bigfloat:to x))))
        (t (list '(%lambert_w simp) x))))

;; An approximation of the k-branch of generalized Lambert W function
;;   k integer
;;   z real or complex lisp float
;; Used as initial guess for Halley's iteration.
;; When W(z) is real, ensure that guess is real.
(defun init-lambert-w-k (k z)
  (let ( ; parameters for k = +/- 1 near branch pont z=-1/%e
        (branch-eps 0.2e0)
        (branch-point (/ -1 %e-val))) ; branch pont z=-1/%e
    (cond
      ; For principal branch k=0, use expression by Winitzki
      ((= k 0) (init-lambert-w-0 z))
      ; For k=1 branch, near branch point z=-1/%e with im(z) <  0
      ((and (= k 1)
            (< (imagpart z) 0)
            (< (abs (- branch-point z)) branch-eps))
        (bigfloat::lambert-branch-approx z))
      ; Default to asymptotic expansion Corless et al (4.20)
      ; W_k(z) = log(z) + 2.pi.i.k - log(log(z)+2.pi.i.k)
      (t (let ((two-pi-i-k (complex 0.0e0 (* 2 pi k))))
                 (+ (log z)
                    two-pi-i-k
                    (* -1 (log (+ (log z) two-pi-i-k )))))))))

;; Complex value of the principal branch of Lambert's W function in
;; the entire complex plane with relative error less than 1%, given
;; standard branch cuts for sqrt(z) and log(z).
;; Winitzki (2003)
(defun init-lambert-w-0 (z)
  (let ((a 2.344e0) (b 0.8842e0) (c 0.9294e0) (d 0.5106e0) (e -1.213e0)
     (y (sqrt (+ (* 2 %e-val z ) 2)) ) )   ; y=sqrt(2*%e*z+2)
    ; w = (2*log(1+B*y)-log(1+C*log(1+D*y))+E)/(1+1/(2*log(1+B*y)+2*A)
     (/
      (+ (* 2 (log (+ 1 (* b y))))
         (* -1 (log (+ 1 (* c (log (+ 1 (* d y)))))))
         e)
      (+ 1
         (/ 1 (+ (* 2 (log (+ 1 (* b y)))) (* 2 a)))))))

;; Approximate k=-1 branch of Lambert's W function over -1/e < z < 0.
;; W(z) is real, so we ensure the starting guess for Halley iteration
;; is also real.
;; Verebic (2010)
(defun init-lambert-w-minus1 (z)
  (cond
    ((not (realp z))
      (merror "z not real in init-lambert-w-minus1"))
    ((or (< z (/ -1 %e-val)) (plusp z))
      (merror "z outside range of approximation in init-lambert-w-minus1"))
    ;; In the region where W(z) is real
    ;; -1/e < z < C, use power series about branch point -1/e ~ -0.36787
    ;; C = -0.3 seems a reasonable crossover point
    ((< z -0.3)
      (bigfloat::lambert-branch-approx z))
    ;; otherwise C <= z < 0, use iteration W(z) ~ ln(-z)-ln(-W(z))
    ;; nine iterations are sufficient over -0.3 <= z < 0
    (t (let* ((ln-z (log (- z))) (maxiter 9) (w ln-z))
         (dotimes (k maxiter w)
            (setq w (- ln-z (log (- w)))))))))

(in-package #-gcl #:bigfloat #+gcl "BIGFLOAT")

;; Approximate Lambert W(k,z) for k=1 and k=-1 near branch point z=-1/%e
;; using power series in y=-sqrt(2*%e*z+2)
;;   for im(z) < 0,  approximates k=1 branch
;;   for im(z) >= 0, approximates k=-1  branch
;;
;; Corless et al (1996) (4.22)
;; Verebic (2010)
;;
;; z is a real or complex bigfloat:
(defun lambert-branch-approx (z)
  (let ((y (- (sqrt (+ (* 2 (%e z) z ) 2)))) ; y=-sqrt(2*%e*z+2)
    (b0 -1) (b1 1) (b2 -1/3) (b3 11/72))
    (+ b0 (* y (+ b1 (* y (+ b2 (* b3 y))))))))

;; Algorithm based in part on Corless et al (1996).
;;
;; It is Halley's iteration applied to w*exp(w).
;;
;;
;;                               w[j] exp(w[j]) - z
;; w[j+1] = w[j] - -------------------------------------------------
;;                                       (w[j]+2)(w[j] exp(w[j]) -z)
;;                  exp(w[j])(w[j]+1) -  ---------------------------
;;                                               2 w[j] + 2
;;
;; The algorithm has cubic convergence.  Once convergence begins, the
;; number of digits correct at step k is roughly 3 times the number
;; which were correct at step k-1.
;;
;; Convergence can stall using convergence test abs(w[j+1]-w[j]) < prec,
;; as happens for generalized_lambert_w(-1,z) near branch point z = -1/%e
;; Therefore also stop iterating if abs(w[j]*exp(w[j]) - z) << abs(z)
(defun lambert-w-k (k z &key (maxiter 50))
  (let ((w (init-lambert-w-k k z)) we w1e delta (prec (* 4 (epsilon z))))
    (dotimes (i maxiter (maxima::merror "lambert-w-k did not converge"))
      (setq we (* w (exp w)))
      (when (<= (abs (- z we)) (* 4 (epsilon z) (abs z))) (return))
      (setq w1e (* (1+ w) (exp w)))
      (setq delta (/ (- we z)
                     (- w1e (/ (* (+ w 2) (- we z)) (+ 2 (* 2 w))))))
      (decf w delta)
      (when (<= (abs (/ delta w)) prec) (return)))
    ;; Check iteration converged to correct branch
    (check-lambert-w-k k w z)
    w))

(defmethod init-lambert-w-k ((k integer) (z number))
  (maxima::init-lambert-w-k k z))

(defmethod init-lambert-w-k ((k integer) (z bigfloat))
  (bfloat-init-lambert-w-k k z))

(defmethod init-lambert-w-k ((k integer) (z complex-bigfloat))
  (bfloat-init-lambert-w-k k z))

(defun bfloat-init-lambert-w-k (k z)
  "Approximate generalized_lambert_w(k,z) for bigfloat: z as initial guess"
  (let ((branch-point -0.36787944117144)) ; branch point -1/%e
    (cond
       ;; if k=-1, z very close to -1/%e and imag(z)>=0, use power series
       ((and (= k -1)
             (or (zerop (imagpart z))
                 (plusp (imagpart z)))
             (< (abs (- z branch-point)) 1e-10))
         (lambert-branch-approx z))
       ;; if k=1, z very close to -1/%e and imag(z)<0, use power series
       ((and (= k 1)
             (minusp (imagpart z))
             (< (abs (- z branch-point)) 1e-10))
         (lambert-branch-approx z))
       ;; Initialize using float value if z is representable as a float
       ((< (abs z) 1.0e100)
         (if (complexp z)
             (bigfloat (lambert-w-k k (cl:complex (float (realpart z) 1.0)
                                                  (float (imagpart z) 1.0))))
             (bigfloat (lambert-w-k k (float z 1.0)))))
       ;; For large z, use Corless et al (4.20)
       ;;              W_k(z) ~ log(z) + 2.pi.i.k - log(log(z)+2.pi.i.k)
       (t
        (let ((log-z (log z)))
          (if (= k 0)
            (- log-z (log log-z))
            (let* ((i (make-instance 'complex-bigfloat :imag (intofp 1)))
                  (two-pi-i-k (* 2 (%pi z) i k)))
              (- (+ log-z two-pi-i-k)
                 (log (+ log-z two-pi-i-k))))))))))


;; Check Lambert W iteration converged to the correct branch
;; W_k(z) + ln W_k(z) = ln z, for k = -1 and z in [-1/e,0)
;;                    = ln z + 2 pi i k, otherwise
;; See Jeffrey, Hare, Corless (1996), eq (12)
;; k integer
;; z, w bigfloat: numbers
(defun check-lambert-w-k (k w z)
  (let ((tolerance #-gcl 1.0e-6
                   #+gcl (cl:float 1/1000000)))
  (if
     (cond
       ;; k=-1 branch with z and w real.
      ((and (= k -1) (realp z) (minusp z) (>= z (/ -1 (%e z))))
       (if (and (realp w)
                (<= w -1)
                (< (abs (+ w (log w) (- (log z)))) tolerance))
           t
           nil))
       (t
         ; i k =  (W_k(z) + ln W_k(z) - ln(z)) / 2 pi
        (let (ik)
          (setq ik (/ (+ w (log w) (- (log z))) (* 2 (%pi z))))
          (if (and (< (realpart ik) tolerance)
                   (< (abs (- k (imagpart ik))) tolerance))
            t
            nil))))
      t
      (maxima::merror "Lambert W iteration converged to wrong branch"))))

(in-package :maxima)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;; Generalized Lambert W function
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defun $generalized_lambert_w (k z)
  (simplify (list '(%generalized_lambert_w) (resimplify k) (resimplify z))))

;;; Set properties to give full support to the parser and display
(defprop $generalized_lambert_w %generalized_lambert_w alias)
(defprop $generalized_lambert_w %generalized_lambert_w verb)
(defprop %generalized_lambert_w $generalized_lambert_w reversealias)
(defprop %generalized_lambert_w $generalized_lambert_w noun)

;;; lambert_w is a simplifying function
(defprop %generalized_lambert_w simp-generalized-lambertw operators)

;;; Derivative of lambert_w
(defprop %generalized_lambert_w
  ((k x)
   nil
   ((mtimes)
    ((mexpt) $%e ((mtimes ) -1 ((%generalized_lambert_w) k x)))
    ((mexpt) ((mplus) 1 ((%generalized_lambert_w) k x)) -1)))
  grad)

;;; Integral of lambert_w
;;; integrate(W(k,x),x) := x*(W(k,x)^2-W(k,x)+1)/W(k,x)
(defprop %generalized_lambert_w
  ((k x)
   nil
   ((mtimes)
    x
    ((mplus)
     ((mexpt) ((%generalized_lambert_w) k x) 2)
     ((mtimes) -1 ((%generalized_lambert_w) k x))
     1)
    ((mexpt) ((%generalized_lambert_w) k x) -1)))
  integral)

(defun simp-generalized-lambertw (expr ignored z)
  (declare (ignore ignored))
  (twoargcheck expr)
  (let ((k (simpcheck (cadr expr) z))
        (x (simpcheck (caddr expr) z)))
    (cond
     ;; Numerical evaluation for real or complex x
     ((and (integerp k) (complex-float-numerical-eval-p x))
       ;; x may be an integer.  eg "generalized_lambert_w(0,1),numer;"
       (if (integerp x)
           (to (bigfloat::lambert-w-k k (bigfloat:to ($float x))))
           (to (bigfloat::lambert-w-k k (bigfloat:to x)))))
     ;; Numerical evaluation for real or complex bigfloat x
     ((and (integerp k) (complex-bigfloat-numerical-eval-p x))
      (to (bigfloat::lambert-w-k k (bigfloat:to x))))
     (t (list '(%generalized_lambert_w simp) k x)))))