jaxgsa

Appearance

PCE ​

python
jaxgsa.pce.analyze(
    problem, X, Y, *,
    order=3,
    ridge=1e-8,
    fit_ratio=0.5,
    batch_size=None,
    n_bootstrap=0,
    conf_level=0.95,
    ci_method="quantile",
    key=None,
    on_invalid="raise",
    verbose=True,
    keep_replicates=False,
) -> PCEResult | IrregularResult

analyze fits an orthogonal polynomial surrogate to arbitrary (X, Y) pairs, then reads the first-, total- and second-order Sobol indices straight off the expansion coefficients (Sudret, 2008). No structured design, and no extra model evaluations once the fit is done. On a smooth response it needs far fewer samples than a Monte-Carlo Sobol estimator.

It only captures what the polynomial can represent, so check cv_rmse before you trust the indices.

A run ​

python
import numpy as np
import jaxgsa
from jaxgsa.sampling import monte_carlo

problem = jaxgsa.Problem(names=("x1", "x2", "x3"), bounds=((-np.pi, np.pi),) * 3)
X = monte_carlo(problem, 2000, seed=0)
Y = np.sin(X[:, 0]) + 7.0 * np.sin(X[:, 1]) ** 2 + 0.1 * X[:, 2] ** 4 * np.sin(X[:, 0])

res = jaxgsa.pce.analyze(problem, X, Y, order=8)
print("S1      ", np.asarray(res.S1).round(4))
print("ST      ", np.asarray(res.ST).round(4))
print("rmse    ", float(res.rmse))
print("cv_rmse ", float(res.cv_rmse))
print("streamed", res.streamed)
jaxgsa.pce.analyze
  problem: D=3 (x1, x2, x3)
    marginals: uniform=3
    dependence: independent
    output: N=2000 runs, T=1 x K=1 output slice
    invalid: none found in 2000 rows (policy 'raise')
  timing:
    fit + estimator (includes compile on the first call): 1.544 s
    order: 8
    fit: single-pass
    batch_size: auto (resolved from the memory budget)
  results: top 3 of 3 parameters by ST
    1. x1  ST=0.5576
    2. x2  ST=0.4425
    3. x3  ST=0.243

S1       [0.3146 0.4424 0.    ]
ST       [0.5576 0.4425 0.243 ]
rmse     0.06354028731584549
cv_rmse  0.08116015046834946
streamed False

The analytic Ishigami values are S1 = [0.314, 0.442, 0] and ST = [0.558, 0.442, 0.244]. A degree-8 expansion on 2000 rows reproduces them to three decimals, and cv_rmse = 0.081 against an output standard deviation of 3.59 says the surrogate is not the limiting factor.

Drop to the default order=3 on the same data and the answer falls apart. The cubic basis cannot represent sin(x2)^2 or x3^4 * sin(x1), and the fit misattributes the variance to x1:

S1 [0.6499 0.0585 0.0009]   ST [0.9396 0.0627 0.2885]
cv_rmse 2.709   explained_variance 0.4437

Nothing in the index arrays says they are wrong. cv_rmse jumping from 0.08 to 2.71, and explained_variance at 0.44, are the only signal. Read them first, every time.

verbose=True is the default and printed the block. Pass verbose=False for a silent run.

Arguments ​

ArgumentDefaultWhat it changes
order3Maximum total polynomial degree. Higher orders capture sharper nonlinearity and larger interactions, but the term count C(D+order, order) grows fast and needs more rows to fit. It is reduced automatically, with a warning, when the term count would exceed fit_ratio * N.
ridge1e-8Tikhonov regularization on the least-squares fit. The default is small enough to be a guard against a singular design and nothing more. Raise it when coefficients look unstable, which happens with noisy Y or near-duplicate rows. A ridge of about 1e-3 or more hides a singular design from the rank check. Constant, duplicate and linearly dependent inputs still raise.
fit_ratio0.5The terms-to-samples ceiling that triggers the order reduction. Lower it to demand more samples per term, which gives a more conservative fit that is harder to overfit. Must be <= 1: a value above 1 asks for more terms than one row per term can support. order is never reduced below 1 (D + 1 terms), so a row count too small even for that raises ValueError instead of fitting an underdetermined system.
batch_sizeNoneRows per batch during the fit. See below.
n_bootstrap0Row resamples. See below.
keyNoneA jax.random key for the resampling. Required when n_bootstrap > 0.
on_invalid"raise"Policy for non-finite rows. One row is one unit, so "drop" removes the (X, Y) pair. This check matters more here than for a reduction-only method: one NaN reaches the QR factorisation and poisons every coefficient, every leave-one-out RMSE and every index.
verboseTruePrints the summary block shown above.
keep_replicatesFalseRetains the per-replicate index arrays on result.ci.

conf_level and ci_method behave as on every other method.

The two fit paths ​

There are two, and they give the same numbers. The single-pass path builds the full design matrix. The streamed path reads rows in batches, so it holds far less memory. Both solve the same ridge least squares by QR and compute the same exact leave-one-out error. They differ only in the order of the float32 sums.

The streamed path engages when batch_size is an integer below N, or when the estimated resident memory of the single-pass fit (design matrix, QR factorization, LOO residuals) exceeds the active memory budget, about 512 MiB by default. Set that budget with jaxgsa.config.set_memory_budget. batch_size >= N is one full block, which is the single-pass fit.

result.streamed reports which path ran, and the verbose block prints it as fit: single-pass or the streamed equivalent. Read it when a fit takes much longer than you expect: True means the budget engaged.

The single-pass path holds two (N, n_terms) arrays at its peak, not three: the design matrix and the leverage solve inside the leave-one-out step. The coefficient solve reads Phi.T @ Y directly into the normal-equations solve, so it never separately builds and holds an (n_terms, N) product across the leave-one-out step. This changes the size at which streaming starts, but not the fitted numbers.

Confidence intervals ​

This is the expensive kind of bootstrap. Every replicate refits the whole expansion, a fresh design matrix and a fresh QR solve, so twenty replicates cost about twenty fits. That is an order of magnitude above the row resample of a method like jaxgsa.pawn, which only re-reduces numbers it already has. Hence the 0 default, and hence a small n_bootstrap when you do want intervals.

The interval measures the sampling variability of (X, Y) propagated through the fit. It does not measure truncation error. An expansion too coarse to represent the model gives a biased index, every replicate inherits the same bias, and the interval stays narrow and wrong. cv_rmse and explained_variance are what report that.

PCEResult ​

FieldShapeMeaning
S1, ST(D,) / (K, D) / (T, K, D)First-order and total Sobol indices.
S2(D, D) / (K, D, D) / (T, K, D, D)Second-order indices.
S1_conf, ST_conf, S2_confleading (2, ...)None when n_bootstrap=0.
coefficients, multi_indexThe fitted expansion, reused by predict.
orderThe degree the fit actually used, after any reduction.
rmse() / (K,) / (T, K)In-sample RMSE of the fit per output slice, in the units of Y.
cv_rmsesameExact leave-one-out RMSE per output slice, in the units of Y.
explained_variancesame1 - rmse**2 / Var(Y) per output slice: the in-sample R^2, so it stays in [0, 1]. NaN on a constant slice.
streamedWhich fit path ran.
problem, invalid, ciProblem, non-finite report, interval provenance.

Operations:

  • result.predict(X_new, batch_size=None) evaluates the fitted expansion on new inputs.
  • result.shapley() derives Shapley effects from the fit. See Shapley effects.
  • result.to_dataset(time_coords=None) gives the labeled xarray view.

All output slices share one basis and are fitted in a single multi-right-hand- side solve. Indices are computed independently per (t, k) slice.

Aggregating outputs ​

result.aggregate(*, weights="variance") collapses the T and K axes into one index per input. With weights="variance" each output counts in proportion to its variance. This is the generalised Sobol' index of Gamboa, Janon, Klein and Lagnoux (2014), also used by Lamboni, Monod and Makowski (2011). It is the share of the total output variance that the input explains. With weights="equal" every output counts the same.

The result keeps its class. Its output axes have size 1. A constant output has weight 0. A result with one output slice comes back unchanged.

The weights come from the variance of the fitted surrogate under the input measure. That is the sum of the squared non-constant coefficients. A slice whose indices are all NaN and whose coefficients are finite (a constant output) gets weight 0. A slice with non-finite coefficients (a failed fit) has an unknown variance, so the aggregate is NaN. aggregate applies to S1, ST and S2.

The aggregated PCEResult is a summary, not a surrogate. Its coefficients are NaN, so predict() and shapley() raise ValueError: call them before aggregate(). Its rmse, cv_rmse and explained_variance are None.

With keep_replicates=True, each bootstrap draw is aggregated and the interval is drawn again from those draws. Otherwise the intervals are NaN. The weights stay fixed at their point estimates.

effective_order ​

python
jaxgsa.pce.effective_order(problem, n_samples, *, order=3, fit_ratio=0.5) -> int

Reports the degree a fit would actually use, given only the sample size. The reduction is real information: an expansion fitted at degree 2 does not represent the cubic effects a caller who asked for degree 3 expects to see.

python
jaxgsa.pce.effective_order(problem, 100, order=8)
4

analyze reports the same number twice, as a warning and as result.order. indices reports it neither way, because it is traceable and so has no result object and no place for a side effect. This function is where an indices caller reads it, and it answers before the model has been run.

What it refuses ​

ValueError when X fails validation against the problem, when Y's layout cannot be resolved against X's row count, when batch_size is given and is not a positive integer, when on_invalid is unknown or refuses the sample, when n_bootstrap > 0 without a key, when fit_ratio > 1, when the row count (after on_invalid has run) cannot fit even the order-1 expansion at the given fit_ratio, and in two structural cases:

It also raises ValueError when an input is constant, when two inputs are equal in every row, and when an input is a linear combination of other inputs (for example c = 1 - a or c = a + b). The check reads X, so it does not depend on the dtype or on ridge. In float32 it also raises when the QR pivot share of the design falls below 1e-5. float64 only warns, below 1e-9.

Near that precision band, float32 can warn about an overfit even when the indices are exact. Example: three Gaussian inputs at order 13 give exact indices, but cv_rmse is about 2 times the standard deviation of Y. The leave-one-out error is the first number to lose precision in float32, so trust the indices and re-run in float64 to confirm.

Traceable core ​

jaxgsa.pce.indices(problem, X, Y, *, order=3, ridge=1e-8, fit_ratio=0.5, batch_size=None) returns a PCEIndices named tuple (S1, ST, S2) of bare arrays with none of the checks, so it composes with jit, vmap and jacrev. It has no result object, so it reports neither the reduced order nor cv_rmse. Use effective_order for the first.

References ​

See the PCE example and the API overview.

Released under the BSD-3-Clause License.