lambertw, all branches

[Albert Chan]

Here is equivalent complex.W version 2, adapted to machine without signed zero.

Without signed zero, this trick is out: W(a, k) = conj(W(conj(a), -k)
Instead, let W(a, k) = x = re(z) + s*im(z)*I , such that s = ±1, arg(z) ≥ 0

s = sign(arg(x)) = ±1            |x| = |z|            arg(x) = s * arg(z)

lua> I.W(-1e-100, -1, true)
(-230.25850929940458-0*I)
(-235.72143712473638-0*I)
(-235.72115887568603-0*I)
(-235.72115887568532-0*I)
(-235.72115887568535-0*I)      true

Above example, s=-1, all x iterations, arg(x) = s * pi
Almost identical, for efficiency: arg(x) ≠ bad = s * -pi

In other words, iteration overshoot allowed, where arg(z) ≥ 0 does not hold.
But, not to the extreme opposite direction, i.e. arg(z) ≠ -pi

We need to add arg(x) edge case, to make W (without signed zero) work the same as I.W

Code:

        h = arg(x)
        if h == bad: h = -h

I use Python mpmath, which does not have signed zero.

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, im(a)<0
    small_imag = k==1 and s or k==-1 and not s
    if abs(h)<.25 and (k==0 or small_imag):
        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
            h = y - (y+a) / log1p((y-r-err)/r)
            y = y - h
            if y == y+h*1e-4: return a/y, True
    s = 1 - 2 * (k<0 or k==0 and s)
    A, T, bad = abs(a), (2*k*pi+arg(a))*j, -s*pi
    x = T   # W rough guess
    if small_imag and A<.6: x = log(-a)
    if k == 0: x = T/2 if abs(a+1)<.6 else log1p(a)
    if s*im(x) < 0: x = conj(x)     # s = sign(im(x))
    k = 40  # max loops
    while 1:
        if verbal: print x
        h = arg(x)
        if h == bad: h = -h
        h = x/(x+1) * (x + mpc(log(abs(x)/A),h) - T)
        x, k = x-h, k-1
        if k==0 or x==x+h*1e-8 or not isfinite(x): break
    if s*im(x) >= 3: return x, k>0
    if verbal: print x
    return x/(x+1) * (log(a/x)+1), k>0

>>> W(-1e-100, -1, True)
-230.258509299405
(-235.721437124736 + 0.0j)
(-235.721158875686 + 0.0j)
(-235.721158875685 + 0.0j)
(mpc(real='-235.72115887568535', imag='0.0'), True)