jaxgsa

Global sensitivity analysis in JAX

jaxgsa helps you find which inputs drive variation in a model's output. You choose input distributions, evaluate your model at sample points, and use jaxgsa to estimate each input's effect. It does not run your model for you.

Install

uv add jaxgsa

Python 3.12 or newer is required. If you do not use uv, install with pip install jaxgsa.

First analysis

This example uses the Ishigami test model, whose sensitivity indices are known. Save it as first_analysis.py and run uv run first_analysis.py:

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

# Draw the structured input values Sobol analysis needs.
design = jaxgsa.sobol.sample(PROBLEM, n_samples=4096, seed=42, verbose=False)

# In your work, replace this call with your model evaluated at design.samples.
Y = evaluate(design.samples)

result = jaxgsa.sobol.analyze(design, Y, verbose=False)
print("S1:", result.S1)
print("ST:", result.ST)

S1 estimates how much output variation an input explains on its own. ST also counts its interactions with other inputs. For this run, x3 has S1 near zero but ST around 0.24: it matters through an interaction with x1. Looking at S1 alone would miss it. These numbers are estimates, so check their uncertainty before making a close decision.

The getting-started tutorial explains each step and how to interpret the result.

Choose a method

jaxgsa provides thirteen analysis methods. A useful starting point depends on the data you can collect and the question you want answered:

The methods guide covers all thirteen methods and their assumptions. The API reference gives signatures and result fields.

Scalar, multi-output, and time-series results are supported. Twelve methods also accept output channels with separate, irregular time grids. See the multi-output and irregular-output examples.

More to read

• Concepts: what sensitivity indices mean.
• Ishigami example: sampling error and convergence.
• Configuration: precision, memory budget, and compilation.
• Benchmarks: measurement methods and performance results.

Browser workbench

The separate browser workbench runs ported analyses locally in the browser. Its source and development instructions are in web/. The Python package and documentation site have separate builds.

Develop

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

See CONTRIBUTING.md. To run the docs locally:

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

The repository also includes a coding agent skill covering the API and method-specific caveats.

Cite jaxgsa

For research, cite the exact jaxgsa version used and the primary paper for each sensitivity method. For v1.0.0, use 10.5281/zenodo.23054912. Use CITATION.cff for software citation metadata and the methods guide for method references.

License

BSD-3-Clause. See LICENSE.