quadax.LeibnizAdjoint

class quadax.LeibnizAdjoint(options: dict = <factory>, options_fwd: dict = <factory>, options_rev: dict = <factory>, options_face: dict = <factory>)Source

Differentiate by the Leibniz rule, either mode, with its own error control.

The derivative is evaluated with its own adaptive solve, so it gets its own error control rather than inheriting the subdivision chosen for the integral; see Adjoints for when that is worth paying for.

Because each mode picks its own subdivision, forward and reverse results agree to quadrature accuracy rather than exactly.

Parameters:

• options (dict, optional) – Options for the derivative’s own solve, replacing the ones the quadrature routine was called with. Default is to solve for the derivative with the same options as the integral.

• options_fwd (dict, optional) – The same, for one direction alone, taking precedence over options. Which vector is being measured differs between the two directions, so norm is the option most likely to want this; see the note below.

• options_rev (dict, optional) – The same, for one direction alone, taking precedence over options. Which vector is being measured differs between the two directions, so norm is the option most likely to want this; see the note below.

• options_face (dict, optional) – Options for the boundary term’s own solves, taking precedence over options and applying to both directions. Over a box the boundary term is an integral over a face rather than a value at a point, and these configure it; in one dimension a face is a point and this does nothing. There is deliberately no forward and reverse pair, see the note below. Default is to solve for the faces with the same options as the integral, and with a rule derived from the integral’s, one axis shorter. rule may be given here to override that, and must then be of dimension ndim - 1.

Notes

The two directions of the derivative do not integrate the same vector, which is what options_fwd and options_rev are for.

Forward mode integrates the tangent of the integrand, a vector of the integrand’s own shape. Reverse mode integrates the cotangent of the arguments being differentiated, raveled into a single flat vector. Its length is the total number of differentiated parameter components, which has nothing to do with the integrand’s shape, and its entries are parameters rather than components of the integral. They appear in the order the quadrature’s own arguments do: the limits, args, and whatever the integrand closes over. Each contributing size entries if it is being differentiated and none at all if it is not. Differentiating a two point interval and a length three args[0] gives a vector of five with the limits first; differentiating args[0] alone gives a vector of three.

A norm weighted per parameter is written against that layout, and belongs in options_rev so that it is never handed a tangent instead, which the two vectors being the same length by coincidence would otherwise hide.

The boundary term is the other way round, which is why options_face has no forward and reverse halves. It is built from the integrand at its primal values, the limits alone being differentiated, so both directions integrate the identical faces and a vector of the integrand’s own shape. For the same reason options_face is merged over the options the quadrature was called with and never over options_fwd or options_rev, so a norm written for the reverse direction’s parameter layout cannot reach a face.

Over a box the boundary term costs up to two adaptive solves per axis whose limits are being differentiated, each one dimension shorter than the integral, and each direction builds its own. Forward mode skips the faces of an axis that is not actually moving; reverse mode owes a cotangent for every limit of a live axis and so cannot. DirectAdjoint remains the cheaper default for cubature.

When differentiating a moving jump or singularity, mark the jump or singularity with a breakpoint, and build that breakpoint from the same parameter that positions the feature. Marking a feature that is genuinely there is never worse than leaving it unmarked, and for derivatives it is frequently the difference between a correct answer and a silently wrong one. The one way marking can hurt is marking something that is not there, which is the third case below.

Differentiating a jump gives a delta, which no quadrature of the integrand’s tangent can represent, so it is recovered from the motion of the breakpoint instead, and that works only when the breakpoint moves with the discontinuity:

step = lambda t, z: jnp.where(t > z[0], 1.0, 0.0)   # jumps at t = z[0]

# correct: one parameter `s` positions both the jump and the breakpoint
f = lambda s: quadgk(step, jnp.array([-1.0, s, 1.0]), (jnp.array([s]),))[0]
jax.grad(f)(0.3)            # -1.0

Two versions that look equivalent are not. Both return the same primal value, 0.7, and both give zero where the derivative is -1:

# WRONG: the jump is unmarked, so nothing tracks it
f = lambda s: quadgk(step, jnp.array([-1.0, 1.0]), (jnp.array([s]),))[0]

# WRONG: marked, but with a constant, which carries no derivative. Tempting,
# because it does help the primal value: it cuts the subdivision from 13
# sub-intervals to 2.
f = lambda s: quadgk(step, jnp.array([-1.0, 0.3, 1.0]), (jnp.array([s]),))[0]

Splitting the feature across two parameters is what the requirement rules out. The derivative with respect to a breakpoint on its own is not a well posed question: the integral is the same whatever the mesh is cut at, so the breakpoint only means anything in combination with the integrand it is marking. Ask for it anyway, by differentiating with respect to [breakpoint, jump location] at [0.3, 0.3], and the answer is [-1, 0]. The total over the two is right, and how it is divided between them is an artifact of having written one feature as two parameters, not a property of the quadrature.

A moving singularity needs the same treatment, for the same reason: unmarked, it slides across a subdivision that does not move with it.

Only a feature that is steep but continuous needs none of this: a sharp tanh differentiates correctly unmarked under either adjoint. Marking one anyway is harmless, the breakpoint then only helps the subdivision, and contributes nothing to the derivative since the integrand has the same limit from both sides of it.

Methods

quadrature(ops, interval, args, consts, ...)

Evaluate the quadrature, differentiating it by the Leibniz rule.

Attributes

options
options_fwd
options_rev
options_face