jaxgsa

Appearance

jaxgsaGlobal Sensitivity Analysis in JAX

Thirteen methods, one interface, eleven of them JIT-compiled. Scalar, multi-output, and time-series models.

Get Started
Choose a Method
API Reference

Thirteen methods, one interface

Every method has analyze(). The design-based ones also have sample(). Results share one output contract and convert to labeled xarray with to_dataset().

Bring your own design, or your old runs

Sobol, eFAST, Morris, and Kucherenko build the sample matrix for you. The other nine read indices off any (X, Y) pairs you already have, including a sweep you ran last year.

Time series in one pass

Pass Y as (N,), (N, K), or (N, T, K). One compiled call returns indices for every timestep and output, so you can watch a ranking change along the trajectory instead of averaging it away.

Gradients are cheap

DGSM bounds the total Sobol index from model derivatives, taken by JAX autodiff. It picks forward or reverse mode from the output shape, so time-series models do not pay T reverse passes.

Correlated inputs

Declare a Gaussian-copula matrix on the Problem. VKOGA and Kucherenko return variance-based indices under that dependence, split into correlated and uncorrelated parts.

Categorical inputs

Declare unordered levels with probs and labels. Sobol, Borgonovo delta, optimal transport, and PAWN handle them. The methods whose indices would depend on code order refuse instead.

It refuses rather than approximates

A method that cannot handle your problem raises a ValueError naming the parameters and the alternatives. A zero-variance output slice returns NaN with a warning that names the slice.

Fast where output size is large

Vectorized estimators replace SALib's per-slice Python loop. RS-HDMR on 50 timesteps by 6 outputs runs 1060x faster than single-process SALib on one M1 Pro core, and 10.9x faster on a single output slice. On a scalar output there is no gain.

Which method should I use? ​

The methods measure different quantities, cost different numbers of model runs, and do not all accept the same problems. Four build their own sampling design, and the other nine work on (X,Y) pairs you already have.

Start at Choosing a Method. It walks three questions: can you still choose where to run the model, what should the number mean, and what is your evaluation budget.

For a high-dimensional study, read Scale and limits before choosing a design. It records the main cost drivers and the settings that move each method's practical limit.

The method capability table is the one place that records which methods accept correlated parameters, which accept categorical parameters, and which report bootstrap confidence intervals. tests/test_docs_matrix.py checks those three columns, plus Own design, against the code.

A note on the speed numbers ​

The gain is vectorization over output slices, so it scales with T×K and vanishes at T×K=1. Any speedup quoted without those two numbers says nothing. The benchmarks guide gives the full tables, the hardware, and the baseline.

jaxgsa's Sobol sampling and analysis workflow follows SALib, reimplemented for JAX.

Released under the BSD-3-Clause License.