Merge pull request #105 from cgeoga/convg_check

`EnzymeRule` for checking convergence in series-type code

Author: heltonmc

ext/BesselsEnzymeCoreExt.jl

@@ -3,6 +3,7 @@ module BesselsEnzymeCoreExt
   using Bessels, EnzymeCore
   using EnzymeCore.EnzymeRules
   using Bessels.Math
+  import Bessels.Math: check_convergence
 
   # A manual method that separately transforms the `val` and `dval`, because
   # sometimes the `val` can converge while the `dval` hasn't, so just using an
@@ -17,18 +18,36 @@ module BesselsEnzymeCoreExt
       Duplicated(ls, dls)
   end
 
-  # This is fixing a straight bug in Enzyme.
+  function EnzymeRules.forward(func::Const{typeof(check_convergence)},
+                               ::Type{Const{Bool}},
+                               t::Duplicated{T}) where{T}
+    check_convergence(t.val) && check_convergence(t.dval)
+  end
+
+  function EnzymeRules.forward(func::Const{typeof(check_convergence)},
+                               ::Type{Const{Bool}},
+                               t::Duplicated{T},
+                               s::Duplicated{T}) where{T}
+    check_convergence(t.val, s.val) && check_convergence(t.dval, s.val)
+  end
+
+  # This will be fixed upstream: see #861 for Enzyme.jl whenever the next
+  # release occurs.
   function EnzymeRules.forward(func::Const{typeof(sinpi)}, 
                                ::Type{<:Duplicated}, 
                                x::Duplicated)
       (sp, cp) = sincospi(x.val)
       Duplicated(sp, pi*cp*x.dval)
   end
 
+  # #861 will probably also mean this can be deleted at the next release of
+  # Enzyme.jl.
   function EnzymeRules.forward(func::Const{typeof(sinpi)}, 
                                ::Type{<:Const}, 
                                x::Const)
       sinpi(x.val)
   end
 
+
+
 end

src/BesselFunctions/besseli.jl

@@ -414,7 +414,6 @@ _besseli(nu::Union{Int16, Float16}, x::Union{Int16, Float16}) = Float16(_besseli
 _besseli(nu::AbstractRange, x::T) where T = besseli!(zeros(T, length(nu)), nu, x)
 
 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)
     abs_x = abs(x)
@@ -602,7 +601,7 @@ function besselix_large_args(v, x::ComplexOrReal{T}) where T
     for i in 1:MaxIter
         t *= -invx * ((4*v^2 - (2i - 1)^2) / i)
         s += t
-        abs(t) <= eps(T) && break
+        Math.check_convergence(t) && break
     end
     return s / sqrt(2 * (π * x))
 end
@@ -624,7 +623,7 @@ function besseli_power_series(v, x::ComplexOrReal{T}) where T
     xx = x * x * T(0.25)
     for i in 0:MaxIter
         s += t
-        abs(t) < eps(T) * abs(s) && break
+        Math.check_convergence(t, s) && break
         t *= xx / ((v + i + 1) * (i + 1))
     end
     return s * ((x/2)^v / gamma(v + 1))

src/BesselFunctions/besselk.jl

@@ -458,7 +458,7 @@ function besselkx_large_args(v, x::ComplexOrReal{T}) where T
     for i in 1:MaxIter
         t *= invx * ((4*v^2 - (2i - 1)^2) / i)
         s += t
-        abs(t) <= eps(T) && break
+        Math.check_convergence(t) && break
     end
     return s * sqrt(π / (2 * x))
 end
@@ -552,7 +552,7 @@ function besselk_power_series(v, x::ComplexOrReal{T}) where T
         s2 += t2
         t1 *= x^2 / (4k * (k - v))
         t2 *= x^2 / (4k * (k + v))
-        abs(t1) < eps(T) && break
+        Math.check_convergence(t1) && break
     end
 
     xpv = (x/2)^v
@@ -590,8 +590,7 @@ function besselk_temme_series(v::T, x::T) where T <: Float64
         term_vp1 = ck * (pk - (k-1) * fk)
         out_v += term_v
         out_vp1 += term_vp1
-        ((abs(term_v) < eps(T)) && (abs(term_vp1) < eps(T))) && break
-
+        (Math.check_convergence(term_v) && Math.check_convergence(term_vp1)) && break
         fk = (k * fk + pk + qk) / (k^2 - v^2)
         pk /= k - v
         qk /= k + v

src/Math/Math.jl

@@ -154,7 +154,10 @@ end
     )
 end
 
-# TODO (cg 2023/05/16 18:09): dispute this cutoff.
 isnearint(x) = abs(x-round(x)) < 1e-5
 
+@inline check_convergence(term::T) where T = abs(term) <= eps(T)
+
+@inline check_convergence(ser::T, term::T) where T = (abs(term) <= eps(T)*ser)
+
 end

test/besseli_enzyme_test.jl

@@ -0,0 +1,27 @@
+using EnzymeCore, Enzyme
+
+dbesseli_dv(v, x) = autodiff(Forward, _v->besseli(_v, x), 
+                             Duplicated, Duplicated(v, 1.0))[2]
+
+dbesseli_dx(v, x) = autodiff(Forward, _x->besseli(v, _x), 
+                             Duplicated, Duplicated(x, 1.0))[2]
+
+
+for line in eachline("data/besseli/enzyme/besseli_enzyme_tests.csv")
+    (v, x, dv, dx) = parse.(Float64, split(line))
+    test_dv = dbesseli_dv(v, x)
+    test_dx = dbesseli_dx(v, x)
+    # TODO (cg 2023/05/30 12:09): temporarily test at lower rtols in the v=0.001
+    # case. The power series code tests for convergence scaled by the series
+    # value itself, which costs a little bit of rtol. When x is about 20, we
+    # just barely hit that edge regime where we switch to the large argument
+    # expansion and that edge zone also costs a digit. Those are things to
+    # discuss addressing in a different PR I think.
+    if v == 0.001 && x < 21.0
+      @test isapprox(dv, test_dv, rtol=1e-11)
+      @test isapprox(dx, test_dx, rtol=1e-11)
+    else
+      @test isapprox(dv, test_dv, rtol=5e-14)
+      @test isapprox(dx, test_dx, rtol=5e-14)
+    end
+end

test/besseli_test.jl

@@ -131,7 +131,8 @@ end
 (v, x) = 12.0, 3.2
 @test besseli(v,x) ≈ 7.1455266650203694069897133431e-7
 
-(v,x) = 13.0, -1.0
+# Negative arguments only allowed with v::Int!
+(v,x) = 13, -1.0
 @test besseli(v,x) ≈ -1.995631678207200756444e-14
 
 (v,x) = 12.6, -3.0
@@ -146,8 +147,7 @@ end
 (v, x) = -12.3, 8.2
 @test besseli(v,x) ≈ 0.267079696793126091886043602895
 
-(v, x) = -14.0, -9.9
+# Negative arguments only allowed with v::Int!
+(v, x) = -14, -9.9
 @test besseli(v,x) ≈ 0.2892290867115615816280234648
 
-(v, x) = -14.6, -10.6
-#@test besseli(v,x) ≈ -0.157582642056898598750175404443 - 0.484989503203097528858271270828*im

test/data/besseli/enzyme/README.md

@@ -0,0 +1,29 @@
+```julia
+using ArbNumerics, DelimitedFiles, SpecialFunctions
+
+# Because you can't just use FiniteDifferences due to the "numerical noise".
+simplefd(f,x,h=ArbReal(1e-40)) = (f(x+h)-f(x))/h
+
+ArbNumerics.setprecision(ArbReal, digits=100)
+arb_besseli(v,x)  = ArbNumerics.besseli(ArbReal(v), ArbReal(x))
+
+vgrid = range(1e-3, 15.0,  length=20) 
+xgrid = range(1e-3, 100.0, length=30)
+vx = vec(collect(Iterators.product(vgrid, xgrid)))
+
+if !isinteractive()
+  ref_values = map(vx) do vxj
+    (v,x) = vxj
+    dx = simplefd(_x->arb_besseli(v, _x), x) 
+    dv = simplefd(_v->arb_besseli(_v, x), v) 
+    Float64.((dv, dx))
+  end
+
+  out_matrix = hcat(getindex.(vx, 1),          # test v argument
+                    getindex.(vx, 2),          # test x argument
+                    getindex.(ref_values, 1),  # test d/dv value
+                    getindex.(ref_values, 2))  # test d/dx value
+
+  writedlm("besseli_enzyme_tests.csv", out_matrix)
+end
+```