Recommended formalism for double quad?

[nmearl]

I have a double integral I'd love to move to using quadax, but it seems that when I attempt to implement it, it crashes the GPU by running out of memory. It's a naive implementation, but I'm curious if there's some recommendation for how to approach this problem. My double quad takes the form e.g.

@jax.jit
def integrand(y, x, *args):
    ...
    return result

def inner_quad(x, y1, y2, *args):
    return quadgk(integrand, [y1, y2], args=(x, ...))

def outer_quad(x1, x2, y1, y2, *args):
    return quadgk(inner_quad, [x1, x2], args=(y1, y2, ...))

Any suggestions would be appreciated.

[f0uriest]

I think that should work, I had to change inner_quad to

return quadgk(integrand, [y1, y2], args=(x, ...))[0]

since we want that to just return the value of the integrand, not the aux info.

Not sure about the OOM issue, I think that would depend on how complicated integrand is and the order of both the inner and outer quads. If you can share a bit more about the error you're getting I can try to help.

[nmearl]

Thanks, and apologies for the late follow-up. It seems that the OOM issue only arises when I attempt to use the double integration (as above) with numpyro. The integrand is pretty involved, but it's essentially a spectrum generator of an accretion disk. A simplified model might be something like:

import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, HMC


@jax.jit
def integrand(r, phi, X, *args):
    ...
    return result


@jax.jit
def inner_quad(r, phi1, phi2, X, *args):
    return quadgk(integrand, [phi1, phi2], args=(X, *args))[0]


@jax.jit
def outer_quad(r1, r2, phi1, phi2, X, *args):
    return quadgk(inner_quad, [r1, r2], args=(phi1, phi2, X, *args))[0]


def disk_model(wave, flux, flux_err):
    r1 = numpyro.sample("disk_inner_radius", dist.LogUniform(1e2, 1e3))
    r2 = numpyro.sample("disk_outer_radius", dist.LogUniform(1e3, 1e5))

    phi1 = -jnp.pi * 0.5
    phi2 = jnp.pi * 0.5

    total_flux = outer_quad(r1, r1, phi1, phi2, wave)

    with numpyro.plate("data", wave.shape[0]):
        numpyro.sample("obs", dist.Normal(total_flux, flux_err), obs=flux)


wave, flux, flux_err = ...

nuts_kernel = NUTS(disk_model)
mcmc = MCMC(
    nuts_kernel,
    num_warmup=500,
    num_samples=2000,
)
mcmc.run(jax.random.PRNGKey(0), wave=wave, flux=flux, flux_err=flux_err)

r1 and r2 are the inner and outer radius of the disk, while phi1 and phi1 are the orientation angle relative to the observer. X are the wavelength bins at which the flux is to be calculated, and is a fixed-size array. The rest of the *args values are float parameters. In the prescription above, the fitting always crashes due to memory, which explodes up to ~50 GB. It can be obviated if I use fixed integration schemes (GaussKronrodRule), but the results are just not good in that case. It'd be great if we can get the double integration with the adaptive methods working.

[f0uriest]

One thought is that it could be due to quadax vectorizing the evaluation of the integrand at order points (default 21 for quadgk), which when you nest the quadrature means its actually evaluating it 21^2 = 441 times at each step. If that's the bottleneck, you could try switching to quadcc and setting order=4 or something (The lowest order I have implemented for quadgk is 15 while quadcc can be any arbitrary even order). If your integrand is the output of an MC simulation which may not be smooth you probably want to use a low order rule anyways.