Add EnzymeCore weakdep and an extension with a custom rule for the Levin transformation

ext/BesselsEnzymeCoreExt.jl

@@ -35,4 +35,17 @@ module BesselsEnzymeCoreExt
       Duplicated(ls, dls)
   end
 
+  # This is fixing a straight bug in Enzyme.
+  function EnzymeRules.forward(func::Const{typeof(sinpi)}, 
+                               ::Type{<:Duplicated}, 
+                               x::Duplicated)
+      Duplicated(sinpi(x.val), pi*cospi(x.val))
+  end
+
+  function EnzymeRules.forward(func::Const{typeof(sinpi)}, 
+                               ::Type{<:Const}, 
+                               x::Const)
+      sinpi(x.val)
+  end
+
 end

src/BesselFunctions/besselk.jl

@@ -654,16 +654,17 @@ function besselk_power_series_temme_basal(v::V, x::Float64) where{V}
 end
 
 function besselk_power_series_int(v, x::Float64)
-    v < zero(v) && return besselk_power_series_int(-v, x)
-    flv = Int(floor(v))
-    _v  = v - flv
+    v = abs(v)
+    (_v, flv) = modf(v)
+    if _v > 1/2
+      (_v, flv) = (_v-one(_v), flv+1)
+    end
     (kv, kvp1) = besselk_power_series_temme_basal(_v, x)
-    abs(v) < 1/2 && return kv
     twodx = 2/x
-    for _ in 1:(flv-1)
+    for _ in 1:flv
         _v += 1
         (kv, kvp1) = (kvp1, muladd(twodx*_v, kvp1, kv))
     end
-    kvp1
+    kv
 end
 

src/Math/Math.jl

@@ -155,6 +155,6 @@ end
 end
 
 # TODO (cg 2023/05/16 18:09): dispute this cutoff.
-isnearint(x) = abs(x-round(x)) < 1e-7
+isnearint(x) = abs(x-round(x)) < 1e-5
 
 end

test/besselk_enzyme_test.jl

@@ -0,0 +1,38 @@
+using EnzymeCore, Enzyme
+import Bessels.BesselFunctions: besselkx_levin
+import Bessels.BesselFunctions: besselk_power_series
+
+dbesselkx_dv(v, x) = autodiff(Forward, _v->besselkx_levin(_v, x, Val(30)), 
+                              Duplicated, Duplicated(v, 1.0))[2]
+
+dbesselkx_dx(v, x) = autodiff(Forward, _x->besselkx_levin(v, _x, Val(30)), 
+                              Duplicated, Duplicated(x, 1.0))[2]
+
+#=
+dbesselk_ps_dv(v, x) = autodiff(Forward, _v->besselk_power_series(_v, x), 
+                                Duplicated, Duplicated(v, 1.0))[2]
+
+dbesselk_ps_dx(v, x) = autodiff(Forward, _x->besselk_power_series(v, _x), 
+                                Duplicated, Duplicated(x, 1.0))[2]
+=#
+
+
+for line in eachline("data/besselk/enzyme/besselkx_levin_enzyme_tests.csv")
+  (v, x, dv, dx) = parse.(Float64, split(line))
+  test_dv = dbesselkx_dv(v, x)
+  test_dx = dbesselkx_dx(v, x)
+  @test isapprox(dv, test_dv, rtol=5e-14)
+  @test isapprox(dx, test_dx, rtol=5e-14)
+end
+
+#=
+for line in eachline("data/besselk/enzyme/besselk_power_series_enzyme_tests.csv")
+  (v, x, dv, dx) = parse.(Float64, split(line))
+  test_dv = dbesselk_ps_dv(v, x)
+  test_dx = dbesselk_ps_dx(v, x)
+  #@test isapprox(dv, test_dv, rtol=5e-14)
+  #@test isapprox(dx, test_dx, rtol=5e-14)
+
+end
+=#
+

test/runtests.jl

@@ -10,4 +10,4 @@ import SpecialFunctions
 @time @testset "gamma" begin include("gamma_test.jl") end
 @time @testset "airy" begin include("airy_test.jl") end
 @time @testset "sphericalbessel" begin include("sphericalbessel_test.jl") end
-@time @testset "enzyme autodiff" begin include("enzyme_test.jl") end
+@time @testset "besselk enzyme autodiff" begin include("besselk_enzyme_test.jl") end