Skip to content

Commit

Permalink
Merge pull request #241 from lxvm/infstability
Browse files Browse the repository at this point in the history 
Make frontend type stable
  • Loading branch information
@ChrisRackauckas
ChrisRackauckas authored Mar 3, 2024
2 parents 278d70f + e6e9a5b commit 973a5e4
Show file tree
Hide file tree
Showing 21 changed files with 524 additions and 367 deletions.
6 changes: 3 additions & 3 deletions Project.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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"
Expand All @@ -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"]
4 changes: 0 additions & 4 deletions docs/src/basics/solve.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
8 changes: 3 additions & 5 deletions ext/IntegralsArblibExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
73 changes: 38 additions & 35 deletions ext/IntegralsCubaExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 ||
Expand All @@ -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))
Expand All @@ -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

Expand Down
102 changes: 50 additions & 52 deletions ext/IntegralsCubatureExt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading

0 comments on commit 973a5e4

Please sign in to comment.