gsaxGlobal Sensitivity Analysis in JAX

GPU-accelerated Sobol indices and RS-HDMR with JIT compilation, vectorized bootstrap, and multi-output support.

Sobol Indices

First-order, total-order, and second-order indices via Saltelli sampling with Sobol quasi-random sequences.

RS-HDMR

Surrogate-based sensitivity analysis that works with any (X, Y) pairs. Includes a built-in emulator for prediction.

Save & Reload Samples

Persist unique sample matrices plus Saltelli reconstruction metadata with `SamplingResult.save()` and reload them later with `gsax.load()`.

Multi-Output & Time-Series

Pass scalar, (N, K), or (N, T, K) outputs. All indices are computed in a single vectorized pass.

Up to 929× Faster than SALib

Fused JIT kernels and vectorized execution replace Python loops. 4.7× faster even on scalar outputs, up to 929× on multi-output workloads.

gsax's Sobol sampling and analysis workflow is heavily drawn from SALib, adapted here into a JAX-first implementation focused on JIT compilation, accelerator execution, and multi-output workloads.