lambertw, all branches

[Albert Chan]

complex.W code version 2, with many improvements.

• W01 function removed, code placed inside W, for branch point cases only.
  This allow for brute force confirmation that it always converge to correct branch.
  Result guaranteed finite, NaN check (to get out of infinite loops) not needed.
• replaced hexfloat 0x1p-13, 0x1p-26 to more readable 1e-4, 1e-8
• W guess improved, for all branches.
• f formula rearranged, to cause catastrophic cancellation, on purpose!
  Imag part (if small) are not very good, but at least it converge quickly.
  We got very good estimated W real part, and a reasonable sized imag part.
• Added convergence confirmation, based from loop counts.
  If we reached solution before max loops, we have convergence.
  If not, result may still be good, but W_k(a) likely purely imaginary.
    
• if |Im(x)| < pi, we are on W branch 0, or branch ±1 with small imag part.
  To be safe, I lower the limit, from pi to 3.
  
  Within this region, we can improve result by y = e^W Newton step.
  With good W real part, one correction is enough, to fix imag part.
  
  y = (a+y) / (ln(y) + 1)
  a/x = (a+a/x) / (ln(a/x) + 1)
  x = x/(x+1) * (ln(a/x) + 1)
   
• Note that all iterations are using equivalent f = x - (ln(a) - ln(x)) - 2*k*pi*i
  W01 code use accurate denominator around branch point, ln(y)+1 = log1p((y-1/e)*e)
  W code use the general form of f, for all k's, but may cost of some accuracy, if |k|≤1
  W final correction (with restriction, see next post) use accurate denominator around 0, ln(y)+1
  

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.W(a, k, verbal)
    a, k = I.new(a), k or 0
    local flip = function(x) return x end   -- identity function
    if k<0 or k==0 and signbit(I.imag(a)) then flip=I.conj; a=flip(a); k=-k end
    local h, small_imag = a+r+err, (k==1 and signbit(I.imag(a)))
    if I.abs(h)<.25 and (k==0 or small_imag) then
        local y = I.sqrt(2*r*h) * (1-2*k)   -- close to branch point
        y = r + y*I.sqrt(1+y/(3*r))         -- with small imag part
        repeat
            if verbal then print(flip(a/y)) end
            h = y - (y+a) / I.log1p((y-r-err)/r)
            y = y - h -- Newton's correction
        until y == y+h*1e-4
        return flip(a/y), true
    end
    local A, T = I.abs(a), (2*k*pi+I.arg(a))*I
    local x = T -- W rough guess
    if small_imag and A<.6 then x = I.log(-a) end
    if k==0 then x = I.abs(a+1)<.6 and T/2 or I.log1p(a) end
    k = 40      -- max loops
    repeat
        if verbal then print(flip(x)) end
        h = x/(x+1) * (x + I.new(log(I.abs(x)/A),I.arg(x)) - T)
        x, k = x-h, k-1
    until k==0 or x == x+h*1e-8 or not I.finite(x)
    if I.imag(x) >= 3 then return flip(x), k>0 end
    if verbal then print(flip(x)) end
    return flip(x/(x+1) * (I.log(a/x)+1)), k>0
end

lua> I.W(3+4*I, 5, true)
(0+32.34322175389954*I)
(-1.8166634488635778+30.71625713741827*I)
(-1.8170058961963056+30.713334138441088*I)
(-1.8170058918466274+30.713334137004896*I)      true

lua> I.W(3+4*I, -5, true)
(0-30.48863131789632*I)
(-1.754507767775116-28.860288696238875*I)
(-1.7547636338945982-28.857101048916533*I)
(-1.7547636279545493-28.85710104741841*I)      true

lua> I.W(-0.3678, 0, true)
(-0.9793607152032429+0*I)
(-0.9793607149578304+0*I)
(-0.9793607149578306+0*I)      true

lua> I.W(-.1+1e-100*I,-1,true)
(-2.3025850929940455-1e-099*I)
(-3.7769077194622533-5.084465616380436e-100*I)
(-3.579124176402325-5.180347740277737e-100*I)
(-3.5771522746958055-5.181454351722726e-100*I)
(-3.5771520639573-5.181454470168089e-100*I)
(-3.577152063957297-5.18145447016809e-100*I)
(-3.577152063957297-1.3880252213229779e-099*I)      true

lua> I.W(-.1+1e-100*I,0,true)
(-0.10536051565782631+1.1111111111111111e-100*I)
(-0.11161906744205384+1.1848854610650868e-100*I)
(-0.11183232962808998+1.1874337731206376e-100*I)
(-0.11183255915869776+1.1874365171431562e-100*I)
(-0.11183255915896297+1.1874365171463266e-100*I)
(-0.11183255915896298+1.259138243719729e-100*I)      true

lua> I.W(-.1+1e-100*I,1,true)
(0+9.42477796076938*I)
(-4.330508133087425+7.394500450320804*I)
(-4.448949088733658+7.307107936743056*I)
(-4.449098178536345+7.307060789152774*I)
(-4.44909817870089+7.3070607892176085*I)      true

lua> x = 5*I
lua> a = x * I.exp(x)
lua> x + x:log() - a:log()
(0+6.283185307179586*I)
lua> I.W(a, 1)
(3.1321400622130957e-068+5*I)      false