Skip to content

Commit fe60194

Browse files
committed
lambertw: small fixes
1 parent e5f2fb5 commit fe60194

File tree

1 file changed

+6
-5
lines changed

1 file changed

+6
-5
lines changed

src/lambertw.jl

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,9 @@ function _lambertw(z::Complex{T}, k::Integer, maxits::Integer) where T<:Real
100100
k != 0 ? w += complex(0, k * 2 * pi) : nothing
101101
end
102102
elseif k == 0 && imag(z) <= pointseven && abs(z) <= pointseven
103-
w = abs(z+ 1//2) < 1//10 ? imag(z) > 0 ? complex(pointseven, pointseven) : complex(pointseven, -pointseven) : z
103+
w = abs(z+ 1//2) < 1//10 ? imag(z) > 0 ?
104+
complex(pointseven, pointseven) :
105+
complex(pointseven, -pointseven) : z
104106
else
105107
if real(z) == convert(T, Inf)
106108
k == 0 && return z
@@ -118,16 +120,15 @@ end
118120
# Use Halley's root-finding method to find
119121
# x = lambertw(z) with initial point x0.
120122
function lambertw_root_finding(z::T, x0::T, maxits::Integer) where T <: Number
121-
two_t = convert(T, 2)
122123
x = x0
123124
lastx = x
124125
lastdiff = zero(real(T))
125126
converged = false
126-
for i in 1:maxits
127+
for _ in 1:maxits
127128
ex = exp(x)
128129
xexz = x * ex - z
129130
x1 = x + 1
130-
x -= xexz / (ex * x1 - (x + two_t) * xexz / (two_t * x1 ))
131+
x -= xexz / (ex * x1 - (x + 2) * xexz / (2 * x1))
131132
xdiff = abs(lastx - x)
132133
if xdiff <= 3 * eps(lastdiff) || lastdiff == xdiff # second condition catches two-value cycle
133134
converged = true
@@ -136,7 +137,7 @@ function lambertw_root_finding(z::T, x0::T, maxits::Integer) where T <: Number
136137
lastx = x
137138
lastdiff = xdiff
138139
end
139-
converged || @warn("lambertw with z=", z, " did not converge in ", maxits, " iterations.")
140+
converged || @warn("lambertw(", z, ") did not converge in ", maxits, " iterations.")
140141
return x
141142
end
142143

0 commit comments

Comments
 (0)