(28/48/50) Lambert W Function

[Albert Chan]

Here is my implementation of accurate x = W(a), both branches.

x * e^x = a
ln(x) + x = ln(a)      → f = x + ln(x/a) = 0

x - f/f' = x - (x + ln(x/a)) * x/(1+x) = (1 - ln(x/a)) * x/(1+x)

W(-1/e) = -1 --> a ≈ -1/e, x ≈ -1 --> (1+x) is tiny

Let 1/e = r + err

1 - ln(x/a) = - ln(x*(r+err)/a) = - log1p((x*err-(a+r)+r*(1+x))/a)

x + ln(x/a) = (x+1) - (1 - ln(x/a))

Terms inside log1p sorted in increasing size, r*(1+x) the biggest.
With this, and some minor changes to my expW function, I get this:

Code:

function W(a, x, verbal)
    x = x or 0  -- default branch 0
    local h, s
    if a <= -0.1 then
        h = a + r + err
        x = sqrt(2*h/r) * (x+x+1)
        x = x * sqrt(1+x/3) -- estimate for e*h = (1+x)*log1p(x) - x
        x = a * log1p(x)/(r*x+h) -- Newton step
        s = function(x) return (x+1) + log1p((x*err-(a+r)+r*(1+x))/a) end
    else
        if a==0 then return (x==0 and 0 or -Inf) end
        x = (x+1) + (x+x+1)*a/4  -- e^W(a) guess
        x = a/(x+a) * (log(x)+1) -- Newton step
        s = function(x) return x + log(x/a) end
    end
    repeat
        if verbal then print(x) end
        h = x/(1+x) * s(x)
        x = x - h   -- Newton's correction
    until x == x+h*0x1p-13 or x ~= x
    return x
end

lua> W(-0.367, 0, true)
-0.9323991847479226
-0.9323991847479282
lua> W(-0.367, -1, true)
-1.0707918867680617
-1.0707918867680521

lua> W(1e10, 0, true)
18.111645253927524
20.02372402639066
20.028685384073526
20.02868541330495
20.028685413304952

Update 1: added special case W₀(0) = 0, W_-1(0) = -∞ (should it be NaN?)

Update 2: removed guard "if h == err then return -1 end", see next post.

Update 3: improve W guess with a Newton's step, based from y = e^W
y = (y+a) / (log(y)+1)      → x = a/y = a/(y+a) * (log(y)+1)

Update 4: change behavior from "W branch -1 if x" to "x default to branch 0"
Assumed input, x = nil/false/0/-1 only: x = nil/false/0 get branch 0; x = -1 get branch -1