# Lambert W function in Maxima

Oct 24th, 2018
294
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
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)))
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)))
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)))))