Tests for besseli+Enzyme that cover the new convergence checker.

Author: Chris Geoga

ext/BesselsEnzymeCoreExt.jl

@@ -21,14 +21,14 @@ module BesselsEnzymeCoreExt
   function EnzymeRules.forward(func::Const{typeof(check_convergence)},
                                ::Type{Const{Bool}},
                                t::Duplicated{T}) where{T}
-    (t.val <= eps(T)) && (t.dval <= eps(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}
-    (t.val <= eps(T)*s.val) && (t.dval <= eps(T)*s.dval)
+    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

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/data/besseli/enzyme/README.md

@@ -0,0 +1,35 @@
+```julia
+using ArbNumerics, DelimitedFiles, SpecialFunctions, Enzyme, Bessels
+
+# 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))
+
+# TODO (cg 2023/05/30 11:44): temporarily, we make vgrid[end]=15.1 instead of
+# 15.0 or some other integer because the integer-order `besseli` derivatives
+# need to be fixed with dispatch, like `besselk`.
+vgrid = range(1e-3, 15.1,  length=20) 
+xgrid = range(1e-3, 100.0, length=30)
+vx = vec(collect(Iterators.product(vgrid, xgrid)))
+
+our_dx(v,x) = autodiff(Forward, _x->Bessels.besseli(v, _x), Duplicated, Duplicated(x, 1.0))[2]
+our_dv(v,x) = autodiff(Forward, _v->Bessels.besseli(_v, x), Duplicated, Duplicated(v, 1.0))[2]
+
+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
+```