gsax 0.1.0__py3-none-any.whl

gsax/__init__.py

@@ -0,0 +1,48 @@
+from gsax.dgsm import DGSMResult
+from gsax.dgsm import analyze as analyze_dgsm
+from gsax.efast import EFASTResult
+from gsax.efast import analyze as analyze_efast
+from gsax.efast import sample as sample_efast
+from gsax.hdmr import HDMREmulator, HDMRResult
+from gsax.hdmr import analyze as analyze_hdmr
+from gsax.hdmr import emulate as emulate_hdmr
+from gsax.hsic import HSICResult
+from gsax.hsic import analyze as analyze_hsic
+from gsax.pawn import PAWNResult
+from gsax.pawn import analyze as analyze_pawn
+from gsax.pce import PCEResult
+from gsax.pce import analyze as analyze_pce
+from gsax.pce import emulate as emulate_pce
+from gsax.problem import GaussianInputSpec, Problem, UniformInputSpec
+from gsax.sampling import SamplingResult, downsample, load, sample, sample_mc, verify_prefix
+from gsax.sobol import SAResult, analyze
+
+__all__ = [
+    "DGSMResult",
+    "EFASTResult",
+    "GaussianInputSpec",
+    "HDMREmulator",
+    "HDMRResult",
+    "HSICResult",
+    "PAWNResult",
+    "PCEResult",
+    "Problem",
+    "SAResult",
+    "SamplingResult",
+    "UniformInputSpec",
+    "analyze",
+    "analyze_dgsm",
+    "analyze_efast",
+    "analyze_hdmr",
+    "analyze_hsic",
+    "analyze_pawn",
+    "analyze_pce",
+    "downsample",
+    "emulate_hdmr",
+    "emulate_pce",
+    "load",
+    "sample",
+    "sample_efast",
+    "sample_mc",
+    "verify_prefix",
+]

gsax/_normalization.py

@@ -0,0 +1,170 @@
+"""Shared output helpers for analysis entrypoints."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import jax.numpy as jnp
+from jax import Array
+
+if TYPE_CHECKING:
+    from gsax.problem import Problem
+
+
+def _default_output_names(K: int, problem: Problem) -> list[str]:
+    """Resolve output coordinate labels, defaulting to y0, y1, ...
+
+    Args:
+        K: Number of output variables.
+        problem: Problem definition (may carry output_names).
+
+    Returns:
+        List of K string labels.
+    """
+    if problem.output_names is not None:
+        if len(problem.output_names) != K:
+            raise ValueError(f"output_names length {len(problem.output_names)} != K={K}")
+        return list(problem.output_names)
+    return [f"y{i}" for i in range(K)]
+
+
+def _prepare_Y(
+    Y: Array,
+) -> tuple[Array, bool, bool]:
+    """Promote Y to a canonical 3-D shape (n_total, T, K).
+
+    Args:
+        Y: Model output array with 1, 2, or 3 dimensions.
+
+    Returns:
+        A tuple ``(Y_3d, squeeze_time, squeeze_output)`` indicating which
+        singleton dimensions were inserted and should be removed later.
+    """
+    squeeze_time = False
+    squeeze_output = False
+    # Normalize to 3-D so downstream kernels always see (samples, T, K)
+    # without shape-dependent branching.
+    if Y.ndim == 1:  # scalar output per sample -- both T and K are singleton
+        Y = Y[:, None, None]
+        squeeze_time = True
+        squeeze_output = True
+    elif Y.ndim == 2:  # multi-output but single timestep -- only T is singleton
+        Y = Y[:, None, :]
+        squeeze_time = True
+    return Y, squeeze_time, squeeze_output
+
+
+def _warn_zero_variance_slices(
+    Y: Array,
+    output_names: tuple[str, ...] | None = None,
+) -> None:
+    """Check for zero-variance output slices before analysis and warn.
+
+    Args:
+        Y: Model output array with shape ``(n_expanded, ...)`` where
+            trailing dims are ``()``, ``(K,)``, or ``(T, K)``.
+        output_names: Optional names for the K output dimension.
+    """
+    import warnings
+
+    # Collapse trailing dims so variance is computed per (t, k) slice.
+    flat = Y.reshape(Y.shape[0], -1)
+    n_outputs = flat.shape[1]
+
+    # Recover K from the original shape to map flat indices back to named outputs.
+    trailing = Y.shape[1:]
+    if len(trailing) == 0:
+        K = 1
+    elif len(trailing) == 1:
+        K = trailing[0]
+    else:
+        K = trailing[1]
+
+    if output_names is not None and len(output_names) != K:
+        raise ValueError(
+            f"len(output_names)={len(output_names)} does not match number of outputs K={K}"
+        )
+
+    def _fmt_k(k: int) -> str:
+        if output_names is not None:
+            return f"k={k} ('{output_names[k]}')"
+        return f"k={k}"
+
+    # Sample variance along axis 0; zero means the output is constant
+    # and Sobol indices become 0/0 = NaN.
+    var_per_slice = jnp.var(flat, axis=0)
+    zero_mask = var_per_slice == 0
+    n_zero = int(jnp.sum(zero_mask))
+
+    if n_zero == 0:
+        return
+
+    if n_outputs == 1:
+        if output_names is not None and len(output_names) == 1:
+            msg = f"gsax: output '{output_names[0]}' has zero variance — all indices will be NaN"
+        else:
+            msg = "gsax: output has zero variance — all indices will be NaN"
+        warnings.warn(msg, stacklevel=2)
+        return
+
+    # Materialize indices eagerly -- this is a rare warning path, not a hot loop.
+    zero_indices = [int(i) for i in jnp.where(zero_mask)[0]]
+
+    if len(trailing) == 1:  # single-timestep: flat index equals output index k
+        labels = [_fmt_k(k) for k in zero_indices]
+        if len(labels) > 5:  # cap displayed labels to keep warnings readable
+            shown = ", ".join(labels[:5])
+            extra = f"... and {len(labels) - 5} more"
+            label_str = f"{shown}, {extra}"
+        else:
+            label_str = ", ".join(labels)
+        warnings.warn(
+            f"gsax: {n_zero}/{n_outputs} output(s) have zero variance "
+            f"({label_str}) — corresponding indices will be NaN",
+            stacklevel=2,
+        )
+    elif len(trailing) == 2:  # multi-timestep: flat index encodes (t, k) in row-major order
+        affected = []
+        for idx in zero_indices:
+            # divmod decomposes flat index into (time_step, output_column)
+            t, k = divmod(idx, K)
+            affected.append(f"(t={t}, {_fmt_k(k)})")
+        if len(affected) > 5:  # cap displayed labels to keep warnings readable
+            shown = ", ".join(affected[:5])
+            extra = f"... and {len(affected) - 5} more"
+            label_str = f"{shown}, {extra}"
+        else:
+            label_str = ", ".join(affected)
+        warnings.warn(
+            f"gsax: {n_zero}/{n_outputs} output slice(s) have zero variance "
+            f"[{label_str}] — corresponding indices will be NaN",
+            stacklevel=2,
+        )
+
+
+def _prenormalize_outputs(Y: Array) -> tuple[Array, Array, Array, Array]:
+    """Standardize outputs over the sample axis.
+
+    Args:
+        Y: Output array with shape ``(N, ...)``. The first axis is treated as
+            the sample axis, and all trailing axes are normalized
+            independently.
+
+    Returns:
+        A tuple ``(Y_norm, y_mean, y_std, safe_scale)`` where:
+            - ``Y_norm`` is the centered/scaled output array.
+            - ``y_mean`` is the per-output-slice mean.
+            - ``y_std`` is the per-output-slice original standard deviation.
+            - ``safe_scale`` is the divisor actually used, with zeros replaced
+              by ``1.0`` to avoid division by zero.
+    """
+    # Centering + scaling stabilizes Sobol variance estimators when output
+    # magnitudes vary across slices (prevents large-magnitude slices from
+    # dominating numerical precision).
+    y_mean = jnp.mean(Y, axis=0)
+    y_std = jnp.std(Y, axis=0)
+    # Replace zero std with 1.0 so division doesn't produce NaN; the
+    # corresponding zero-variance output slices remain all-zero after scaling.
+    safe_scale = jnp.where(y_std == 0, jnp.ones_like(y_std), y_std)
+    Y_norm = (Y - y_mean) / safe_scale
+    return Y_norm, y_mean, y_std, safe_scale

gsax/_transforms.py

@@ -0,0 +1,56 @@
+"""Shared input transforms for expansion methods (HDMR, PCE)."""
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+from gsax.problem import Problem
+
+
+def cdf_to_unit_interval(X: Array, problem: Problem) -> Array:
+    """Map each column of X to [0, 1] via its marginal CDF.
+
+    Uniform inputs use an affine map. Gaussian inputs (truncated or
+    not) are mapped through their normal CDF. The result is always in
+    [0, 1], suitable for B-spline bases or further affine mapping to
+    [-1, 1] for Legendre polynomials.
+
+    Args:
+        X: (N, D) input samples in physical units.
+        problem: Problem with per-dimension distribution specs.
+
+    Returns:
+        (N, D) array with each column in [0, 1].
+    """
+    from jax.scipy.stats import norm as jax_norm
+    from scipy.stats import truncnorm
+
+    D = problem.num_vars
+    cols = []  # one column per dimension, stacked at the end into (N, D)
+
+    # CDF-to-unit-interval mapping: apply F(x) per dimension so each column
+    # lands in (0, 1).  For Gaussian inputs, standardised truncation bounds
+    # a=(lo-mu)/sigma, b=(hi-mu)/sigma follow scipy's truncnorm convention.
+    # Clipping output to (1e-12, 1-1e-12) keeps values in the open interval
+    # so downstream inverse transforms (ppf) never receive exactly 0 or 1.
+    for d in range(D):
+        dist, first, second, lo, hi = problem.input_specs[d]
+        if dist == "uniform":  # affine map (x - lo)/(hi - lo) is the CDF of U(lo, hi)
+            cols.append((X[:, d] - first) / (second - first))
+        # Truncated Gaussian -- need scipy because JAX lacks a truncnorm CDF.
+        elif lo is not None or hi is not None:
+            mean, variance = first, second
+            std = float(jnp.sqrt(variance))
+            a = -np.inf if lo is None else (lo - mean) / std
+            b = np.inf if hi is None else (hi - mean) / std
+            u = jnp.asarray(truncnorm.cdf(np.asarray(X[:, d]), a=a, b=b, loc=mean, scale=std))
+            cols.append(jnp.clip(u, 1e-12, 1.0 - 1e-12))
+        else:  # unbounded Gaussian -- stays on device via jax.scipy
+            mean, variance = first, second
+            std = float(jnp.sqrt(variance))
+            u = jax_norm.cdf(X[:, d], loc=mean, scale=std)
+            cols.append(jnp.clip(u, 1e-12, 1.0 - 1e-12))
+
+    return jnp.column_stack(cols)  # reassemble per-dimension columns into (N, D)

gsax/benchmarks/__init__.py

@@ -0,0 +1,22 @@
+"""Analytical benchmark functions with known Sobol indices.
+
+Each submodule provides:
+    - ``PROBLEM``: a :class:`~gsax.Problem` definition.
+    - ``evaluate(X)``: the benchmark function (JAX-compatible).
+    - ``ANALYTICAL_S1``, ``ANALYTICAL_ST``, ``ANALYTICAL_S2``: precomputed
+      indices for the default parameters.
+    - ``analytical_indices(...)``: compute indices for custom parameters.
+
+Example::
+
+    from gsax.benchmarks import ishigami
+    from gsax import sample, analyze
+
+    sr = sample(ishigami.PROBLEM, 4096)
+    Y = ishigami.evaluate(sr.samples)
+    result = analyze(sr, Y)
+"""
+
+from gsax.benchmarks import ishigami, linear, oakley_ohagan, sobol_g
+
+__all__ = ["ishigami", "linear", "oakley_ohagan", "sobol_g"]

gsax/benchmarks/ishigami.py

@@ -0,0 +1,111 @@
+"""Ishigami test function for sensitivity analysis benchmarking.
+
+A standard benchmark with 3 input parameters (x1, x2, x3) and known
+analytical Sobol indices, commonly used to validate sensitivity analysis methods.
+"""
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+from gsax.problem import Problem
+
+# All three inputs are uniform on [-pi, pi], the standard domain for
+# the trigonometric terms in the Ishigami function.
+PROBLEM = Problem.from_dict(
+    {
+        "x1": (-np.pi, np.pi),
+        "x2": (-np.pi, np.pi),
+        "x3": (-np.pi, np.pi),
+    }
+)
+
+
+def analytical_indices(
+    A: float = 7.0, B: float = 0.1
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Compute analytical first-order, total-order, and second-order Sobol indices.
+
+    For the Ishigami function ``f(x) = sin(x1) + A*sin^2(x2) + B*x3^4*sin(x1)``
+    with ``x_i ~ U[-pi, pi]``, the ANOVA decomposition is available in closed
+    form using Gaussian moment identities for uniform trig integrals.
+
+    The main-effect variances are:
+
+    - ``V1 = (1 + B*pi^4/5)^2 / 2``
+    - ``V2 = A^2 / 8``
+    - ``V3 = 0``
+
+    The only pairwise interaction is between x1 and x3:
+
+    - ``V13 = 8 * B^2 * pi^8 / 225``
+
+    And the total variance:
+
+    - ``V = 1/2 + B*pi^4/5 + B^2*pi^8/18 + A^2/8``
+
+    Args:
+        A: Model parameter controlling the second-order term.
+        B: Model parameter controlling the higher-order interaction term.
+
+    Returns:
+        ``(S1, ST, S2)`` where S1 and ST are ``(3,)`` arrays and S2 is
+        a ``(3, 3)`` symmetric matrix with NaN on the diagonal.
+    """
+    pi4 = np.pi**4
+    pi8 = np.pi**8
+
+    # Main-effect variances from the ANOVA decomposition.
+    V1 = 0.5 * (1.0 + B * pi4 / 5.0) ** 2
+    V2 = A**2 / 8.0
+    V3 = 0.0
+
+    # x1-x3 interaction variance (the B*x3^4*sin(x1) cross-term).
+    V13 = 8.0 * B**2 * pi8 / 225.0
+
+    # Total variance = sum of all partial variances.
+    VY = 0.5 + B * pi4 / 5.0 + B**2 * pi8 / 18.0 + A**2 / 8.0
+
+    Vi = np.array([V1, V2, V3])
+    S1 = Vi / VY
+
+    # Total-order: main effect + all interactions involving that input.
+    ST = np.array([V1 + V13, V2, V13]) / VY
+
+    # Second-order interaction matrix: only (0, 2) and (2, 0) are nonzero.
+    S2 = np.full((3, 3), np.nan)
+    S2[0, 1] = 0.0
+    S2[1, 0] = 0.0
+    S2[0, 2] = V13 / VY
+    S2[2, 0] = V13 / VY
+    S2[1, 2] = 0.0
+    S2[2, 1] = 0.0
+
+    return S1, ST, S2
+
+
+# Precomputed analytical solutions for A=7, B=0.1.
+# x3 has zero first-order effect (enters only through the B*x3^4*sin(x1) interaction),
+# making it a good test for methods that must distinguish S1=0 from ST>0.
+ANALYTICAL_S1, ANALYTICAL_ST, ANALYTICAL_S2 = analytical_indices()
+
+
+def evaluate(X: Array, A: float = 7.0, B: float = 0.1) -> Array:
+    """Evaluate the Ishigami function.
+
+    f(x) = sin(x1) + A*sin^2(x2) + B*x3^4*sin(x1)
+
+    Args:
+        X: Input array of shape ``(N, 3)`` with columns x1, x2, x3.
+        A: Model parameter controlling the second-order term.
+        B: Model parameter controlling the higher-order interaction term.
+
+    Returns:
+        Array of shape ``(N,)`` with the function output for each sample.
+    """
+    # sin(x1): first-order effect of x1
+    # A*sin^2(x2): purely additive in x2 (large A=7 makes it the dominant first-order term)
+    # B*x3^4*sin(x1): x1-x3 interaction (small B=0.1 keeps it a minority contribution)
+    return jnp.sin(X[:, 0]) + A * jnp.sin(X[:, 1]) ** 2 + B * X[:, 2] ** 4 * jnp.sin(X[:, 0])

gsax/benchmarks/linear.py

@@ -0,0 +1,94 @@
+"""Linear additive model for sensitivity analysis benchmarking.
+
+The simplest possible benchmark: a weighted sum of independent uniform
+inputs. Because the model is purely additive, first-order indices equal
+total-order indices and all second-order interactions are exactly zero.
+
+Useful for:
+    - Verifying that a SA method correctly identifies zero interactions.
+    - Sanity-checking that S1 == ST when there are no interactions.
+    - Testing convergence rates (exact analytical solution is trivial).
+"""
+
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+from gsax.problem import Problem
+
+# Increasing weights (1, 2, 3) test whether the method correctly ranks
+# input importance: x3 should have the largest index, x1 the smallest.
+DEFAULT_COEFFS = (1.0, 2.0, 3.0)
+DEFAULT_BOUNDS = ((0.0, 1.0), (0.0, 1.0), (0.0, 1.0))
+
+PROBLEM = Problem(
+    names=tuple(f"x{i + 1}" for i in range(len(DEFAULT_COEFFS))),
+    bounds=DEFAULT_BOUNDS,
+)
+
+
+def evaluate(
+    X: Array,
+    coeffs: tuple[float, ...] = DEFAULT_COEFFS,
+) -> Array:
+    """Evaluate the linear additive model.
+
+    .. math::
+        f(\\mathbf{x}) = \\sum_{j=1}^{D} c_j \\, x_j
+
+    Args:
+        X: Input array of shape ``(N, D)``.
+        coeffs: Coefficient per dimension.
+
+    Returns:
+        Array of shape ``(N,)`` with function values.
+    """
+    # Matrix-vector multiply: vectorized weighted sum across all N samples at once.
+    return X @ jnp.asarray(coeffs)
+
+
+def analytical_indices(
+    coeffs: tuple[float, ...] = DEFAULT_COEFFS,
+    bounds: tuple[tuple[float, float], ...] = DEFAULT_BOUNDS,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Compute analytical Sobol indices for the linear additive model.
+
+    For ``Y = sum c_j x_j`` with independent inputs:
+
+    .. math::
+        V_j = c_j^2 \\operatorname{Var}(x_j), \\quad
+        S_{1,j} = S_{T,j} = V_j / V(Y), \\quad
+        S_{2,jk} = 0
+
+    Args:
+        coeffs: Coefficient per dimension.
+        bounds: ``(low, high)`` bounds for each uniform input.
+
+    Returns:
+        ``(S1, ST, S2)`` where S1 == ST (no interactions) and S2 is
+        all-zero off-diagonal with NaN on the diagonal.
+    """
+    c = np.asarray(coeffs, dtype=float)
+    D = len(c)
+    # Var(Uniform[lo, hi]) = (hi - lo)^2 / 12
+    var_x = np.array([(hi - lo) ** 2 / 12.0 for lo, hi in bounds])
+    # Variance propagation for linear functions: Vi = c_j^2 * Var(x_j)
+    Vi = c**2 * var_x
+    VY = Vi.sum()
+
+    # Purely additive model => no interactions => S1 = ST and sum(S1) = 1
+    S1 = Vi / VY
+    ST = S1.copy()
+
+    # Diagonal is NaN by convention (S2_jj is undefined); off-diagonals are
+    # exactly zero because a purely additive model has no interactions.
+    S2 = np.full((D, D), np.nan)
+    for j in range(D):
+        for k in range(j + 1, D):
+            S2[j, k] = 0.0
+            S2[k, j] = 0.0
+
+    return S1, ST, S2
+
+
+ANALYTICAL_S1, ANALYTICAL_ST, ANALYTICAL_S2 = analytical_indices()

gsax/benchmarks/oakley_ohagan.py

@@ -0,0 +1,189 @@
+# ruff: noqa: E501
+"""Oakley & O'Hagan (2004) 15-dimensional Gaussian-input benchmark.
+
+A canonical benchmark for variance-based sensitivity analysis with
+Gaussian inputs. The function combines linear, trigonometric, and
+quadratic terms so that first-order, total-order, and pairwise
+interaction Sobol indices are all available in closed form.
+
+.. math::
+    f(\\mathbf{x}) = \\mathbf{a}_1^\\top \\mathbf{x}
+        + \\mathbf{a}_2^\\top \\sin(\\mathbf{x})
+        + \\mathbf{a}_3^\\top \\cos(\\mathbf{x})
+        + \\mathbf{x}^\\top M \\mathbf{x}
+
+with :math:`x_i \\sim \\mathcal{N}(0, \\sigma^2)` i.i.d.
+
+The coefficients are the published values from Oakley & O'Hagan (2004,
+JRSS-B 66:751-769), embedded as literals so no external data file is
+needed.
+
+References:
+    Oakley, J. E. and O'Hagan, A. (2004). Probabilistic sensitivity
+    analysis of complex models: a Bayesian approach. J. R. Statist.
+    Soc. B, 66(3):751-769.
+"""
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+from jax import Array
+
+from gsax.problem import Problem
+
+# 15 dimensions — high enough to stress-test scalability and to span
+# a wide range of importance levels (some inputs are nearly inert).
+D = 15
+DEFAULT_SIGMA = 1.0
+
+# Published coefficients from Oakley & O'Hagan (2004, Table 1 / psa_example.txt).
+# fmt: off
+_M = np.array([
+    [-0.0225, -0.185, 0.134, 0.369, 0.172, 0.137, -0.44, -0.0814, 0.713, -0.444, 0.504, -0.0241, -0.0459, 0.217, 0.0559],
+    [0.257, 0.0538, 0.258, 0.238, -0.591, -0.0816, -0.287, 0.416, 0.498, 0.0839, -0.111, 0.0332, -0.14, -0.031, -0.223],
+    [-0.056, 0.195, 0.0955, -0.286, -0.144, 0.224, 0.145, 0.29, 0.231, -0.319, -0.29, -0.21, 0.431, 0.0244, 0.0449],
+    [0.664, 0.431, 0.299, -0.162, -0.315, -0.39, 0.177, 0.058, 0.172, 0.135, -0.353, 0.251, -0.0188, 0.365, -0.325],
+    [-0.121, 0.125, 0.107, 0.0466, -0.217, 0.195, -0.0655, 0.0244, -0.0968, 0.194, 0.334, 0.313, -0.0836, -0.253, 0.373],
+    [-0.284, -0.328, -0.105, -0.221, -0.137, -0.144, -0.115, 0.224, -0.0304, -0.515, 0.0173, 0.039, 0.361, 0.309, 0.05],
+    [-0.0779, 0.00375, 0.887, -0.266, -0.0793, -0.0427, -0.187, -0.356, -0.175, 0.0887, 0.4, -0.056, 0.137, 0.215, -0.0113],
+    [-0.0923, 0.592, 0.0313, -0.0331, -0.243, -0.0998, 0.0345, 0.0951, -0.338, 0.00639, -0.612, 0.0813, 0.887, 0.143, 0.148],
+    [-0.132, 0.529, 0.127, 0.0451, 0.584, 0.373, 0.114, -0.295, -0.57, 0.463, -0.0941, 0.14, -0.386, -0.449, -0.146],
+    [0.0581, -0.323, 0.0931, 0.0724, -0.569, 0.526, 0.237, -0.0118, 0.0718, 0.0783, -0.134, 0.227, 0.144, -0.452, -0.556],
+    [0.661, 0.346, 0.141, 0.519, -0.28, -0.16, -0.0684, -0.204, 0.0697, 0.231, -0.0444, -0.165, 0.216, 0.00427, -0.0874],
+    [0.316, -0.0276, 0.134, 0.135, 0.054, -0.174, 0.175, 0.0603, -0.179, -0.311, -0.254, 0.0258, -0.43, -0.623, -0.034],
+    [-0.29, 0.0341, 0.0349, -0.121, 0.026, -0.335, -0.414, 0.0532, -0.271, -0.0263, 0.41, 0.266, 0.156, -0.187, 0.0199],
+    [-0.244, -0.441, 0.0126, 0.249, 0.0711, 0.246, 0.175, 0.00853, 0.251, -0.147, -0.0846, 0.369, -0.3, 0.11, -0.757],
+    [0.0415, -0.26, 0.464, -0.361, -0.95, -0.165, 0.00309, 0.0528, 0.225, 0.384, 0.456, -0.186, 0.00823, 0.167, 0.16],
+], dtype=float)
+
+# Coefficients for linear, sin, and cos terms respectively.
+# Magnitudes increase toward higher indices, creating a natural importance
+# gradient: x11-x15 dominate, x1-x5 are nearly inert, x6-x10 intermediate.
+_A1 = np.array([0.0118, 0.0456, 0.2297, 0.0393, 0.1177, 0.3865, 0.3897, 0.6061, 0.6159, 0.4005, 1.0741, 1.1474, 0.788, 1.1242, 1.1982], dtype=float)
+_A2 = np.array([0.4341, 0.0887, 0.0512, 0.3233, 0.1489, 1.036, 0.9892, 0.9672, 0.8977, 0.8083, 1.8426, 2.4712, 2.3946, 2.0045, 2.2621], dtype=float)
+_A3 = np.array([0.1044, 0.2057, 0.0774, 0.273, 0.1253, 0.7526, 0.857, 1.0331, 0.8388, 0.797, 2.2145, 2.0382, 2.4004, 2.0541, 1.9845], dtype=float)
+# fmt: on
+
+# Reference S1 values from the original paper; useful for cross-checking
+# our analytical derivation against the published table.
+PUBLISHED_S1 = np.array(
+    [
+        0.00156,
+        0.000186,
+        0.001307,
+        0.003045,
+        0.002905,
+        0.023035,
+        0.024151,
+        0.026517,
+        0.046036,
+        0.014945,
+        0.101823,
+        0.135708,
+        0.101989,
+        0.105169,
+        0.122818,
+    ]
+)
+
+# Gaussian inputs (not uniform) -- one of few SA benchmarks with non-uniform distributions.
+PROBLEM = Problem.from_dict(
+    {
+        f"x{i + 1}": {"dist": "gaussian", "mean": 0.0, "variance": DEFAULT_SIGMA**2}
+        for i in range(D)
+    }
+)
+
+
+def evaluate(X: Array) -> Array:
+    """Evaluate the Oakley & O'Hagan function.
+
+    Args:
+        X: Input array of shape ``(N, 15)`` with ``x_i ~ N(0, sigma^2)``.
+
+    Returns:
+        Array of shape ``(N,)`` with function values.
+    """
+    Xj = jnp.asarray(X)
+    # Four additive components: linear a1^T x, trigonometric a2^T sin(x) + a3^T cos(x),
+    # and the quadratic form x^T M x which introduces all pairwise interactions.
+    return (
+        Xj @ jnp.asarray(_A1)
+        + jnp.sin(Xj) @ jnp.asarray(_A2)
+        + jnp.cos(Xj) @ jnp.asarray(_A3)
+        # Batched quadratic form: einsum contracts x_i * M_ij * x_j per sample row.
+        + jnp.einsum("ni,ij,nj->n", Xj, jnp.asarray(_M), Xj)
+    )
+
+
+def analytical_indices(
+    sigma: float = DEFAULT_SIGMA,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Compute analytical first-order, total-order, and second-order Sobol indices.
+
+    Each input enters through the block ``(x, sin x, cos x, x^2)``
+    (main effect) plus off-diagonal quadratic cross-terms (pairwise
+    interactions). The main-effect variance is ``c_i' Sigma c_i`` where
+    ``Sigma`` is the covariance of that block under ``N(0, sigma^2)``,
+    and pairwise interaction variance is ``(M_ij + M_ji)^2 * sigma^4``.
+
+    Args:
+        sigma: Standard deviation of each input.
+
+    Returns:
+        ``(S1, ST, S2)`` where S1 and ST are ``(15,)`` arrays and S2 is
+        a ``(15, 15)`` symmetric matrix with NaN on the diagonal.
+    """
+    s2 = sigma**2
+    # Gaussian moment-generating-function identities for X ~ N(0, s2):
+    # E[sin X] = 0, E[cos X] = exp(-s2/2), E[sin^2 X] = (1 - exp(-2s2))/2
+    es = np.exp(-0.5 * s2)  # = E[cos X], used in cross-covariances
+    es2 = np.exp(-2.0 * s2)  # appears in Var[sin X] and Var[cos X]
+
+    # Covariance matrix of the feature vector (x, sin x, cos x, x^2) for X~N(0,s2).
+    # Block-diagonal: (x, sin x) decouple from (cos x, x^2) because odd/even symmetry.
+    # Each input's main-effect variance is Vi = c_i^T * cov_block * c_i
+    # where c_i = [a1_i, a2_i, a3_i, M_ii].
+    cov_block = np.array(
+        [
+            [s2, s2 * es, 0.0, 0.0],
+            [s2 * es, (1 - es2) / 2, 0.0, 0.0],
+            [0.0, 0.0, (1 + es2) / 2 - es**2, -(s2**2) * es],
+            [0.0, 0.0, -(s2**2) * es, 2 * s2**2],
+        ]
+    )
+
+    # Main-effect variance for each input: quadratic form c_i^T Sigma c_i
+    Vi = np.array(
+        [
+            np.array([_A1[i], _A2[i], _A3[i], _M[i, i]])
+            @ cov_block
+            @ np.array([_A1[i], _A2[i], _A3[i], _M[i, i]])
+            for i in range(D)
+        ]
+    )
+
+    # Off-diagonal cross-terms x_i*x_k from the quadratic form contribute
+    # interaction variance Vpair_ik = ((M_ik + M_ki) * s2)^2 (symmetrized).
+    Vpair = ((_M + _M.T) * s2) ** 2
+    np.fill_diagonal(Vpair, 0.0)  # diagonal already captured in Vi via M_ii
+
+    # Total variance = sum of all main effects + sum of all unique pairwise interactions.
+    total_var = float(Vi.sum() + np.triu(Vpair, 1).sum())
+    S1 = Vi / total_var
+    # ST_i = main effect + all pairwise interactions involving input i
+    ST = (Vi + Vpair.sum(axis=1)) / total_var
+
+    # Second-order interaction matrix: S2_ij = Vpair_ij / total_var.
+    S2 = np.full((D, D), np.nan)
+    for j in range(D):
+        for k in range(j + 1, D):
+            val = Vpair[j, k] / total_var
+            S2[j, k] = val
+            S2[k, j] = val
+
+    return S1, ST, S2
+
+
+ANALYTICAL_S1, ANALYTICAL_ST, ANALYTICAL_S2 = analytical_indices()