Merge pull request #241 from lxvm/infstability

Make frontend type stable

Author: ChrisRackauckas

Project.toml

@@ -9,6 +9,7 @@ HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MonteCarloIntegration = "4886b29c-78c9-11e9-0a6e-41e1f4161f7b"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 
@@ -46,8 +47,8 @@ HCubature = "1.5.2"
 LinearAlgebra = "1.10"
 MCIntegration = "0.4.2"
 MonteCarloIntegration = "0.2"
-Pkg = "1.10"
 QuadGK = "2.9"
+Random = "1.10"
 Reexport = "1.0"
 SafeTestsets = "0.1"
 SciMLBase = "2.24"
@@ -67,11 +68,10 @@ FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MCIntegration = "ea1e2de9-7db7-4b42-91ee-0cd1bf6df167"
-Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Arblib", "StaticArrays", "FiniteDiff", "Pkg", "SafeTestsets", "Test", "Distributions", "ForwardDiff", "Zygote", "ChainRulesCore", "FastGaussQuadrature", "Cuba", "Cubature", "MCIntegration"]
+test = ["Aqua", "Arblib", "StaticArrays", "FiniteDiff", "SafeTestsets", "Test", "Distributions", "ForwardDiff", "Zygote", "ChainRulesCore", "FastGaussQuadrature", "Cuba", "Cubature", "MCIntegration"]

docs/src/basics/solve.md

@@ -3,7 +3,3 @@
 ```@docs
 solve(prob::IntegralProblem, alg::SciMLBase.AbstractIntegralAlgorithm)
 ```
-
-Additionally, the extra keyword arguments are splatted to the library calls, so
-see the documentation of the integrator library for all the extra details.
-These extra keyword arguments are not guaranteed to act uniformly.

ext/IntegralsArblibExt.jl

@@ -15,21 +15,19 @@ function Integrals.__solvebp_call(
 
     if isinplace(prob)
         res = Acb(0)
-        y_ = similar(prob.f.integrand_prototype, typeof(res))
+        @assert res isa eltype(prob.f.integrand_prototype) "Arblib require inplace prototypes to store Acb elements"
+        y_ = similar(prob.f.integrand_prototype)
         f_ = (y, x; kws...) -> (prob.f(y_, x, p; kws...); Arblib.set!(y, only(y_)))
         val = Arblib.integrate!(f_, res, lb, ub, atol = abstol, rtol = reltol,
             check_analytic = alg.check_analytic, take_prec = alg.take_prec,
             warn_on_no_convergence = alg.warn_on_no_convergence, opts = alg.opts)
-        SciMLBase.build_solution(
-            prob, alg, val, get_radius(val), retcode = ReturnCode.Success)
     else
         f_ = (x; kws...) -> only(prob.f(x, p; kws...))
         val = Arblib.integrate(f_, lb, ub, atol = abstol, rtol = reltol,
             check_analytic = alg.check_analytic, take_prec = alg.take_prec,
             warn_on_no_convergence = alg.warn_on_no_convergence, opts = alg.opts)
-        SciMLBase.build_solution(
-            prob, alg, val, get_radius(val), retcode = ReturnCode.Success)
     end
+    SciMLBase.build_solution(prob, alg, val, get_radius(val), retcode = ReturnCode.Success)
 end
 
 function get_radius(ball)

ext/IntegralsCubaExt.jl

@@ -18,8 +18,12 @@ function Integrals.__solvebp_call(prob::IntegralProblem, alg::AbstractCubaAlgori
         throw(ArgumentError("Cuba.jl only supports real-valued integrands"))
     # we could support other types by multiplying by the jacobian determinant at the end
 
-    if prob.f isa BatchIntegralFunction
-        nvec = min(maxiters, prob.f.max_batch)
+    f = prob.f
+    prototype = Integrals.get_prototype(prob)
+    if f isa BatchIntegralFunction
+        fsize = size(prototype)[begin:(end - 1)]
+        ncomp = prod(fsize)
+        nvec = min(maxiters, f.max_batch)
         # nvec == 1 in Cuba will change vectors to matrices, so we won't support it when
         # batching
         nvec > 1 ||
@@ -33,24 +37,21 @@ function Integrals.__solvebp_call(prob::IntegralProblem, alg::AbstractCubaAlgori
             scale = x -> scale_x!(view(_x, :, 1:size(x, 2)), ub, lb, x)
         end
 
-        if isinplace(prob)
-            fsize = size(prob.f.integrand_prototype)[begin:(end - 1)]
-            y = similar(prob.f.integrand_prototype, fsize..., nvec)
-            ax = map(_ -> (:), fsize)
-            f = function (x, dx)
-                dy = @view(y[ax..., begin:(begin + size(dx, 2) - 1)])
-                prob.f(dy, scale(x), p)
-                dx .= reshape(dy, :, size(dx, 2)) .* vol
+        if isinplace(f)
+            ax = ntuple(_ -> (:), length(fsize))
+            _f = let y_ = similar(prototype, fsize..., nvec)
+                function (u, _y)
+                    y = @view(y_[ax..., begin:(begin + size(_y, 2) - 1)])
+                    f(y, scale(u), p)
+                    _y .= reshape(y, size(_y)) .* vol
+                end
             end
         else
-            y = mid isa Number ? prob.f(typeof(mid)[], p) :
-                prob.f(Matrix{typeof(mid)}(undef, length(mid), 0), p)
-            fsize = size(y)[begin:(end - 1)]
-            f = (x, dx) -> dx .= reshape(prob.f(scale(x), p), :, size(dx, 2)) .* vol
+            _f = (u, y) -> y .= reshape(f(scale(u), p), size(y)) .* vol
         end
-        ncomp = prod(fsize)
     else
         nvec = 1
+        ncomp = length(prototype)
 
         if mid isa Real
             scale = x -> scale_x(ub, lb, only(x))
@@ -59,58 +60,60 @@ function Integrals.__solvebp_call(prob::IntegralProblem, alg::AbstractCubaAlgori
             scale = x -> scale_x!(_x, ub, lb, x)
         end
 
-        if isinplace(prob)
-            y = similar(prob.f.integrand_prototype)
-            f = (x, dx) -> dx .= vec(prob.f(y, scale(x), p)) .* vol
+        if isinplace(f)
+            _f = let y = similar(prototype)
+                (u, _y) -> begin
+                    f(y, scale(u), p)
+                    _y .= vec(y) .* vol
+                end
+            end
         else
-            y = prob.f(mid, p)
-            f = (x, dx) -> dx .= Iterators.flatten(prob.f(scale(x), p)) .* vol
+            _f = (u, y) -> y .= Iterators.flatten(f(scale(u), p)) .* vol
         end
-        ncomp = length(y)
     end
 
-    if alg isa CubaVegas
-        out = Cuba.vegas(f, ndim, ncomp; rtol = reltol,
+    out = if alg isa CubaVegas
+        Cuba.vegas(_f, ndim, ncomp; rtol = reltol,
             atol = abstol, nvec = nvec,
             maxevals = maxiters,
             flags = alg.flags, seed = alg.seed, minevals = alg.minevals,
             nstart = alg.nstart, nincrease = alg.nincrease,
             gridno = alg.gridno)
     elseif alg isa CubaSUAVE
-        out = Cuba.suave(f, ndim, ncomp; rtol = reltol,
+        Cuba.suave(_f, ndim, ncomp; rtol = reltol,
             atol = abstol, nvec = nvec,
             maxevals = maxiters,
             flags = alg.flags, seed = alg.seed, minevals = alg.minevals,
             nnew = alg.nnew, nmin = alg.nmin, flatness = alg.flatness)
     elseif alg isa CubaDivonne
-        out = Cuba.divonne(f, ndim, ncomp; rtol = reltol,
+        Cuba.divonne(_f, ndim, ncomp; rtol = reltol,
             atol = abstol, nvec = nvec,
             maxevals = maxiters,
             flags = alg.flags, seed = alg.seed, minevals = alg.minevals,
             key1 = alg.key1, key2 = alg.key2, key3 = alg.key3,
             maxpass = alg.maxpass, border = alg.border,
             maxchisq = alg.maxchisq, mindeviation = alg.mindeviation)
     elseif alg isa CubaCuhre
-        out = Cuba.cuhre(f, ndim, ncomp; rtol = reltol,
+        Cuba.cuhre(_f, ndim, ncomp; rtol = reltol,
             atol = abstol, nvec = nvec,
             maxevals = maxiters,
             flags = alg.flags, minevals = alg.minevals, key = alg.key)
     end
 
     # out.integral is a Vector{Float64}, but we want to return it to the shape of the integrand
-    if prob.f isa BatchIntegralFunction
-        if y isa AbstractVector
-            val = out.integral[1]
+    val = if f isa BatchIntegralFunction
+        if prototype isa AbstractVector
+            out.integral[1]
         else
-            val = reshape(out.integral, fsize)
+            reshape(out.integral, fsize)
         end
     else
-        if y isa Real
-            val = out.integral[1]
-        elseif y isa AbstractVector
-            val = out.integral
+        if prototype isa Real
+            out.integral[1]
+        elseif prototype isa AbstractVector
+            out.integral
         else
-            val = reshape(out.integral, size(y))
+            reshape(out.integral, size(prototype))
         end
     end
 

ext/IntegralsCubatureExt.jl

@@ -13,143 +13,141 @@ function Integrals.__solvebp_call(prob::IntegralProblem,
     mid = (lb + ub) / 2
 
     # we get to pick fdim or not based on the IntegralFunction and its output dimensions
-    y = if prob.f isa BatchIntegralFunction
-        isinplace(prob.f) ? prob.f.integrand_prototype :
-        mid isa Number ? prob.f(eltype(mid)[], p) :
-        prob.f(Matrix{eltype(mid)}(undef, length(mid), 0), p)
-    else
-        # we evaluate the oop function to decide whether the output should be vectorized
-        isinplace(prob.f) ? prob.f.integrand_prototype : prob.f(mid, p)
-    end
+    f = prob.f
+    prototype = Integrals.get_prototype(prob)
 
-    @assert eltype(y)<:Real "Cubature.jl is only compatible with real-valued integrands"
+    @assert eltype(prototype)<:Real "Cubature.jl is only compatible with real-valued integrands"
 
-    if prob.f isa BatchIntegralFunction
-        if y isa AbstractVector # this branch could be omitted since the following one should work similarly
-            if isinplace(prob)
+    if f isa BatchIntegralFunction
+        if prototype isa AbstractVector # this branch could be omitted since the following one should work similarly
+            if isinplace(f)
                 # dx is a Vector, but we provide the integrand a vector of the same type as
                 # y, which needs to be resized since the number of batch points changes.
-                dy = similar(y)
-                f = (x, dx) -> begin
-                    resize!(dy, length(dx))
-                    prob.f(dy, x, p)
-                    dx .= dy
+                _f = let y = similar(prototype)
+                    (u, v) -> begin
+                        resize!(y, length(v))
+                        f(y, u, p)
+                        v .= y
+                    end
                 end
             else
-                f = (x, dx) -> (dx .= prob.f(x, p))
+                _f = (u, v) -> (v .= f(u, p))
             end
             if mid isa Number
                 if alg isa CubatureJLh
-                    val, err = Cubature.hquadrature_v(f, lb, ub;
+                    val, err = Cubature.hquadrature_v(_f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters)
                 else
-                    val, err = Cubature.pquadrature_v(f, lb, ub;
+                    val, err = Cubature.pquadrature_v(_f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters)
                 end
             else
                 if alg isa CubatureJLh
-                    val, err = Cubature.hcubature_v(f, lb, ub;
+                    val, err = Cubature.hcubature_v(_f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters)
                 else
-                    val, err = Cubature.pcubature_v(f, lb, ub;
+                    val, err = Cubature.pcubature_v(_f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters)
                 end
             end
-        elseif y isa AbstractArray
-            bfsize = size(y)[begin:(end - 1)]
-            bfdim = prod(bfsize)
-            if isinplace(prob)
+        elseif prototype isa AbstractArray
+            fsize = size(prototype)[begin:(end - 1)]
+            fdim = prod(fsize)
+            if isinplace(f)
                 # dx is a Matrix, but to provide a buffer of the same type as y, we make
                 # would like to make views of a larger buffer, but CubatureJL doesn't set
                 # a hard limit for max_batch, so we allocate a new buffer with the needed size
-                f = (x, dx) -> begin
-                    dy = similar(y, bfsize..., size(dx, 2))
-                    prob.f(dy, x, p)
-                    dx .= reshape(dy, bfdim, size(dx, 2))
+                _f = let fsize = fsize
+                    (u, v) -> begin
+                        y = similar(prototype, fsize..., size(v, 2))
+                        f(y, u, p)
+                        v .= reshape(y, fdim, size(v, 2))
+                    end
                 end
             else
-                f = (x, dx) -> (dx .= reshape(prob.f(x, p), bfdim, size(dx, 2)))
+                _f = (u, v) -> (v .= reshape(f(u, p), fdim, size(v, 2)))
             end
             if mid isa Number
                 if alg isa CubatureJLh
-                    val_, err = Cubature.hquadrature_v(bfdim, f, lb, ub;
+                    val_, err = Cubature.hquadrature_v(fdim, _f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters, error_norm = alg.error_norm)
                 else
-                    val_, err = Cubature.pquadrature_v(bfdim, f, lb, ub;
+                    val_, err = Cubature.pquadrature_v(fdim, _f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters, error_norm = alg.error_norm)
                 end
             else
                 if alg isa CubatureJLh
-                    val_, err = Cubature.hcubature_v(bfdim, f, lb, ub;
+                    val_, err = Cubature.hcubature_v(fdim, _f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters, error_norm = alg.error_norm)
                 else
-                    val_, err = Cubature.pcubature_v(bfdim, f, lb, ub;
+                    val_, err = Cubature.pcubature_v(fdim, _f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters, error_norm = alg.error_norm)
                 end
             end
-            val = reshape(val_, bfsize...)
+            val = reshape(val_, fsize...)
         else
             error("BatchIntegralFunction integrands must be arrays for Cubature.jl")
         end
     else
-        if y isa Real
+        if prototype isa Real
             # no inplace in this case, since the integrand_prototype would be mutable
-            f = x -> prob.f(x, p)
+            _f = u -> f(u, p)
             if lb isa Number
                 if alg isa CubatureJLh
-                    val, err = Cubature.hquadrature(f, lb, ub;
+                    val, err = Cubature.hquadrature(_f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters)
                 else
-                    val, err = Cubature.pquadrature(f, lb, ub;
+                    val, err = Cubature.pquadrature(_f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters)
                 end
             else
                 if alg isa CubatureJLh
-                    val, err = Cubature.hcubature(f, lb, ub;
+                    val, err = Cubature.hcubature(_f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters)
                 else
-                    val, err = Cubature.pcubature(f, lb, ub;
+                    val, err = Cubature.pcubature(_f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters)
                 end
             end
-        elseif y isa AbstractArray
-            fsize = size(y)
-            fdim = length(y)
+        elseif prototype isa AbstractArray
+            fsize = size(prototype)
+            fdim = length(prototype)
             if isinplace(prob)
-                dy = similar(y)
-                f = (x, v) -> (prob.f(dy, x, p); v .= vec(dy))
+                _f = let y = similar(prototype)
+                    (u, v) -> (f(y, u, p); v .= vec(y))
+                end
             else
-                f = (x, v) -> (v .= vec(prob.f(x, p)))
+                _f = (u, v) -> (v .= vec(f(u, p)))
             end
             if mid isa Number
                 if alg isa CubatureJLh
-                    val_, err = Cubature.hquadrature(fdim, f, lb, ub;
+                    val_, err = Cubature.hquadrature(fdim, _f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters, error_norm = alg.error_norm)
                 else
-                    val_, err = Cubature.pquadrature(fdim, f, lb, ub;
+                    val_, err = Cubature.pquadrature(fdim, _f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters, error_norm = alg.error_norm)
                 end
             else
                 if alg isa CubatureJLh
-                    val_, err = Cubature.hcubature(fdim, f, lb, ub;
+                    val_, err = Cubature.hcubature(fdim, _f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters, error_norm = alg.error_norm)
                 else
-                    val_, err = Cubature.pcubature(fdim, f, lb, ub;
+                    val_, err = Cubature.pcubature(fdim, _f, lb, ub;
                         reltol = reltol, abstol = abstol,
                         maxevals = maxiters, error_norm = alg.error_norm)
                 end