lambertw, all branches

[Albert Chan]

complex.W version 3, adapted for system without signed zero.
This time I made a customized LN(x) to honor 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 abs(h/y) < 1e-9: return a/y
    s = 1 - 2 * (k<0 or k==0 and s)
    def LN(x, bad=-s*pi):
        x = ln(x)
        return conj(x) if 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 small_imag: 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
    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

Trivia: W(a < -1/e) is complex, W₀(a) = conj(W_-1(a))

x + ln(x) = ln_k(a)

ln₀(a) = ln(|a|) - pi*j + 2*pi*j = ln₁(conj(a))

W₀(a) = W₁(conj(a)) = conj(W_-1(a))

>>> a = -pi/2
>>> W(a, 0, True)
(0.375191803678016 + 1.59508875325509j)
(0.0128176527641888 + 1.58759064221342j)
(-5.58573510823744e-5 + 1.57084664181356j)
(6.00809277877961e-10 + 1.57079632603836j)
(5.72587285769658e-17 + 1.5707963267949j)
(1.30059124689582e-17 + 1.5707963267949j)
>>> W(a, -1, True)
(0.451582705289455 - 1.5707963267949j)
(0.0200654830289762 - 1.59320451813763j)
(-9.57300550001525e-5 - 1.57091376376939j)
(1.39642133562001e-9 - 1.5707963231279j)
(-3.14444301571107e-17 - 1.5707963267949j)
(-5.01152163329281e-17 - 1.5707963267949j)

Confirm x * exp(x) = a, both at the same time.

(±pi/2*j) * exp(±pi/2*j) = (±pi/2*j) * (±j) = pi/2 * j² = -pi/2

Udpate 1/19/24: added |a| = ∞ or 0 edge cases

p2> W(inf), W(-inf)
((+inf + 0.0j), (+inf + 3.14159265358979j))

p2> W(0, -1)
(-inf - 3.14159265358979j)

Default setup is 0 == +0, that is why we get W_-1(+0) = -∞ - pi*I
To get W_-1(0) = -∞, we need special case for k=-1, 0 == −0 + 0*I

Real part -0, Imag part +0 ... it is just too confusing.
I am keeping 0 == +0, for consistency. ln(0) = ln(+0) = -∞

--> W_k(+0) = 0 if k==0 else -∞ + (2*k-sign(k))*pi*I

Of course, we are free to patch it to make it pick what we wanted.
But then, what is W_k(0) in general ? Is 0 still meant +0 + 0*I ?
This is a weakness of not having signed zero.

Update 2/01/24

lyuka branch relative limit reduced, from 1e-12 to 1e-9 (see post #14)