lambertw, all branches

[Albert Chan]

Python W code (post 16)

Quote:s = (im(a) < 0)
small_imag = k==1 and s or k==-1 and not s
...
s = 1 - 2 * (k<0 or k==0 and s)

It is very convoluted. Let's make a table

Code:

old s       updated s
im(a)<0  k>0   k<0   k=0                   
True     +1    -1    -1
False    +1    -1    +1

It is simpler to have old s = - sign(im(a) or 1) = ±1
Then, small_imag = (k==s), too simple to even need a variable for it.

Previous post customized LN and ability to enter W guess also included.

Python W, version 4

Code:

from mpmath import *    # no signed zero support

r, err = 0.36787944117144233, -1.2428753672788363E-17

def W(a, k=0, verbal=False):
    h, s = a+r+err, -sign(im(a) or 1)
    if abs(h)<.25 and (k==0 or k==s):
        y = sqrt(2*r*h) * (1-2*k*k) # close to branch point
        y = r + y*sqrt(1+y/(3*r))   # with small imag part
        while 1:
            if verbal: print a/y
            x = y
            y = (y+a)/log1p((y-r-err)/r)
            if abs((y-x)/y) < 1e-9: return a/y
    def LN(x, bad = -sign(k) or s):
        x = ln(x)
        return conj(x) if sign(im(x))==bad else x
    A, T = abs(a), mpc(0,2*k*pi+arg(a))
    x = T          # W rough guess
    if k==0: x = (A+100)/(A+50); x = log1p(a*x)/x
    if k==s: x = LN(-a) + (0 if A<.5 else T/2)
    T = T + log(A) # = ln_k(a)
    if not isfinite(A): return T
    if A==0: return a if k==0 else T - sign(k)*pi*j
    if verbal and verbal!=True: x = verbal
    while 1:
        if verbal: print x
        h = x/(x+1) * (x + LN(x) - T)
        x = x - h
        if not (abs(h/x) >= 1e-6): break
    if verbal: print x
    h = (x-a*exp(-x))/(x+1) # final touch
    return x-h if isfinite(h) else x