Advertisement
Bisqwit

Lambert W function in Maxima

Oct 24th, 2018
294
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Lisp 11.20 KB | None | 0 0
  1. (defun $lambert_w (z)
  2.   (simplify (list '(%lambert_w) (resimplify z))))
  3.  
  4. ;;; Set properties to give full support to the parser and display
  5. (defprop $lambert_w %lambert_w alias)
  6. (defprop $lambert_w %lambert_w verb)
  7. (defprop %lambert_w $lambert_w reversealias)
  8. (defprop %lambert_w $lambert_w noun)
  9.  
  10. ;;; lambert_w is a simplifying function
  11. (defprop %lambert_w simp-lambertw operators)
  12.  
  13. ;;; Derivative of lambert_w
  14. (defprop %lambert_w
  15.   ((x)
  16.    ((mtimes)
  17.     ((mexpt) $%e ((mtimes ) -1 ((%lambert_w) x)))
  18.     ((mexpt) ((mplus) 1 ((%lambert_w) x)) -1)))
  19.   grad)
  20.  
  21. ;;; Integral of lambert_w
  22. ;;; integrate(W(x),x) := x*(W(x)^2-W(x)+1)/W(x)
  23. (defprop %lambert_w
  24.   ((x)
  25.    ((mtimes)
  26.     x
  27.     ((mplus)
  28.      ((mexpt) ((%lambert_w) x) 2)
  29.      ((mtimes) -1 ((%lambert_w) x))
  30.      1)
  31.     ((mexpt) ((%lambert_w) x) -1)))
  32.   integral)
  33.  
  34. (defun simp-lambertw (x yy z)
  35.   (declare (ignore yy))
  36.   (oneargcheck x)
  37.   (setq x (simpcheck (cadr x) z))
  38.   (cond ((equal x 0) 0)
  39.         ((equal x 0.0) 0.0)
  40.         ((zerop1 x) ($bfloat 0))        ;bfloat case
  41.         ((alike1 x '$%e)
  42.          ;; W(%e) = 1
  43.          1)
  44.         ((alike1 x '((mtimes simp) ((rat simp) -1 2) ((%log simp) 2)))
  45.          ;; W(-log(2)/2) = -log(2)
  46.          '((mtimes simp) -1 ((%log simp) 2)))
  47.         ((alike1 x '((mtimes simp) -1 ((mexpt simp) $%e -1)))
  48.          ;; W(-1/e) = -1
  49.          -1)
  50.         ((alike1 x '((mtimes) ((rat) -1 2) $%pi))
  51.          ;; W(-%pi/2) = %i*%pi/2
  52.          '((mtimes simp) ((rat simp) 1 2) $%i $%pi))
  53.         ;; numerical evaluation
  54.         ((complex-float-numerical-eval-p x)
  55.           ;; x may be an integer.  eg "lambert_w(1),numer;"
  56.           (if (integerp x)
  57.             (to (bigfloat::lambert-w-k 0 (bigfloat:to ($float x))))
  58.             (to (bigfloat::lambert-w-k 0 (bigfloat:to x)))))
  59.         ((complex-bigfloat-numerical-eval-p x)
  60.          (to (bigfloat::lambert-w-k 0 (bigfloat:to x))))
  61.         (t (list '(%lambert_w simp) x))))
  62.  
  63. ;; An approximation of the k-branch of generalized Lambert W function
  64. ;;   k integer
  65. ;;   z real or complex lisp float
  66. ;; Used as initial guess for Halley's iteration.
  67. ;; When W(z) is real, ensure that guess is real.
  68. (defun init-lambert-w-k (k z)
  69.   (let ( ; parameters for k = +/- 1 near branch pont z=-1/%e
  70.         (branch-eps 0.2e0)
  71.         (branch-point (/ -1 %e-val))) ; branch pont z=-1/%e
  72.     (cond
  73.       ; For principal branch k=0, use expression by Winitzki
  74.       ((= k 0) (init-lambert-w-0 z))
  75.       ; For k=1 branch, near branch point z=-1/%e with im(z) <  0
  76.       ((and (= k 1)
  77.             (< (imagpart z) 0)
  78.             (< (abs (- branch-point z)) branch-eps))
  79.         (bigfloat::lambert-branch-approx z))
  80.       ; Default to asymptotic expansion Corless et al (4.20)
  81.       ; W_k(z) = log(z) + 2.pi.i.k - log(log(z)+2.pi.i.k)
  82.       (t (let ((two-pi-i-k (complex 0.0e0 (* 2 pi k))))
  83.                  (+ (log z)
  84.                     two-pi-i-k
  85.                     (* -1 (log (+ (log z) two-pi-i-k )))))))))
  86.  
  87. ;; Complex value of the principal branch of Lambert's W function in
  88. ;; the entire complex plane with relative error less than 1%, given
  89. ;; standard branch cuts for sqrt(z) and log(z).
  90. ;; Winitzki (2003)
  91. (defun init-lambert-w-0 (z)
  92.   (let ((a 2.344e0) (b 0.8842e0) (c 0.9294e0) (d 0.5106e0) (e -1.213e0)
  93.      (y (sqrt (+ (* 2 %e-val z ) 2)) ) )   ; y=sqrt(2*%e*z+2)
  94.     ; w = (2*log(1+B*y)-log(1+C*log(1+D*y))+E)/(1+1/(2*log(1+B*y)+2*A)
  95.      (/
  96.       (+ (* 2 (log (+ 1 (* b y))))
  97.          (* -1 (log (+ 1 (* c (log (+ 1 (* d y)))))))
  98.          e)
  99.       (+ 1
  100.          (/ 1 (+ (* 2 (log (+ 1 (* b y)))) (* 2 a)))))))
  101.  
  102. ;; Approximate k=-1 branch of Lambert's W function over -1/e < z < 0.
  103. ;; W(z) is real, so we ensure the starting guess for Halley iteration
  104. ;; is also real.
  105. ;; Verebic (2010)
  106. (defun init-lambert-w-minus1 (z)
  107.   (cond
  108.     ((not (realp z))
  109.       (merror "z not real in init-lambert-w-minus1"))
  110.     ((or (< z (/ -1 %e-val)) (plusp z))
  111.       (merror "z outside range of approximation in init-lambert-w-minus1"))
  112.     ;; In the region where W(z) is real
  113.     ;; -1/e < z < C, use power series about branch point -1/e ~ -0.36787
  114.     ;; C = -0.3 seems a reasonable crossover point
  115.     ((< z -0.3)
  116.       (bigfloat::lambert-branch-approx z))
  117.     ;; otherwise C <= z < 0, use iteration W(z) ~ ln(-z)-ln(-W(z))
  118.     ;; nine iterations are sufficient over -0.3 <= z < 0
  119.     (t (let* ((ln-z (log (- z))) (maxiter 9) (w ln-z))
  120.          (dotimes (k maxiter w)
  121.             (setq w (- ln-z (log (- w)))))))))
  122.  
  123. (in-package #-gcl #:bigfloat #+gcl "BIGFLOAT")
  124.  
  125. ;; Approximate Lambert W(k,z) for k=1 and k=-1 near branch point z=-1/%e
  126. ;; using power series in y=-sqrt(2*%e*z+2)
  127. ;;   for im(z) < 0,  approximates k=1 branch
  128. ;;   for im(z) >= 0, approximates k=-1  branch
  129. ;;
  130. ;; Corless et al (1996) (4.22)
  131. ;; Verebic (2010)
  132. ;;
  133. ;; z is a real or complex bigfloat:
  134. (defun lambert-branch-approx (z)
  135.   (let ((y (- (sqrt (+ (* 2 (%e z) z ) 2)))) ; y=-sqrt(2*%e*z+2)
  136.     (b0 -1) (b1 1) (b2 -1/3) (b3 11/72))
  137.     (+ b0 (* y (+ b1 (* y (+ b2 (* b3 y))))))))
  138.  
  139. ;; Algorithm based in part on Corless et al (1996).
  140. ;;
  141. ;; It is Halley's iteration applied to w*exp(w).
  142. ;;
  143. ;;
  144. ;;                               w[j] exp(w[j]) - z
  145. ;; w[j+1] = w[j] - -------------------------------------------------
  146. ;;                                       (w[j]+2)(w[j] exp(w[j]) -z)
  147. ;;                  exp(w[j])(w[j]+1) -  ---------------------------
  148. ;;                                               2 w[j] + 2
  149. ;;
  150. ;; The algorithm has cubic convergence.  Once convergence begins, the
  151. ;; number of digits correct at step k is roughly 3 times the number
  152. ;; which were correct at step k-1.
  153. ;;
  154. ;; Convergence can stall using convergence test abs(w[j+1]-w[j]) < prec,
  155. ;; as happens for generalized_lambert_w(-1,z) near branch point z = -1/%e
  156. ;; Therefore also stop iterating if abs(w[j]*exp(w[j]) - z) << abs(z)
  157. (defun lambert-w-k (k z &key (maxiter 50))
  158.   (let ((w (init-lambert-w-k k z)) we w1e delta (prec (* 4 (epsilon z))))
  159.     (dotimes (i maxiter (maxima::merror "lambert-w-k did not converge"))
  160.       (setq we (* w (exp w)))
  161.       (when (<= (abs (- z we)) (* 4 (epsilon z) (abs z))) (return))
  162.       (setq w1e (* (1+ w) (exp w)))
  163.       (setq delta (/ (- we z)
  164.                      (- w1e (/ (* (+ w 2) (- we z)) (+ 2 (* 2 w))))))
  165.       (decf w delta)
  166.       (when (<= (abs (/ delta w)) prec) (return)))
  167.     ;; Check iteration converged to correct branch
  168.     (check-lambert-w-k k w z)
  169.     w))
  170.  
  171. (defmethod init-lambert-w-k ((k integer) (z number))
  172.   (maxima::init-lambert-w-k k z))
  173.  
  174. (defmethod init-lambert-w-k ((k integer) (z bigfloat))
  175.   (bfloat-init-lambert-w-k k z))
  176.  
  177. (defmethod init-lambert-w-k ((k integer) (z complex-bigfloat))
  178.   (bfloat-init-lambert-w-k k z))
  179.  
  180. (defun bfloat-init-lambert-w-k (k z)
  181.   "Approximate generalized_lambert_w(k,z) for bigfloat: z as initial guess"
  182.   (let ((branch-point -0.36787944117144)) ; branch point -1/%e
  183.     (cond
  184.        ;; if k=-1, z very close to -1/%e and imag(z)>=0, use power series
  185.        ((and (= k -1)
  186.              (or (zerop (imagpart z))
  187.                  (plusp (imagpart z)))
  188.              (< (abs (- z branch-point)) 1e-10))
  189.          (lambert-branch-approx z))
  190.        ;; if k=1, z very close to -1/%e and imag(z)<0, use power series
  191.        ((and (= k 1)
  192.              (minusp (imagpart z))
  193.              (< (abs (- z branch-point)) 1e-10))
  194.          (lambert-branch-approx z))
  195.        ;; Initialize using float value if z is representable as a float
  196.        ((< (abs z) 1.0e100)
  197.          (if (complexp z)
  198.              (bigfloat (lambert-w-k k (cl:complex (float (realpart z) 1.0)
  199.                                                   (float (imagpart z) 1.0))))
  200.              (bigfloat (lambert-w-k k (float z 1.0)))))
  201.        ;; For large z, use Corless et al (4.20)
  202.        ;;              W_k(z) ~ log(z) + 2.pi.i.k - log(log(z)+2.pi.i.k)
  203.        (t
  204.         (let ((log-z (log z)))
  205.           (if (= k 0)
  206.             (- log-z (log log-z))
  207.             (let* ((i (make-instance 'complex-bigfloat :imag (intofp 1)))
  208.                   (two-pi-i-k (* 2 (%pi z) i k)))
  209.               (- (+ log-z two-pi-i-k)
  210.                  (log (+ log-z two-pi-i-k))))))))))
  211.  
  212.  
  213. ;; Check Lambert W iteration converged to the correct branch
  214. ;; W_k(z) + ln W_k(z) = ln z, for k = -1 and z in [-1/e,0)
  215. ;;                    = ln z + 2 pi i k, otherwise
  216. ;; See Jeffrey, Hare, Corless (1996), eq (12)
  217. ;; k integer
  218. ;; z, w bigfloat: numbers
  219. (defun check-lambert-w-k (k w z)
  220.   (let ((tolerance #-gcl 1.0e-6
  221.                    #+gcl (cl:float 1/1000000)))
  222.   (if
  223.      (cond
  224.        ;; k=-1 branch with z and w real.
  225.       ((and (= k -1) (realp z) (minusp z) (>= z (/ -1 (%e z))))
  226.        (if (and (realp w)
  227.                 (<= w -1)
  228.                 (< (abs (+ w (log w) (- (log z)))) tolerance))
  229.            t
  230.            nil))
  231.        (t
  232.          ; i k =  (W_k(z) + ln W_k(z) - ln(z)) / 2 pi
  233.         (let (ik)
  234.           (setq ik (/ (+ w (log w) (- (log z))) (* 2 (%pi z))))
  235.           (if (and (< (realpart ik) tolerance)
  236.                    (< (abs (- k (imagpart ik))) tolerance))
  237.             t
  238.             nil))))
  239.       t
  240.       (maxima::merror "Lambert W iteration converged to wrong branch"))))
  241.  
  242. (in-package :maxima)
  243.  
  244. ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  245. ;;; Generalized Lambert W function
  246. ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  247.  
  248. (defun $generalized_lambert_w (k z)
  249.   (simplify (list '(%generalized_lambert_w) (resimplify k) (resimplify z))))
  250.  
  251. ;;; Set properties to give full support to the parser and display
  252. (defprop $generalized_lambert_w %generalized_lambert_w alias)
  253. (defprop $generalized_lambert_w %generalized_lambert_w verb)
  254. (defprop %generalized_lambert_w $generalized_lambert_w reversealias)
  255. (defprop %generalized_lambert_w $generalized_lambert_w noun)
  256.  
  257. ;;; lambert_w is a simplifying function
  258. (defprop %generalized_lambert_w simp-generalized-lambertw operators)
  259.  
  260. ;;; Derivative of lambert_w
  261. (defprop %generalized_lambert_w
  262.   ((k x)
  263.    nil
  264.    ((mtimes)
  265.     ((mexpt) $%e ((mtimes ) -1 ((%generalized_lambert_w) k x)))
  266.     ((mexpt) ((mplus) 1 ((%generalized_lambert_w) k x)) -1)))
  267.   grad)
  268.  
  269. ;;; Integral of lambert_w
  270. ;;; integrate(W(k,x),x) := x*(W(k,x)^2-W(k,x)+1)/W(k,x)
  271. (defprop %generalized_lambert_w
  272.   ((k x)
  273.    nil
  274.    ((mtimes)
  275.     x
  276.     ((mplus)
  277.      ((mexpt) ((%generalized_lambert_w) k x) 2)
  278.      ((mtimes) -1 ((%generalized_lambert_w) k x))
  279.      1)
  280.     ((mexpt) ((%generalized_lambert_w) k x) -1)))
  281.   integral)
  282.  
  283. (defun simp-generalized-lambertw (expr ignored z)
  284.   (declare (ignore ignored))
  285.   (twoargcheck expr)
  286.   (let ((k (simpcheck (cadr expr) z))
  287.         (x (simpcheck (caddr expr) z)))
  288.     (cond
  289.      ;; Numerical evaluation for real or complex x
  290.      ((and (integerp k) (complex-float-numerical-eval-p x))
  291.        ;; x may be an integer.  eg "generalized_lambert_w(0,1),numer;"
  292.        (if (integerp x)
  293.            (to (bigfloat::lambert-w-k k (bigfloat:to ($float x))))
  294.            (to (bigfloat::lambert-w-k k (bigfloat:to x)))))
  295.      ;; Numerical evaluation for real or complex bigfloat x
  296.      ((and (integerp k) (complex-bigfloat-numerical-eval-p x))
  297.       (to (bigfloat::lambert-w-k k (bigfloat:to x))))
  298.      (t (list '(%generalized_lambert_w simp) k x)))))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement