Multidimentional integrals in julia: speed and automatic differentiation

[misha_mikhasenko]

I would like to differentiate through a two-dimensional integral.

• seems to be not possible for c++ wrappers (e.g. with the great Cuba.jl)
• Julia native integrators are way slower (HCubature.jl, NIntegration.jl)
• nested QuadGK.jl is also not fast enough.

Therefore, two questions:

1. Are there tricks to differentiate Cuba?
2. Is there a way to speed up Julia native implementations?

Here is MWE and tests

using BenchmarkTools

using HCubature
using NIntegration
using QuadGK
using Cuba

function model(cosθ,ϕ; pars)
    α1 = α2 = pars[1]-1
    v = (exp(-α1*(cosθ+1)) + exp(-α2*(1-cosθ)))*(1+sin(ϕ))^(1.2)
    return abs2(v)
end
#
# using Plots
# plot(cosθ -> model(cosθ,π/4; pars = [2.3]), -1, 1)
#
I_quadgk(pars) = quadgk(cosθ->quadgk(ϕ->model(cosθ,ϕ;pars=pars),-π,π)[1],-1.0, 1.0)[1]
I_nintegration(pars) = nintegrate((cosθ,ϕ,z)->model(cosθ,ϕ;pars=pars), (-1.0, -π, 0.0), (1.0, π, 1.0))[1]
I_hcubature(pars) =  hcubature(x->model(x[1],x[2]; pars= pars), [-1.0, -π], [1.0, π])[1]
I_cuba(pars) = cuhre((x,f)->f[1]=model(2x[1]-1, π*(2x[2]-1);pars=pars),2,1).integral[1]*(4π)
#
@btime I_cuba([2.3])
@btime I_quadgk([2.3])
@btime I_hcubature([2.3])
@btime I_nintegration([2.3])

Results are:

1) Cuba:          72.500 μs (1305 allocations: 61.28 KiB)
2) 2x QuadGK:    160.800 μs (827 allocations: 19.44 KiB)
3) NIntegration: 303.800 μs (209 allocations: 9.91 KiB)
4) HCubature:      2.215 ms (23787 allocations: 850.86 KiB)

I found it quite surprising with respect to the number of github stars on the packages

EDIT: replaced parameters to a vector.

[misha_mikhasenko]

Pin the experts: @stevengj, @giordano, @pabloferz.
Any thoughts are welcome!

[simeonschaub]

Shouldn’t it be fairly straightforward to just use an analytical derivative here? If you’re using this inside other code using e.g. Zygote, you could also write a custom adjoint for Cuba.jl.

[misha_mikhasenko]

My working problem is more complicated than that.
However, trying Zygote is an interesting idea. Could maybe sketch an example?
Let’s say, the parameter is a vector of dim2. Thanks in advance!

[ChrisRackauckas]

@simeonschaub

Shouldn’t it be fairly straightforward to just use an analytical derivative here? If you’re using this inside other code using e.g. Zygote, you could also write a custom adjoint for Cuba.jl.

Yes. FWIW, we made Quadrature.jl parameter explicit to make it easy to write the adjoint, I just haven’t gotten to it yet. It would be fairly straightforward to do though.

[simeonschaub]

With Zygote, you could solve it like this, by differentiating through the integrand:

using Cuba, Zygote

function model(cosθ,ϕ; αp1 = 2.3)
    α1 = α2 = αp1-1.0
    v = (exp(-α1*(cosθ+1)) + exp(-α2*(1-cosθ)))*(1+sin(ϕ))^(1.2)
    return abs2(v)
end

integrand(x, αp1) = model(2x[1]-1, π*(2x[2]-1);αp1=αp1)
I_cuba(αp1) = cuhre((x, f) -> f[1]=integrand(x, αp1),2,1).integral[1]*(4π)
Zygote.@adjoint I_cuba(αp1) = I_cuba(αp1), Δ -> (Δ * cuhre((x, f) -> f[1]=Zygote.pullback(a -> integrand(x, a), αp1)[2](1)[1], 2, 1).integral[1] * (4π),)

And then get the gradient:

julia> Zygote.gradient(x -> I_cuba(x[1]), [2.3])
([-13.368478696191737],)

It should be fairly easy to replace αp1 with a vector if you have multiple parameters.

[simeonschaub]

If you only a few parameters, it might be also worth considering using ForwardDiff for the integrand instead of Zygote.pullback, since that will then probably be faster.

[misha_mikhasenko]

Nice, I will try that. Thanks a lot

[stevengj]

Regarding the benchmark — it isn’t an entirely fair comparison, since you probably aren’t computing the integrals to the same accuracy in the different routines. For example, in NIntegration I think the default relative error is 1e-6, whereas in HCubature the default is sqrt(ε) ≈ 1.5e-8

(Even if you set the same error tolerances in all the approaches, because they estimate errors in different ways you might get substantially different accuracy. One way to compare would be to compute the “exact” integral using a very low tolerance in one of the routines, and then check out what tolerance is needed in each routine to roughly obtain a given relative error, e.g. 1e-6.)

That being said, I suspect that there may have been a performance regression in HCubature — it didn’t use to require so many allocations if I recall correctly (it uses StaticArrays internally so its temporary arrays are not supposed to be heap-allocated.) I’ll have to look into it.

Since you have a smooth integrand in 2d, you might also want to try the pcubature function from Cubature.jl.

[misha_mikhasenko]

A good point, I have not thought about accuracy, indeed.
I tend to assume pragmatically that the returned value is exact, but of course it is not the case.

Thank you for suggestion trying Cubature, I was puzzled by the difference between Cubature and HCubature.