lambertw, all branches

[Albert Chan]

I am keeping code simple. Edge cases, example W₀(0) = 0, are not implemented.

Code:

r, err = 0.36787944117144233, -1.2428753672788363E-17

require'complex'
function complex.log1p(x) local y = x-(-1); return I.log(y) - ((y-1-x)/y+0) end

function complex.W01(a, K, verbal, y)   -- K = abs(W branch) = 0 or 1
    K = K or 0
    local h, s = a+r+err, function(y) return I.log1p((y-r-err)/r) end
    if not y then
        y = I.sqrt(2*r*h)*(1-2*abs(K))  -- assume close to branch point
        y = r + y * I.sqrt(1+y/(3*r))   -- assume W imag part tiny
    end
    for i = 1, 10 do
        if verbal then print(a/y, 'W01') end
        h = y - (y+a) / s(y)
        y = y - h -- Newton's correction
        if y == y+h*0x1p-13 or not I.finite(y) then return a/y end
    end
    return NaN
end

function complex.W(a, k, verbal, x)
    a, k = I.new(a), k or 0
    local K, h, s = abs(k), a+r+err
    if I.abs(h)<.25 and (K==0 or K==1 and signbit(k*I.imag(a))) then
        return complex.W01(a, K, verbal), 'W01'
    end
    local A, T = I.abs(a), 2*k*pi+I.arg(a)
    s = function(x) return x + I.new(log(I.abs(x)/A), I.arg(x)-T) end
    if not x then     -- W rough guess
        if K > 1 then    x = 2*k*pi*I
        elseif K==1 then x = A==1 and 1 or I.log(-a)
        else --[[K==0]]  x = A==1 and 1 or I.log1p(a)
        end
    end
    for i = 1, 10 do
        if verbal then print(x) end
        h = x/(x-(-1)) * s(x)
        x = x - h -- Newton's correction
        if x == x+h*0x1p-13 or not I.finite(x) then return x end
    end
    return complex.W01(a, K, verbal, a/I.real(x)), 'W01'
end

W01 is solving W branch 0 or an ambiguous ±1 branch (with the small imag part)
The code is solving for y*ln(y) = a --> y = e^W(a) --> W(a) = a/y

We can "pull out" conjugate: W01(conj(a),k) = conj(W01(a,k))

lua> I.W(3+4*I, 5, true)
(0+31.41592653589793*I)
(-1.815554248884793+30.714634540472343*I)
(-1.8170058914565277+30.71333413897562*I)
(-1.8170058918466272+30.713334137004896*I)
(-1.8170058918466274+30.713334137004896*I)

lua> I.W(3+4*I, -5, true)
(-0-31.41592653589793*I)
(-1.7565805546696363-28.861921295458465*I)
(-1.7547636407605829-28.85710103800315*I)
(-1.7547636279545493-28.85710104741841*I)
(-1.7547636279545493-28.85710104741841*I)

lua> I.W(-0.3,-1,true)
(-1.7824743623491122+0*I)      W01
(-1.7813374873395937+0*I)      W01
(-1.7813370234217047+0*I)      W01
(-1.781337023421628+0*I)        W01

lua> I.W(-0.03,-1,true)
(-3.506557897319982-0*I)
(-5.261733759227882-0*I)
(-5.144793893666586-0*I)
(-5.1444827237861634-0*I)
(-5.144482721515682-0*I)
(-5.144482721515682-0*I)

lua> I.W(-.1+1e-100*I, -1, true)
(-2.3025850929940455-1e-099*I)
(-3.7769077194622533+1.2592576024726166e-099*I)
(-3.579124176402325-8.545840728943592*I)
(-3.7689826579869887-0.8337217972782689*I)
(-3.5535464061313635-0.025256342802877585*I)
(-3.5771468383315437+6.535012002211599e-005*I)
(-3.5770902410326704-8.721224618630954*I)
(-3.7811926340275566-0.8444302692850991*I)
(-3.5536175209072614-0.02647513850865879*I)
(-3.5771431199160215+6.824715919089008e-005*I)
(-3.5770875001635236+0*I)      W01
(-3.5771520641833883-1.3880349425219542e-099*I)      W01
(-3.5771520639572962-1.3880252213229777e-099*I)      W01
(-3.5771520639572976-1.388025221322979e-099*I)        W01

W code f=0 have trouble converging. Switched to W01 help a lot.
W01 use W last guess, but dropped the unreliable imaginary part.