quadax.romberg

Romberg integration of a callable function or method.

Returns the integral of fun (a function of one variable) over interval.

Refines a uniform mesh over the whole interval and accelerates the sequence of trapezoidal sums by Richardson extrapolation. Suited to smooth integrands on a finite interval, where the extrapolation is worth orders of magnitude and the cost is competitive with the adaptive routines. Often has less compile and dispatch overhead compared to the locally adaptive routines which can mean lower wall clock time for cheap integrands.

Not recommended for infinite intervals or non-smooth integrands or those with other localized features. The uniform mesh cannot refine towards a difficulty, so an integrand with a local feature pays for the entire interval to resolve one point. For these cases quadgk() is preferred.

Parameters:

fun (callable) – Function to integrate, should have a signature of the form fun(x, *args) -> float, Array. Should be JAX transformable.
interval (array-like) – Lower and upper limits of integration. Use np.inf to denote infinite intervals. Its dtype sets the working precision: the integrand is called with an x of this dtype, and the result follows it unless the integrand upcasts. Integer types or python floats fall back to the JAX default. Must be real; complex integrands are supported, complex limits are not.
args (tuple) – additional arguments passed to fun
full_output (bool, optional) – If True, return the full state of the integrator. See below for more information.
epsabs (float) – Absolute and relative tolerances. The algorithm terminates once its estimate of the error in the current approximation I falls below max(epsabs, epsrel*|I|), which takes at least two refinement levels whatever the tolerance and divmin, since the first level’s estimate has no earlier one to be judged against. Default is the square root of the machine precision of the working dtype, ie of interval, or of the integrand’s own dtype if that is the coarser of the two.
epsrel (float) – Absolute and relative tolerances. The algorithm terminates once its estimate of the error in the current approximation I falls below max(epsabs, epsrel*|I|), which takes at least two refinement levels whatever the tolerance and divmin, since the first level’s estimate has no earlier one to be judged against. Default is the square root of the machine precision of the working dtype, ie of interval, or of the integrand’s own dtype if that is the coarser of the two.
divmax (int, optional) – Maximum order of extrapolation. Default is 20. Total number of function evaluations will be at most 2**divmax + 1
norm (int, callable) – Norm to use for measuring error for vector valued integrands. No effect if the integrand is scalar valued. If an int, uses p-norm of the given order, otherwise should be callable.
extrapolate (bool, optional) – Whether to accelerate convergence by Richardson extrapolation, which is what makes this Romberg’s method rather than plain repeated bisection. On by default, and worth leaving on: it is where nearly all of this routine’s accuracy comes from, buying several orders of magnitude on a smooth integrand from the same nodes. Turning it off leaves the same nodes and the same halving schedule, reading the un-extrapolated trapezoidal sum instead. That is the more conservative reading where the integrand is not smooth enough for the error expansion to hold, but on such an integrand this routine is the wrong choice anyways.
adjoint (AbstractAdjoint, optional) – How to compute derivatives of the quadrature. Default is DirectAdjoint(), which gives the exact derivative of the discretized problem, and is the cheaper option for a cheap integrand. LeibnizAdjoint gives the derivative its own error control (ie, can better approximate the true continuous derivative), and is faster when the integrand is expensive or divmax is generous; see Adjoints for when that is worth paying for.
batch_size (int, optional) – Maximum number of points at which to evaluate the integrand in parallel. Defaults to 2**divmin, which is one batch for the whole starting grid and exactly one for the refinement level after it. Each refinement level doubles the number of new points, so raising this together with divmin is usually worth a lot on GPU/TPU, at the cost of peak memory scaling with it. Levels with fewer new points than one batch are padded up to a full batch, so a level costs batch_size evaluations however few points it places; that padding is what keeps a single batch shape traced for every level rather than one per level. Clipped to the largest number of points any one level places.
divmin (int, optional) –
Number of halvings the run starts from, default 4: it begins on a grid of 2**divmin intervals rather than working up to one. Mirrors divmax, and must not exceed it.

That starting grid contains every coarser grid of the halving sequence, so the coarser rows of the table are filled from the same evaluations and no extrapolation is lost - the table is the one a run from divmin=0 would have built by the time it reached the same mesh. What is bought is that the early work happens in one batch instead of a handful of levels evaluating a few points each, which is where the default schedule wastes time on GPU/TPU. What is paid is a floor of 2**divmin + 1 evaluations even on an integrand that needed fewer.
throw (bool, optional) – Whether to raise an error if the routine does not converge. If True, a run that terminates for any reason other than reaching the requested tolerance raises with the message its status carries. If False, the default, that status is reported on the returned info and left to the caller to act on.

Returns:

y (float, Array) – Approximation to the integral
info (QuadratureInfo) – Named tuple with the following fields:
- err : (float) Estimate of the error in the approximation. Built from how far the estimate moved over the last few refinement levels rather than over the last one alone, plus the tail that movement’s own contraction rate implies, and floored at the precision the integrand can be summed to.
- neval : (int) Total number of function evaluations.
- status : (int) Code for why the routine terminated, one of quadax.STATUS. STATUS.normal (0) means the requested tolerances were reached; every other code names a difficulty, whose message is print(quadax.STATUS[status]). Where a run meets more than one condition the most severe is reported.
- info : (dict or None) Other information returned by the algorithm. Only present if full_output is True. Contains the following:
  - table : (ndarray, size(divmax+1, divmax+1, …)) Estimate of the integral from each level of discretization and each step of extrapolation. With extrapolate=False only the first column is filled.

Notes

The number of new points a refinement level places is only known at run time, so there is no single shape to vectorize over. Integrand evaluations are made in fixed size batches of batch_size, defaulting to one at a time; raise it to get parallelism on GPU/TPU.