add array support

Author: heltonmc

src/besseli.jl

@@ -164,11 +164,11 @@ end
 
 Modified Bessel function of the second kind of order nu, ``I_{nu}(x)``.
 """
-besseli(nu::Real, x::Real) = _besseli(nu, float(x))
+besseli(nu, x::Real) = _besseli(nu, float(x))
 
 _besseli(nu, x::Float16) = Float16(_besseli(nu, Float32(x)))
 
-function _besseli(nu, x::T) where T <: Union{Float32, Float64}
+function _besseli(nu::T, x::T) where T <: Union{Float32, Float64}
     isinteger(nu) && return _besseli(Int(nu), x)
     ~isfinite(x) && return x
     abs_nu = abs(nu)
@@ -205,6 +205,13 @@ function _besseli(nu::Integer, x::T) where T <: Union{Float32, Float64}
     end
 end
 
+function _besseli(nu::AbstractRange, x::T) where T
+    (nu[1] >= 0 && step(nu) == 1) || throw(ArgumentError("nu must be >= 0 with step(nu)=1"))
+    out = Vector{T}(undef, length(nu))
+    out[end-1], out[end] = _besseli(nu[end-1], x), _besseli(nu[end], x)
+    return besselk_down_recurrence!(out, x, nu) 
+end
+
 """
     besseli_positive_args(nu, x::T) where T <: Union{Float32, Float64}
 

src/besselj.jl

@@ -206,11 +206,11 @@ end
 Bessel function of the first kind of order nu, ``J_{nu}(x)``.
 nu and x must be real where nu and x can be positive or negative.
 """
-besselj(nu::Real, x::Real) = _besselj(nu, float(x))
+besselj(nu, x::Real) = _besselj(nu, float(x))
 
 _besselj(nu, x::Float16) = Float16(_besselj(nu, Float32(x)))
 
-function _besselj(nu, x::T) where T <: Union{Float32, Float64}
+function _besselj(nu::T, x::T) where T <: Union{Float32, Float64}
     isinteger(nu) && return _besselj(Int(nu), x)
     abs_nu = abs(nu)
     abs_x = abs(x)
@@ -255,6 +255,18 @@ function _besselj(nu::Integer, x::T) where T <: Union{Float32, Float64}
     end
 end
 
+function _besselj(nu::AbstractRange, x::T) where T
+    (nu[1] >= 0 && step(nu) == 1) || throw(ArgumentError("nu must be >= 0 with step(nu)=1"))
+    out = Vector{T}(undef, length(nu))
+    if nu[end] < x
+        out[1], out[2] = _besselj(nu[1], x), _besselj(nu[2], x)
+        return besselj_up_recurrence!(out, x, nu)
+    else
+        out[end-1], out[end] = _besselj(nu[end-1], x), _besselj(nu[end], x)
+        return besselj_down_recurrence!(out, x, nu)
+    end
+end
+
 """
     besselj_positive_args(nu, x::T) where T <: Float64
 

src/besselk.jl

@@ -187,11 +187,11 @@ end
 
 Modified Bessel function of the second kind of order nu, ``K_{nu}(x)``.
 """
-besselk(nu::Real, x::Real) = _besselk(nu, float(x))
+besselk(nu, x::Real) = _besselk(nu, float(x))
 
 _besselk(nu, x::Float16) = Float16(_besselk(nu, Float32(x)))
 
-function _besselk(nu, x::T) where T <: Union{Float32, Float64}
+function _besselk(nu::T, x::T) where T <: Union{Float32, Float64}
     isinteger(nu) && return _besselk(Int(nu), x)
     abs_nu = abs(nu)
     abs_x = abs(x)
@@ -216,6 +216,13 @@ function _besselk(nu::Integer, x::T) where T <: Union{Float32, Float64}
     end
 end
 
+function _besselk(nu::AbstractRange, x::T) where T
+    (nu[1] >= 0 && step(nu) == 1) || throw(ArgumentError("nu must be >= 0 with step(nu)=1"))
+    out = Vector{T}(undef, length(nu))
+    out[1], out[2] = _besselk(nu[1], x), _besselk(nu[2], x)
+    return besselk_up_recurrence!(out, x, nu)
+end
+
 """
     besselk_positive_args(x::T) where T <: Union{Float32, Float64}
 

src/bessely.jl

@@ -243,11 +243,11 @@ end
 Bessel function of the second kind of order nu, ``Y_{nu}(x)``.
 nu and x must be real where nu and x can be positive or negative.
 """
-bessely(nu::Real, x::Real) = _bessely(nu, float(x))
+bessely(nu, x::Real) = _bessely(nu, float(x))
 
 _bessely(nu, x::Float16) = Float16(_bessely(nu, Float32(x)))
 
-function _bessely(nu, x::T) where T <: Union{Float32, Float64}
+function _bessely(nu::T, x::T) where T <: Union{Float32, Float64}
     isnan(nu) || isnan(x) && return NaN
     isinteger(nu) && return _bessely(Int(nu), x)
     abs_nu = abs(nu)
@@ -296,6 +296,13 @@ function _bessely(nu::Integer, x::T) where T <: Union{Float32, Float64}
     end
 end
 
+function _bessely(nu::AbstractRange, x::T) where T
+    (nu[1] >= 0 && step(nu) == 1) || throw(ArgumentError("nu must be >= 0 with step(nu)=1"))
+    out = Vector{T}(undef, length(nu))
+    out[1], out[2] = _bessely(nu[1], x), _bessely(nu[2], x)
+    return besselj_up_recurrence!(out, x, nu)
+end
+
 #####
 #####  `bessely` for positive arguments and orders
 #####

src/recurrence.jl

@@ -2,20 +2,29 @@
 # outputs both (bessel(x, nu_end), bessel(x, nu_end+1)
 # x = 0.1; k0 = besselk(0,x); k1 = besselk(1,x);
 # besselk_up_recurrence(x, k1, k0, 1, 5) will give besselk(5, x)
-@inline function besselk_up_recurrence(x, jnu, jnum1, nu_start, nu_end)
+function besselk_up_recurrence(x, jnu, jnum1, nu_start, nu_end)
     x2 = 2 / x
     while nu_start < nu_end + 0.5 # avoid inexact floating points when nu isa float
         jnum1, jnu = jnu, muladd(nu_start*x2, jnu, jnum1)
         nu_start += 1
     end
     return jnum1, jnu
 end
+function besselk_up_recurrence!(out, x::T, nu_range) where T
+    x2 = 2 / x
+    k = 3
+    for nu in nu_range[2:end-1]
+        out[k] = muladd(nu*x2, out[k-1], out[k-2])
+        k += 1
+    end
+    return out
+end
 
 # forward recurrence relation for besselj and bessely
 # outputs both (bessel(x, nu_end), bessel(x, nu_end+1)
 # x = 0.1; y0 = bessely0(x); y1 = bessely1(x);
 # besselj_up_recurrence(x, y1, y0, 1, 5) will give bessely(5, x)
-@inline function besselj_up_recurrence(x::T, jnu, jnum1, nu_start, nu_end) where T
+function besselj_up_recurrence(x::T, jnu, jnum1, nu_start, nu_end) where T
     x2 = 2 / x
     while nu_start < nu_end + 0.5 # avoid inexact floating points when nu isa float
         jnum1, jnu = jnu, muladd(nu_start*x2, jnu, -jnum1)
@@ -27,12 +36,21 @@ end
     # NaN inputs need to be checked before getting to this routine 
     return isnan(jnum1) ? (-T(Inf), -T(Inf)) : (jnum1, jnu)
 end
+function besselj_up_recurrence!(out, x::T, nu_range) where T
+    x2 = 2 / x
+    k = 3
+    for nu in nu_range[2:end-1]
+        out[k] = muladd(nu*x2, out[k-1], -out[k-2])
+        k += 1
+    end
+    return out
+end
 
 # backward recurrence relation for besselj and bessely
 # outputs both (bessel(x, nu_end), bessel(x, nu_end-1)
 # x = 0.1; j10 = besselj(10, x); j11 = besselj(11, x);
 # besselj_down_recurrence(x, j10, j11, 10, 1) will give besselj(1, x)
-@inline function besselj_down_recurrence(x, jnu, jnup1, nu_start, nu_end)
+function besselj_down_recurrence(x, jnu, jnup1, nu_start, nu_end)
     x2 = 2 / x
     while nu_start > nu_end - 0.5
         jnup1, jnu = jnu, muladd(nu_start*x2, jnu, -jnup1)
@@ -41,6 +59,16 @@ end
     return jnup1, jnu
 end
 
+function besselj_down_recurrence!(out, x::T, nu_range) where T
+    x2 = 2 / x
+    k = length(nu_range) - 2
+    for nu in nu_range[end-1:-1:2]
+        out[k] = muladd(nu*x2, out[k+1], -out[k+2])
+        k -= 1
+    end
+    return out
+end
+
 #=
 # currently not used
 # backward recurrence relation for besselk and besseli
@@ -56,3 +84,12 @@ end
     return jnup1, jnu
 end
 =#
+function besselk_down_recurrence!(out, x::T, nu_range) where T
+    x2 = 2 / x
+    k = length(nu_range) - 2
+    for nu in nu_range[end-1:-1:2]
+        out[k] = muladd(nu*x2, out[k+1], out[k+2])
+        k -= 1
+    end
+    return out
+end

test/besseli_test.jl

@@ -94,7 +94,7 @@ t = [besselix(m, x) for m in m, x in x]
 @test t ≈ [SpecialFunctions.besselix(m, x) for m in m, x in x]
 @test besselix(10, Float16(1.0)) isa Float16
 
-## Tests for besselk
+## Tests for besseli
 
 ## test all numbers and orders for 0<nu<100
 x = [0.01, 0.05, 0.1, 0.2, 0.4, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.91, 0.92, 0.93, 0.95, 0.96, 0.97, 0.98, 0.99, 0.995, 0.999, 1.0, 1.001, 1.01, 1.05, 1.1, 1.2, 1.4, 1.6, 1.8, 1.9, 2.5, 3.0, 3.5, 4.0]
@@ -106,6 +106,12 @@ for v in nu, xx in x
     @test isapprox(besseli(Float32(v), Float32(xx)), Float32(sf))
 end
 
+# test nu_range
+@test besseli(0:50, 2.0) ≈ SpecialFunctions.besseli.(0:50, 2.0) rtol=1e-13
+@test besseli(0.5:1:10.5, 2.0) ≈ SpecialFunctions.besseli.(0.5:1:10.5, 2.0) rtol=1e-13
+
+### need to fix method ambiguities for other functions ###### 
+
 # test Inf
 @test isinf(besseli(2, Inf))
 

test/besselj_test.jl

@@ -114,6 +114,12 @@ for v in nu, xx in x
     @test isapprox(Bessels.besseljy_positive_args(v, xx)[1], Float64(sf), rtol=5e-11)
 end
 
+# test nu_range
+@test besselj(0:50, 2.0) ≈ SpecialFunctions.besselj.(0:50, 2.0) rtol=1e-11
+@test besselj(0:50, 100.0) ≈ SpecialFunctions.besselj.(0:50, 100.0) rtol=1e-11
+@test besselj(0.5:1:10.5, 2.0) ≈ SpecialFunctions.besselj.(0.5:1:10.5, 2.0) rtol=1e-11
+@test besselj(0.5:1:10.5, 40.0) ≈ SpecialFunctions.besselj.(0.5:1:10.5, 40.0) rtol=1e-11
+
 # test Float16 and Float32
 @test besselj(10, Float16(1.0)) isa Float16
 @test besselj(10.2f0, 1.0f0) isa Float32

test/besselk_test.jl

@@ -126,6 +126,10 @@ for v in nu, xx in x
     @test isapprox(besselk(Float32(v), Float32(xx)), Float32(sf))
 end
 
+# test nu_range
+@test besselk(0:50, 2.0) ≈ SpecialFunctions.besselk.(0:50, 2.0) rtol=1e-13
+@test besselk(0.5:1:10.5, 12.0) ≈ SpecialFunctions.besselk.(0.5:1:10.5, 12.0) rtol=1e-13
+
 # test Float16
 @test besselk(10, Float16(1.0)) isa Float16
 @test besselkx(10, Float16(1.0)) isa Float16

test/bessely_test.jl

@@ -94,6 +94,12 @@ for v in nu, xx in x
     @test isapprox(bessely(Float32(v), Float32(xx)), Float32(sf))
 end
 
+# test nu_range
+@test bessely(0:50, 2.0) ≈ SpecialFunctions.bessely.(0:50, 2.0) rtol=1e-11
+@test bessely(0:50, 100.0) ≈ SpecialFunctions.bessely.(0:50, 100.0) rtol=1e-11
+@test bessely(0.5:1:10.5, 2.0) ≈ SpecialFunctions.bessely.(0.5:1:10.5, 2.0) rtol=1e-11
+@test bessely(0.5:1:10.5, 40.0) ≈ SpecialFunctions.bessely.(0.5:1:10.5, 40.0) rtol=1e-11
+
 # test Float16
 @test bessely(10, Float16(1.0)) isa Float16
 @test bessely(10.2f0, 1.0f0) isa Float32

test/sphericalbessel_test.jl

@@ -86,8 +86,8 @@ x = NaN
 @test Bessels.sphericalbesselk(Float16(1.4), Float16(1.2)) isa Float16
 @test Bessels.sphericalbesselk(1.0f0, 1.2f0) isa Float32
 
-@test Bessels.sphericalbesseli(Float16(1.4), Float16(1.2)) isa Float16
-@test Bessels.sphericalbesseli(1.0f0, 1.2f0) isa Float32
+@test Bessels.sphericalbesseli(1, Float16(1.2)) isa Float16
+@test Bessels.sphericalbesseli(2, 1.2f0) isa Float32
 
 for x in 0.5:1.5:100.0, v in [0, 1, 2, 3, 4, 5.5, 8.2, 10]
     @test isapprox(Bessels.sphericalbesselj(v, x), SpecialFunctions.sphericalbesselj(v, x), rtol=1e-12)