aquakin.dgsm#

aquakin.dgsm(fn, ranges, *, input_names=None, output_names=None, n_samples=64, seed=0, diff=DifferentiationConfig(mode='reverse', method='stable', check_finite=True, adjoint_max_steps=100000, adjoint_low_memory=False), batched=True)[source]#

Derivative-based global sensitivity measure via autodiff + Sobol QMC.

Estimates, for each uncertain input z_j,

nu_j = E_z[ (d fn / d z_j)^2 ]

by averaging the squared partial derivative over scrambled-Sobol quasi-random points in the input ranges. nu_j bounds the Sobol total-order index (see DGSMResult.sobol_total_bound), so it is the AD analogue of a variance-based Sobol total index, obtained from derivatives rather than a variance decomposition.

The derivatives are exact (no finite-difference truncation) and reuse the differentiable model, so the same machinery serves the calibration and identifiability analyses. The cost depends on ad_mode and on the number of outputs m and inputs d:

  • ad_mode="reverse" (default) forms the per-sample sensitivities with m reverse-mode passes (one per output), each independent of d. Best when there are few outputs relative to inputs and the adjoint is cheap. Works with any reactor adjoint.

  • ad_mode="forward" forms them with d forward-mode tangents pushed through a single solve, independent of m. Best when there are many outputs, or when the reverse adjoint is expensive – e.g. a stiff solve whose differentiated step must be capped (dtmax), which inflates the reverse pass. The reactor inside ``fn`` must then be built with adjoint=aquakin.forward_adjoint() (dgsm cannot set the adjoint for you, because fn constructs the reactor): the default RecursiveCheckpointAdjoint registers a custom_vjp that rejects forward-mode autodiff.

Both modes return identical sensitivities (to machine precision); ad_mode is purely a performance choice. For a single scalar output reverse is almost always cheaper; the forward advantage appears for multi-output screening of a stiff model.

Parameters:

  • fn (callable) – Maps an input vector (shape (d,)) to either a scalar JAX value or a vector of m outputs (shape (m,)). Must be jax-differentiable in the requested mode. For a reactor study, fn typically maps the uncertain inputs into a parameter vector / initial state, calls reactor.solve and reduces the solution to the output(s). If the model is stiff, build the reactor with a suitable dtmax so the differentiated solve stays finite. dgsm does not own the solve (your fn builds the reactor and chooses the t_eval), so it cannot apply a time_unit conversion for you: any t_eval / t_span inside fn must be in the model’s native time unit, or fn must pass time_unit= to its own reactor.solve call.

  • ranges (array-like, shape (d, 2)) – [lower, upper] bound for each input; sampling is uniform within.

  • input_names (list[str], optional) – Names for reporting; defaults to ["z0", "z1", ...].

  • output_names (list[str], optional) – Names for the m outputs when fn is vector-valued; defaults to ["output0", ...]. Ignored for a scalar fn.

  • n_samples (int, optional) – Target number of quasi-random points; rounded to the nearest power of two (Sobol sequences are balanced at powers of two). Increase until std_error is small relative to the ranking gaps.

  • seed (int, optional) – Seed for the scrambled-Sobol sampler. Fixing it (the default 0) makes the analysis exactly reproducible.

  • diff (DifferentiationConfig, optional) – Autodiff configuration. mode ({“reverse”, “forward”}) selects the direction used to form the per-sample sensitivities (see above).

  • batched (bool, optional) – When True (default) the whole sample is pushed through one jax.vmap dispatch and finiteness is filtered once on the stacked result – one device->host transfer instead of one per point. Set False to evaluate point-by-point (lower peak memory for a large screen). Both give identical results.

Returns:

A single DGSMResult when fn is scalar-valued, or a list of results (one per output, in order, each carrying its output_name) when fn is vector-valued.

Return type:

DGSMResult or list[DGSMResult]

Examples

>>> def fn(z):                       # output sensitive to z0, not z1
...     return 3.0 * z[0] + 0.0 * z[1]
>>> res = aquakin.dgsm(fn, [(0.0, 1.0), (0.0, 1.0)], input_names=["a", "b"])
>>> res.ranked()[0][0]
'a'
On this page
Show Source