1+ module IntegralsCubaExt
2+
3+ using Integrals, Cuba
4+ import Integrals: transformation_if_inf, scale_x, scale_x!, CubaVegas, AbstractCubaAlgorithm,
5+ CubaSUAVE, CubaDivonne, CubaCuhre
6+
7+ function Integrals. __solvebp_call (prob:: IntegralProblem , alg:: AbstractCubaAlgorithm ,
8+ sensealg,
9+ domain, p;
10+ reltol = 1e-8 , abstol = 1e-8 ,
11+ maxiters = alg isa CubaSUAVE ? 1000000 : typemax (Int))
12+ @assert maxiters>= 1000 " maxiters for $alg should be larger than 1000"
13+ lb, ub = domain
14+ mid = (lb + ub) / 2
15+ ndim = length (mid)
16+ (vol = prod (map (- , ub, lb))) isa Real ||
17+ throw (ArgumentError (" Cuba.jl only supports real-valued integrands"
18+ # we could support other types by multiplying by the jacobian determinant at the end
19+
20+ if prob. f isa BatchIntegralFunction
21+ # nvec == 1 in Cuba will change vectors to matrices, so we won't support it when
22+ # batching
23+ (nvec = prob. f. max_batch) > 1 ||
24+ throw (ArgumentError (" BatchIntegralFunction must take multiple batch points"
25+
26+ if mid isa Real
27+ _x = zeros (typeof (mid), prob. f. max_batch)
28+ scale = x -> scale_x! (resize! (_x, length (x)), ub, lb, vec (x))
29+ else
30+ _x = zeros (eltype (mid), length (mid), prob. f. max_batch)
31+ scale = x -> scale_x! (view (_x, :, 1 : size (x, 2 )), ub, lb, x)
32+ end
33+
34+ if isinplace (prob)
35+ fsize = size (prob. f. integrand_prototype)[begin : (end - 1 )]
36+ y = similar (prob. f. integrand_prototype, fsize... , nvec)
37+ ax = map (_ -> (:), fsize)
38+ f = function (x, dx)
39+ dy = @view (y[ax... , begin : (begin + size (dx, 2 ) - 1 )])
40+ prob. f (dy, scale (x), p)
41+ dx .= reshape (dy, :, size (dx, 2 )) .* vol
42+ end
43+ else
44+ y = mid isa Number ? prob. f (typeof (mid)[], p) :
45+ prob. f (Matrix {typeof(mid)} (undef, length (mid), 0 ), p)
46+ fsize = size (y)[begin : (end - 1 )]
47+ f = (x, dx) -> dx .= reshape (prob. f (scale (x), p), :, size (dx, 2 )) .* vol
48+ end
49+ ncomp = prod (fsize)
50+ else
51+ nvec = 1
52+
53+ if mid isa Real
54+ scale = x -> scale_x (ub, lb, only (x))
55+ else
56+ _x = zeros (eltype (mid), length (mid))
57+ scale = x -> scale_x! (_x, ub, lb, x)
58+ end
59+
60+ if isinplace (prob)
61+ y = similar (prob. f. integrand_prototype)
62+ f = (x, dx) -> dx .= vec (prob. f (y, scale (x), p)) .* vol
63+ else
64+ y = prob. f (mid, p)
65+ f = (x, dx) -> dx .= Iterators. flatten (prob. f (scale (x), p)) .* vol
66+ end
67+ ncomp = length (y)
68+ end
69+
70+ if alg isa CubaVegas
71+ out = Cuba. vegas (f, ndim, ncomp; rtol = reltol,
72+ atol = abstol, nvec = nvec,
73+ maxevals = maxiters,
74+ flags = alg. flags, seed = alg. seed, minevals = alg. minevals,
75+ nstart = alg. nstart, nincrease = alg. nincrease,
76+ gridno = alg. gridno)
77+ elseif alg isa CubaSUAVE
78+ out = Cuba. suave (f, ndim, ncomp; rtol = reltol,
79+ atol = abstol, nvec = nvec,
80+ maxevals = maxiters,
81+ flags = alg. flags, seed = alg. seed, minevals = alg. minevals,
82+ nnew = alg. nnew, nmin = alg. nmin, flatness = alg. flatness)
83+ elseif alg isa CubaDivonne
84+ out = Cuba. divonne (f, ndim, ncomp; rtol = reltol,
85+ atol = abstol, nvec = nvec,
86+ maxevals = maxiters,
87+ flags = alg. flags, seed = alg. seed, minevals = alg. minevals,
88+ key1 = alg. key1, key2 = alg. key2, key3 = alg. key3,
89+ maxpass = alg. maxpass, border = alg. border,
90+ maxchisq = alg. maxchisq, mindeviation = alg. mindeviation)
91+ elseif alg isa CubaCuhre
92+ out = Cuba. cuhre (f, ndim, ncomp; rtol = reltol,
93+ atol = abstol, nvec = nvec,
94+ maxevals = maxiters,
95+ flags = alg. flags, minevals = alg. minevals, key = alg. key)
96+ end
97+
98+ # out.integral is a Vector{Float64}, but we want to return it to the shape of the integrand
99+ if prob. f isa BatchIntegralFunction
100+ if y isa AbstractVector
101+ val = out. integral[1 ]
102+ else
103+ val = reshape (out. integral, fsize)
104+ end
105+ else
106+ if y isa Real
107+ val = out. integral[1 ]
108+ elseif y isa AbstractVector
109+ val = out. integral
110+ else
111+ val = reshape (out. integral, size (y))
112+ end
113+ end
114+
115+ SciMLBase. build_solution (prob, alg, val, out. error,
116+ chi = out. probability, retcode = ReturnCode. Success)
117+ end
118+
119+ end
0 commit comments