Merge pull request #94 from sharanry/sy/type_instability

Make function transformation type stable

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "Integrals"
 uuid = "de52edbc-65ea-441a-8357-d3a637375a31"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "3.1.1"
+version = "3.1.2"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -38,7 +38,8 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["SciMLSensitivity", "FiniteDiff", "Pkg", "SafeTestsets", "Test"]
+test = ["SciMLSensitivity", "StaticArrays", "FiniteDiff", "Pkg", "SafeTestsets", "Test"]

src/Integrals.jl

@@ -38,7 +38,7 @@ function scale_x(ub, lb, x)
 end
 
 function v_inf(t)
-    return t ./ (1 .- t .^ 2)
+    return map(t -> t / (1 - t^2), t)
 end
 
 function v_semiinf(t, a, upto_inf)
@@ -76,10 +76,13 @@ function transform_inf(t, p, f, lb, ub)
     semilw = lbb .& .!ubb
 
     function v(t)
-        return t .* _none + v_inf(t) .* _inf + v_semiinf(t, lb, 1) .* semiup +
-               v_semiinf(t, ub, 0) .* semilw
+        t .* _none + v_inf(t) .* _inf + v_semiinf(t, lb, 1) .* semiup +
+        v_semiinf(t, ub, 0) .* semilw
     end
-    jac = Zygote.@ignore ForwardDiff.jacobian(x -> v(x), t)
+    jac = ChainRulesCore.@ignore_derivatives ForwardDiff.jacobian(x -> v(x),
+                                                                  t |> Vector)::Matrix{
+                                                                                       eltype(t)
+                                                                                       }
     j = det(jac)
     f(v(t), p) * (j)
 end
@@ -180,9 +183,9 @@ function __solvebp_call(prob::IntegralProblem, ::HCubatureJL, sensealg, lb, ub,
 
     if isinplace(prob)
         dx = zeros(prob.nout)
-        f = (x) -> (prob.f(dx, x, p); dx)
+        f = x -> (prob.f(dx, x, prob.p); dx)
     else
-        f = (x) -> prob.f(x, p)
+        f = x -> prob.f(x, prob.p)
     end
     @assert prob.batch == 0
 
@@ -207,16 +210,16 @@ function __solvebp_call(prob::IntegralProblem, alg::VEGAS, sensealg, lb, ub, p,
     if prob.batch == 0
         if isinplace(prob)
             dx = zeros(prob.nout)
-            f = (x) -> (prob.f(dx, x, p); dx)
+            f = x -> (f(dx, x, p); dx)
         else
-            f = (x) -> prob.f(x, p)
+            f = x -> prob.f(x, prob.p)
         end
     else
         if isinplace(prob)
             dx = zeros(prob.batch)
-            f = (x) -> (prob.f(dx, x', p); dx)
+            f = x -> (f(dx, x', p); dx)
         else
-            f = (x) -> prob.f(x', p)
+            f = x -> f(x', p)
         end
     end
     val, err, chi = vegas(f, lb, ub, rtol = reltol, atol = abstol,
@@ -232,7 +235,7 @@ function ChainRulesCore.rrule(::typeof(__solvebp), prob, alg, sensealg, lb, ub,
         y = typeof(Δ) <: Array{<:Number, 0} ? Δ[1] : Δ
         if isinplace(prob)
             dx = zeros(prob.nout)
-            _f = (x) -> prob.f(dx, x, p)
+            _f = x -> prob.f(dx, x, p)
             if sensealg.vjp isa ZygoteVJP
                 dfdp = function (dx, x, p)
                     _, back = Zygote.pullback(p) do p
@@ -252,7 +255,7 @@ function ChainRulesCore.rrule(::typeof(__solvebp), prob, alg, sensealg, lb, ub,
                 error("TODO")
             end
         else
-            _f = (x) -> prob.f(x, p)
+            _f = x -> prob.f(x, p)
             if sensealg.vjp isa ZygoteVJP
                 if prob.batch > 0
                     dfdp = function (x, p)

test/inf_integral_tests.jl

@@ -1,4 +1,4 @@
-using Integrals, Distributions, Test
+using Integrals, Distributions, Test, StaticArrays
 
 μ = [0.00, 0.00]
 Σ = [0.4 0.0; 0.00 0.4]
@@ -35,3 +35,19 @@ prob = IntegralProblem(f, 0, Inf)
 sol = solve(prob, HCubatureJL(), reltol = 1e-3, abstol = 1e-3)
 @test (pi / 2 - sol.u)^2 < 1e-6
 @test_nowarn @inferred Integrals.transformation_if_inf(prob, Val(true))
+
+# Type stability
+μ = [0.00, 0.00]
+Σ = [0.4 0.0; 0.00 0.4]
+d = MvNormal(μ, Σ)
+m2 = let d = d
+    (x, p) -> pdf(d, x)
+end
+
+prob = IntegralProblem(m2, SVector(-Inf, -Inf), SVector(Inf, Inf))
+@test_nowarn @inferred solve(prob, HCubatureJL(); do_inf_transformation = Val(true))
+
+prob = @test_nowarn @inferred Integrals.transformation_if_inf(prob, Val(true))
+@test_nowarn @inferred Integrals.__solvebp_call(prob, HCubatureJL(),
+                                                Integrals.ReCallVJP(Integrals.ZygoteVJP()),
+                                                prob.lb, prob.ub, prob.p)