lambertw, all branches

[Albert Chan]

(04-07-2023 02:47 PM)Albert Chan Wrote:  (*) Based from tests, solve for y*ln(y)=a is better than x*e^x=a, even with help of Halley's method.

Not when W guess is extremely close to true root.
We can use Newton's method, f = x*e^x - a, for finishing touches

Since we are now using W loops for guess, we might as well make it slightly worse.
Instead of log(abs(x)/A), make it log(abs(x)) - log(A), literally, f = x + ln(x) - ln_k(a)
Code is easier to read, and cheaper to compute. (last term is a constant)

Instead of comparing *both* complex parts converged, we test for relative errors.
This solved the problem if any complex parts is bouncing around zero.
For symmetry, I used relative error test for Y loop as well.

With relative error test, it should not get stuck in infinte loops → loop limits removed

complex.W code, version 3

Code:

r, err = 0.36787944117144233, -1.2428753672788363E-17

require'complex'
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 = I(a:real() + r + err, a:imag())
    local small_imag = (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
        until I.abs(h/y) < 1e-9
        return flip(a/y)
    end
    local A, T = I.abs(a), I.new(0,2*k*pi+I.arg(a))
    local x = T     -- W rough guess
    if k==0  then x = (A+100)/(A+50); x = I.log1p(a*x) / x end
    if small_imag then x = I.log(-a) + (A<.5 and 0 or T/2) end
    T = T + log(A)  -- = ln_k(a)
    if not finite(A) then return flip(T) end
    if A==0 then return flip(k==0 and a or T-pi*I) end
    if verbal and verbal~=true then x = flip(verbal) end
    repeat
        if verbal then print(flip(x)) end
        h = x/(x+1) * (x + I.log(x) - T)
        x = x - h
    until not (I.abs(h/x) >= 1e-6)
    if verbal then print(flip(x)) end
    h = (x-a*I.exp(-x))/(x+1)   -- final touch
    return flip(I.finite(h) and x-h or x)
end

(04-09-2023 03:59 AM)Albert Chan Wrote:  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

complex.W version 2 was not able to converge after 40 loops.
Redo with version 3:

lua> a
(4.794621373315692+1.4183109273161312*I)
lua> I.W(a, 1, true)
(0+6.570796326794897*I)
(-0.033367069043767406+4.994921913968373*I)
(1.97521636071743e-005+5.000010560789135*I)
(3.753082249819371e-012+5.000000000009095*I)
(1.0141772755203484e-016+5*I)
(-6.474631941441244e-017+5*I)

We need more precisions to get real part right.
For reference, with hexfloat of a: W₁(a) ≈ 4.51974e-018 + 5.00000 I

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

lua> I.W(inf), I.W(-inf)
(inf+0*I)      (inf+3.141592653589793*I)

lua> z = 0
lua> I.W(-z,-1), I.W(z,-1)
(-inf-0*I)      (-inf-3.141592653589793*I)

Update 1/29/24

< h = (a + r + err)
> h = I(a:real() + r + err, a:imag())      -- keep -0*I sign

Update 2/01/24

lyuka branch relative limit reduced, from 1e-12 to 1e-9
Although convergence near -1/e is hard, lyuka (improved) formula is very accurate there.
Further away, guess formula is not as good, but convergence O(rel_err^2) constant < 1

I also removed lua implementation of log1p(x). it had moved to C code