jaxgsa

Global Sensitivity Analysis in JAX

jaxgsa tells you which of your model's inputs drive its output. You give it input samples and the outputs your model produced for them. It returns sensitivity indices that rank the inputs and show their interactions.

Thirteen methods share one interface and one output contract. Regular outputs can be scalar, multi-output, or time-series; twelve methods also accept bucketed channels with their own irregular time grids. Eleven methods are JIT-compiled and vectorized over regular output axes. A model with 50 timesteps and 6 outputs then costs one compiled pass, not 300 Python loop iterations. Those eleven also export a traceable indices(), so you can put the estimator itself under jit, vmap and grad. The other two, kucherenko and vkoga, run on the host in NumPy and SciPy by design. They have no traceable core.

One method, DGSM, needs a model written in JAX so it can take derivatives instead of running many samples. It costs about one gradient per sample point on a scalar output. For a model with several outputs, jaxgsa picks whichever of forward-mode or reverse-mode automatic differentiation is cheaper: reverse mode computes one output's full gradient per pass, so it gets expensive as the output count grows, while forward mode computes one input's effect on every output per pass, so it gets expensive as the input count grows instead. Later sections call the output shape (N, T, K), for N samples, T time steps, and K output channels; jaxgsa compares T * K against the number of inputs to pick the cheaper mode.

jaxgsa does not run your model. You control that. It builds designs and reads indices off the results.

Install

pip install jaxgsa
# or
uv add jaxgsa

Python 3.12 or newer. The runtime dependencies are jax, jaxlib, numpy, scipy, and xarray. Optional extras: examples (matplotlib) and dev (pytest, ruff, ty, SALib, POT).

Develop the web workbench

The browser workbench is a separate Vite app under web/. It runs the ported analysis code locally in the browser; it does not start a Python server.

cd web
npm install
npm run dev          # start the local Vite server

The web checks are:

cd web
npm run typecheck
npm run lint
npm test
npm run build

The VitePress documentation site is a separate project under docs/. Run its commands from docs/, or use npm --prefix docs ... from the repository root:

npm --prefix docs install
npm --prefix docs run docs:dev
npm --prefix docs run docs:check

The web workbench and the documentation site have separate dependency installs and development servers.

Citing jaxgsa

If you use jaxgsa in research, cite the exact version that produced your results. The repository's CITATION.cff contains the current software-citation metadata. If a version-specific DOI is available, use that DOI; otherwise cite the corresponding Git tag or GitHub release. Also cite the primary paper for each sensitivity method you use, as listed in the methods guide.

For version 0.9.1, before a DOI is available, the citation is:

Pessina, D., and Papathanasiou, M. M. (2026). jaxgsa: Global Sensitivity
Analysis in JAX (Version 0.9.1) [Computer software]. Imperial College London.
https://github.com/DanielePessina/jaxgsa

Coding agent skill

The repository ships an agent skill that teaches a coding agent the jaxgsa API: the sampling designs, all thirteen methods, and the caveats that decide whether an index means anything. Install it into a project with

npx skills add https://github.com/DanielePessina/jaxgsa

The skill lives at skills/jaxgsa/SKILL.md, with one reference file per method family under skills/jaxgsa/reference/. It needs no extra setup once installed.

Quickstart

Sobol indices on the Ishigami function, which has known analytic indices.

import jaxgsa
from jaxgsa.benchmarks.ishigami import PROBLEM, evaluate

# 1. Build a Saltelli design. `samples` is the (n_runs, D) array to run your model on.
design = jaxgsa.sobol.sample(PROBLEM, n_samples=16384, seed=42)

# 2. Run the model. Here it is the Ishigami test function.
Y = evaluate(design.samples)  # shape (16384,)

# 3. Read the indices off the design and the outputs.
result = jaxgsa.sobol.analyze(design, Y)

print("S1:", result.S1)
print("S2 x1-x3:", result.S2[0, 2])

Both sample and analyze print a summary by default. This is what the script writes, verbatim, apart from the timing line, which depends on your machine:

jaxgsa.sobol.sample: D=3, mode=second-order, base_n=2048, requested_runs>=16384, n_runs=16384, n_expanded=16384, duplicates_removed=0 (0.0%), scramble=True
jaxgsa.sobol.analyze
  problem: D=3 (x1, x2, x3)
    marginals: uniform=3
    correlation: independent
    output: N=16384 runs, T=1 x K=1 output slice
    invalid: none found in 2048 Saltelli groups (policy 'raise')
  timing:
    estimators (includes compile on the first call): 0.7828 s
    slice_chunk_size: 1 (resolved from the memory budget)
    estimator: saltelli-jansen
  results: top 3 of 3 parameters by ST
    1. x1  ST=0.5559
    2. x2  ST=0.4414
    3. x3  ST=0.2415
S1: [ 0.308098   0.4440502 -0.0113217]
S2 x1-x3: 0.24699646

Pass verbose=False to any analyze() or sample() call to silence it. The summary is worth reading once per new problem, because the invalid and marginals lines catch the mistakes that silently ruin an analysis.

Every index is a fraction of the output variance. S1 is an input's direct effect. ST also counts every interaction the input takes part in. Reading the numbers above:

• x2 has the largest direct effect, 0.444 of the variance, and S1 equals ST, so it acts alone.
• x1 has S1 = 0.308 but ST = 0.556. The gap is interaction.
• x3 has S1 = -0.011, which is zero plus estimator noise. On its own it does nothing. Yet ST = 0.242. result.S2[0, 2] = 0.247 names the partner. The whole effect of x3 runs through x1. Fixing x3 at its nominal value would still change the output, so you cannot drop it.

A negative S1 is not a bug. The Saltelli estimator is a difference of Monte Carlo means, so a true zero comes out slightly either side of zero. Read it as "no direct effect", and read its size as your noise floor.

The analytic answers are S1 = (0.3139, 0.4424, 0) and ST = (0.5576, 0.4424, 0.2437). At n_samples=16384 every index above is within 0.012 of the truth. At n_samples=4096 the worst error is 0.07, on ST for x1. Monte Carlo error falls with the square root of the sample count, so budget for it. Quadrupling the runs roughly halves the error.

When you cannot choose the sample points

Sobol indices normally require the Saltelli design: a specific pattern of sample points, built and evaluated ahead of time. Nine of the thirteen methods skip that requirement. They work on whatever (X, Y) pairs you already have, including runs from an old sweep. Polynomial chaos, one of the nine, fits a polynomial surrogate to those pairs. The polynomials it uses are mutually orthogonal, which lets it read exact Sobol indices straight off the fitted coefficients, with no extra integration step.

import jaxgsa
from jaxgsa.benchmarks.ishigami import PROBLEM, evaluate

X = jaxgsa.sampling.monte_carlo(PROBLEM, n=2000, seed=0)  # any (N, D) points
Y = evaluate(X)

result = jaxgsa.pce.analyze(PROBLEM, X, Y, order=6, verbose=False)
print("S1:", result.S1)
print("ST:", result.ST)
print("LOO RMSE:", result.loo_rmse, "  output sd:", Y.std())

S1: [3.1933489e-01 4.4209677e-01 1.2573625e-04]
ST: [0.5573787  0.44309324 0.23836787]
LOO RMSE: 0.52724934   output sd: 3.5930953

2000 unstructured points land within 0.0054 of every analytic index. The Saltelli run above needed 16384 and did worse. That is the surrogate paying off, and it is why PCE is the right first try for a smooth model.

The catch is that you are now trusting a fit. loo_rmse is the number that decides whether to trust it. Leave-one-out error of 0.527 against an output standard deviation of 3.59 means the surrogate reproduces about 98% of the variance. If loo_rmse approaches the output standard deviation, the indices describe the surrogate and not your model. Raise order, add samples, or switch to jaxgsa.hdmr, whose B-spline basis handles kinks that polynomials cannot.

Multi-output and time series

Every method takes Y as scalar (N,), multi-output (N, K), or time-series (N, T, K). The output axes are never inferred or transposed, and one call covers all of them.

import jax.numpy as jnp

import jaxgsa
from jaxgsa import Problem

problem = Problem.from_dict(
    {"amplitude": (0.5, 1.5), "decay": (0.1, 1.0), "freq": (1.0, 3.0)}
)
t = jnp.linspace(0.1, 5.0, 50)


def model(X):  # X is (N, 3)
    a, k, w = X[:, 0:1], X[:, 1:2], X[:, 2:3]
    return (a * jnp.exp(-k * t) * jnp.sin(w * t))[:, :, None]  # (N, T=50, K=1)


design = jaxgsa.sobol.sample(problem, n_samples=8192, seed=0, verbose=False)
result = jaxgsa.sobol.analyze(design, model(jnp.asarray(design.samples)), verbose=False)

print("S1 shape:", result.S1.shape)  # (T, K, D)
print("t=0.6  S1:", result.S1[5, 0])
print("t=4.5  S1:", result.S1[44, 0])

S1 shape: (50, 1, 3)
t=0.6  S1: [0.6211161  0.18043584 0.16506985]
t=4.5  S1: [-0.00563779 -0.01124666  0.4224577 ]

150 index sets, one compiled pass, one set of model runs. The ranking flips along the trajectory. Early on, amplitude explains 62% of the variance. By t = 4.5 the signal has decayed and only the phase is left, so frequency explains 42% and amplitude explains nothing. A single index averaged over time would have hidden both facts.

Watch for zero-variance slices. If your model returns a constant at t = 0, the indices there are 0/0 and jaxgsa returns NaN with a JaxgsaWarning that names the slice.

If output channels do not share a time grid, pass a list or dict of (times, values) pairs instead. Automatic bucketing analyzes each channel on its own grid and preserves those coordinates without padding or fabricated values:

Y = [
    (t_concentration, concentration),  # values: (N, T_concentration)
    (t_diameter, diameter),             # values: (N, T_diameter)
]
result = jaxgsa.sobol.analyze(design, Y, verbose=False)
ds = result.to_dataset()

The ragged form automatically selects bucketing. There is no mask mode because padding a common grid would add work without helping the per-channel estimators; DGSM requires a fixed Jacobian layout and does not accept the ragged form. See the irregular output grids guide.

The thirteen methods

Every method exposes analyze(). The design-based ones also expose sample(). Results carry to_dataset() for labeled xarray output. Every method except eFAST and HSIC reports bootstrap confidence intervals through n_bootstrap.

The methods guide carries the estimators, the references, and a capability table that a test checks against the code.

Correlated and categorical inputs

Declare a Gaussian-copula correlation matrix on the Problem and jaxgsa.sampling.monte_carlo draws from it. Declare a categorical parameter with {"dist": "categorical", "probs": [...], "labels": [...]} and samples carry integer level codes.

Neither is universally supported, and jaxgsa refuses rather than approximates. Sobol, Morris, eFAST, PCE, DGSM, and PCE-backed Shapley raise a ValueError on a correlated problem that names the alternatives. Every method whose indices would depend on the arbitrary order of category codes raises on a categorical one. Sobol is fine with categoricals, because the Saltelli column-swap scheme never looks at the values.

See correlated inputs and categorical inputs.

Performance

The performance suite covers all thirteen methods across scalar, high-sample, high-dimensional, multi-output, time-series, and bootstrap workloads. On an Apple M1 Pro CPU, the current performance branch reduced the sum of warm median analysis times across 27 isolated cases from 0.921 s to 0.430 s relative to master.

The largest isolated changes were HSIC, from 153.0 ms to 23.1 ms, and VKOGA, from 348.3 ms to 27.1 ms. A 30-parameter Sobol analysis with second-order indices fell from 7.37 ms to 3.45 ms. These are analysis timings; model evaluation and sampling are excluded.

Against SALib 1.5.2 on the same machine, scalar Sobol without bootstrap still favours SALib (0.2 ms against 0.6 ms). With 300 output slices and second-order indices, jaxgsa takes 7.4 ms against 262.8 ms. The gain comes from compiling and vectorizing work across output slices, so every comparison must state its shape and whether compilation is included.

Full tables, methodology, and the script are in the benchmarks guide.

uv run --extra dev benchmark_salib.py

Documentation

• Getting started
• Methods guide, including the capability table
• Scaling to large problems, including high-dimensional planning guidance
• Configuration, including 64-bit floats and the persistent compilation cache
• API reference
• Examples, one page per method

One configuration note is worth repeating here. JAX defaults to float32 and silently downcasts float64 arrays. For precision-sensitive Sobol or HSIC work, call jax.config.update("jax_enable_x64", True) before you create the first array.

Development

git clone https://github.com/DanielePessina/jaxgsa.git
cd jaxgsa
uv sync --extra dev
uv run pytest

See CONTRIBUTING.md.

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

License

BSD-3-Clause. See LICENSE.