1
1
# ### Lambert W function ####
2
2
3
3
"""
4
- lambertw(z::Complex{T}, k::V=0, maxits=1000) where {T<:Real, V<:Integer}
5
- lambertw(z::T, k::V=0, maxits=1000) where {T<:Real, V<:Integer}
4
+ lambertw(z::Number, k::Integer=0, maxits=1000)
6
5
7
6
Compute the `k`th branch of the Lambert W function of `z`. If `z` is real, `k` must be
8
7
either `0` or `-1`. For `Real` `z`, the domain of the branch `k = -1` is `[-1/e, 0]` and the
@@ -26,16 +25,16 @@ julia> lambertw(Complex(-10.0, 3.0), 4)
26
25
-0.9274337508660128 + 26.37693445371142im
27
26
```
28
27
"""
29
- lambertw (z, k:: Integer = 0 , maxits:: Integer = 1000 ) = _lambertw (float (z), k, maxits)
28
+ lambertw (z:: Number , k:: Integer = 0 , maxits:: Integer = 1000 ) = _lambertw (float (z), k, maxits)
30
29
31
30
# lambertw(e + 0im, k) is ok for all k
32
31
# Maybe this should return a float. But, this should cause no type instability in any case
33
- function _lambertw (:: typeof (MathConstants. e), k, maxits)
32
+ function _lambertw (:: typeof (MathConstants. e), k:: Integer , maxits:: Integer )
34
33
k == 0 && return 1
35
34
throw (DomainError (k))
36
35
end
37
- _lambertw (x:: Irrational , k, maxits) = _lambertw (float (x), k, maxits)
38
- function _lambertw (x:: Union{Integer, Rational} , k, maxits)
36
+ _lambertw (x:: Irrational , k:: Integer , maxits:: Integer ) = _lambertw (float (x), k, maxits)
37
+ function _lambertw (x:: Union{Integer, Rational} , k:: Integer , maxits:: Integer )
39
38
if k == 0
40
39
x == 0 && return float (zero (x))
41
40
x == 1 && return convert (typeof (float (x)), omega) # must be a more efficient way
45
44
46
45
# ## Real x
47
46
48
- function _lambertw (x:: Real , k, maxits)
47
+ function _lambertw (x:: Real , k:: Integer , maxits:: Integer )
49
48
k == 0 && return lambertw_branch_zero (x, maxits)
50
49
k == - 1 && return lambertw_branch_one (x, maxits)
51
50
throw (DomainError (k, " lambertw: real x must have branch k == 0 or k == -1" ))
54
53
# Real x, k = 0
55
54
# There is a magic number here. It could be noted, or possibly removed.
56
55
# In particular, the fancy initial condition selection does not seem to help speed.
57
- function lambertw_branch_zero (x:: T , maxits) where T<: Real
56
+ function lambertw_branch_zero (x:: T , maxits:: Integer ) where T<: Real
58
57
isfinite (x) || return x
59
58
one_t = one (T)
60
59
oneoe = - inv (convert (T, MathConstants. e)) # The branch point
@@ -72,7 +71,7 @@ function lambertw_branch_zero(x::T, maxits) where T<:Real
72
71
end
73
72
74
73
# Real x, k = -1
75
- function lambertw_branch_one (x:: T , maxits) where T<: Real
74
+ function lambertw_branch_one (x:: T , maxits:: Integer ) where T<: Real
76
75
oneoe = - inv (convert (T, MathConstants. e))
77
76
x == oneoe && return - one (T) # W approaches -1 as x -> -1/e from above
78
77
oneoe < x || throw (DomainError (x)) # branch domain exludes x < -1/e
83
82
84
83
# ## Complex z
85
84
86
- _lambertw (z:: Complex{<:Integer} , k, maxits) = _lambertw (float (z), k, maxits)
85
+ _lambertw (z:: Complex{<:Integer} , k:: Integer , maxits:: Integer ) = _lambertw (float (z), k, maxits)
87
86
# choose initial value inside correct branch for root finding
88
- function _lambertw (z:: Complex{T} , k, maxits) where T<: Real
89
- one_t = one (T)
87
+ function _lambertw (z:: Complex{T} , k:: Integer , maxits:: Integer ) where T<: Real
90
88
local w:: Complex{T}
91
89
pointseven = 7 // 10
92
90
if abs (z) <= inv (convert (T, MathConstants. e))
120
118
121
119
# Use Halley's root-finding method to find
122
120
# x = lambertw(z) with initial point x0.
123
- function lambertw_root_finding (z:: T , x0:: T , maxits) where T <: Number
121
+ function lambertw_root_finding (z:: T , x0:: T , maxits:: Integer ) where T <: Number
124
122
two_t = convert (T, 2 )
125
123
x = x0
126
124
lastx = x
@@ -203,7 +201,7 @@ Base.BigFloat(::Irrational{:ω}) = omega_const(BigFloat)
203
201
# solve for α₂. We get α₂ = 0.
204
202
# Compute array of coefficients μ in (4.22).
205
203
# m[1] is μ₀
206
- function compute_branch_point_coeffs (T:: DataType , n:: Int )
204
+ function compute_branch_point_coeffs (T:: Type{<:Number} , n:: Integer )
207
205
a = Vector {T} (undef, n)
208
206
m = Vector {T} (undef, n)
209
207
259
257
# Why is wser5 omitted ?
260
258
# p is the argument to the series which is computed
261
259
# from x before calling `branch_point_series`.
262
- function branch_point_series (p, x)
260
+ function branch_point_series (p:: Real , x:: Real )
263
261
x < 4e-11 && return wser3 (p)
264
262
x < 1e-5 && return wser7 (p)
265
263
x < 1e-3 && return wser12 (p)
@@ -273,7 +271,7 @@ function branch_point_series(p, x)
273
271
end
274
272
275
273
# These may need tuning.
276
- function branch_point_series (p:: Complex{T} , z) where T<: Real
274
+ function branch_point_series (p:: Complex{T} , z:: Complex{T} ) where T<: Real
277
275
x = abs (z)
278
276
x < 4e-11 && return wser3 (p)
279
277
x < 1e-5 && return wser7 (p)
@@ -287,13 +285,13 @@ function branch_point_series(p::Complex{T}, z) where T<:Real
287
285
return wser290 (p)
288
286
end
289
287
290
- function _lambertw0 (x) # 1 + W(-1/e + x) , k = 0
288
+ function _lambertw0 (x:: Number ) # 1 + W(-1/e + x) , k = 0
291
289
ps = 2 * MathConstants. e * x
292
290
series_arg = sqrt (ps)
293
291
branch_point_series (series_arg, x)
294
292
end
295
293
296
- function _lambertwm1 (x) # 1 + W(-1/e + x) , k = -1
294
+ function _lambertwm1 (x:: Number ) # 1 + W(-1/e + x) , k = -1
297
295
ps = 2 * MathConstants. e * x
298
296
series_arg = - sqrt (ps)
299
297
branch_point_series (series_arg, x)
0 commit comments