VKOGA (correlated-input indices)

python

jaxgsa.vkoga.analyze(
    problem, X, Y, *,
    gamma=None,
    ridge=None,
    max_centers=None,
    n_folds=10,
    n_outer=512,
    n_inner=128,
    n_variance=8192,
    n_bootstrap=0,
    conf_level=0.95,
    ci_method="quantile",
    key,
    batch_size=None,
    on_invalid="raise",
    verbose=True,
    keep_replicates=False,
) -> VKOGAResult | IrregularResult

analyze fits a Vectorial Kernel Orthogonal Greedy Algorithm surrogate to given (X, Y) data, then estimates the five correlated indices of Li et al. (2010) against that surrogate under independent inputs, a Gaussian copula, or a custom sampler declared on the problem. This surrogate-based sensitivity workflow follows Hilhorst et al. (2024). It is the only method in jaxgsa that separates correlated from uncorrelated contributions out of a single fitted surrogate, and the only one that gets a correlated total without re-running the model.

The kernel centres use the target-independent P-greedy rule of De Marchi et al. (2005): selection maximises the power function without reading Y. The Newton basis and regularised coefficient fit follow Santin & Haasdonk (2021). The VKOGA name comes from Wirtz & Haasdonk (2013), whose target-dependent selection differs from the rule implemented here.

The design-based alternative is jaxgsa.kucherenko, which evaluates the real model on a conditional-copula design instead of querying a surrogate.

Training data must cover the analysis distribution

The training data and the analysis distribution are separate choices. The analysis distribution is the one in problem (independent, a Gaussian copula, or a custom sampler). The surrogate must be accurate wherever that distribution puts its mass.

Both nested estimators draw from conditionals of the joint distribution. They never draw X_i from its whole marginal while the other inputs stay fixed. So every surrogate evaluation lies where the joint distribution puts mass. Training rows drawn from the analysis distribution therefore work.

An independent, space-filling design is a safe default in two cases. Use it when the correlation is not fixed yet. Use it also when you reuse one fit under several correlations. Avoid training data that misses regions the analysis covers. For example, do not train under a correlation of +0.8 and analyse under -0.8.

One measured case: Ishigami, Gaussian copula with rho_ab = 0.8 and rho_bc = -0.4, 400 training rows, 4 seeds, compared with Kucherenko at 2^16 base points. The RMSE of the indices was 0.007 when the surrogate trained on copula draws and 0.010 when it trained on independent draws. Check cv_rmse and ask whether training covers the points the analysis will query.

A run

python

import jax
jax.config.update("jax_enable_x64", True)

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)
rho = np.array([[1.0, 0.7, 0.0], [0.7, 1.0, 0.0], [0.0, 0.0, 1.0]])

X = monte_carlo(problem, 512, 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.vkoga.analyze(
    problem.with_correlation(rho), X, Y, key=jax.random.key(0), max_centers=120
)
for name in ("S_TC", "S_TU", "S_U", "S_C", "S_IU"):
    print(name, np.asarray(getattr(res, name)).round(4))

jaxgsa.vkoga.analyze
  problem: D=3 (x1, x2, x3)
    marginals: uniform=3
    dependence: gaussian copula
    output: N=512 runs, T=1 x K=1 output slice
    invalid: none found in 512 rows (policy 'raise')
  timing:
    fit + compute: 5.86 s
    n_centers: 120
    gamma: 7.533
    ridge: 0.01
    batch_size: auto (resolved from the memory budget)
  results: top 3 of 3 parameters by S_TC
    1. x2  S_TC=0.4913
    2. x1  S_TC=0.3447
    3. x3  S_TC=0.0009869

S_TC [0.3447 0.4913 0.001 ]
S_TU [0.4323 0.4185 0.247 ]
S_U [0.2535 0.401  0.0009]
S_C [0.0912 0.0903 0.    ]
S_IU [0.1788 0.0175 0.2461]

x3 has S_TC = 0.001 and S_TU = 0.247. It has essentially no first-order effect and a large interaction effect, so it is the case where prioritisation and fixing disagree completely. x1 and x2 each pick up about 0.09 of S_C, the share that arrives through their 0.7 correlation with each other. x3 correlates with nothing and its S_C is 0.

The same problem run through jaxgsa.kucherenko on the real model gives S1 = [0.325, 0.525, -0.001] against S_TC = [0.345, 0.491, 0.001], and ST = [0.412, 0.425, 0.243] against S_TU = [0.432, 0.419, 0.247]. The two methods are estimating the same two quantities, one through a surrogate and one through 28672 model runs.

Index reference

The correlated/uncorrelated decomposition follows Li et al. (2010), with the surrogate-based estimators described by Hilhorst et al. (2024).

Every index has shape (..., D). The leading axes follow the shape contract: (D,) for Y of shape (N,), (K, D) for (N, K), (T, K, D) for (N, T, K).

Index	Definition	Reading
`S_TC`	$V (E (Y ∣ X_{i})) / V (Y)$	Total correlated. What $X_{i}$ explains through itself plus what it explains through its correlation with the others. Use it to prioritise.
`S_TU`	$E (V (Y ∣ X_{\sim i})) / V (Y)$	Total uncorrelated. What only $X_{i}$ can explain, correlated pathways removed. Use it to decide what can be fixed.
`S_U`	independent part of $S_{T C}$	The contribution of $X_{i}$ alone.
`S_C`	`S_TC - S_U`	The correlation-borne part. It can be negative when a correlation opposes a direct effect.
`S_IU`	`S_TU - S_U`	Independent interactions.

Under independent inputs the five collapse to the familiar picture: S_TC is the first-order Sobol' index $S_{1}$ , S_TU is the total index $S_{T}$ , S_U equals S_TC, and S_C is zero.

Dependency structure

The dependence comes from the problem alone. analyze has no correlation argument. Declare a Gaussian copula with problem.with_correlation(R). The matrix must be symmetric with a unit diagonal. A non-positive-definite matrix is repaired by clipping its negative eigenvalues and renormalising the diagonal, with a JaxgsaWarning. This is a cheap heuristic, not Higham's nearest correlation matrix; the two can disagree by a visible margin. A problem with no correlation is treated as independent.

To fit a matrix from observed data use jaxgsa.sampling.fit_correlation(problem, X_data) and attach it with problem.with_correlation(...). Making that a separate step keeps it explicit which sample the copula came from. To analyse one fit under several correlations, call analyze once per problem.with_correlation(R).

The matrix used is result.problem.correlation. A custom sampler has no copula matrix.

Custom input dependence

Declare the sampler on the problem, with problem.with_sampler(f, sampler_dimension=k) or Problem(..., sampler=f). A problem with a sampler cannot also declare a correlation. The callback contract is the same as Kucherenko's. It transforms package-generated uniform points into joint or conditional physical samples. It must be deterministic for identical arguments. Dependence is not estimated from X.

X, Y train the surrogate. The sampler defines the integration distribution. Its joint marginals must match Problem, because the declared marginal CDFs still define the coordinate transform and component-function basis.

python

problem = jaxgsa.Problem(names=..., bounds=..., sampler=f, sampler_dimension=2)
res = jaxgsa.vkoga.analyze(problem, X_train, Y_train, key=jax.random.key(0))
res.dependence   # "custom"

A custom sampler and the Gaussian copula share one estimator. The custom path costs about the same time. The default sizes round up to powers of two for both. The required key controls integration and bootstrap reproducibility. The extra S_U, S_C, S_IU fields inherit the fitted additive-component approximation and do not establish a unique causal decomposition.

See the runnable custom-dependence example.

Fit and estimator controls

Argument	Default	What it changes
`gamma`	`None`	Gaussian RBF shape parameter. `None` cross-validates over a grid.
`ridge`	`None`	RKHS regularisation. `None` cross-validates over a grid.
`max_centers`	`None`	Cap on greedily selected kernel centres. `None` means 300, itself capped at `N`. Lower it to cut fit time; too low and the surrogate underfits, which `cv_rmse` reports.
`n_folds`	`10`	Folds for the hyperparameter cross-validation. At least 2.
`n_outer`	`512`	Outer (conditioning) sample size per parameter. Rounded up to the next power of two for Sobol' balance. Raise this when `S_TU` looks noisy: it draws a fresh inner block per outer point, so its sampling error falls as `n_outer` grows. It does not help `S_TC`, which reuses one inner block across every outer point.
`n_inner`	`128`	Inner (conditional) sample size per outer point, rounded up to a power of two. `S_TC` drops the i.i.d. inner-noise correction and leans on a large shared inner block, so do not lower `n_inner` to buy speed there: a small `n_inner` inflates a small `S_TC`, and raising `n_outer` will not fix it.
`n_variance`	`8192`	Sample size for the output variance and the component-function fit, rounded up to a power of two.
`key`	required, keyword-only	A `jax.random` key driving the quasi-random index integration. No default, because the indices are a Monte-Carlo estimate. Use `jax.random.key(0)` for reproducibility.
`batch_size`	`None`	Rows per device call. It bounds every row-wise step: the surrogate evaluations behind `predict`, and the nested conditional draws. `None` derives one from the memory budget. It never changes the answer.
`on_invalid`	`"raise"`	One training row is one unit, so `"drop"` removes the `(X, Y)` pair. The check runs before the hyperparameter search on purpose: one non-finite `Y` makes every cross-validation score non-finite, and the caller would otherwise meet a `RuntimeError` about failed kernel solves that names the wrong cause. `"drop"` needs at least `n_folds` rows to survive.
`verbose`	`True`	Prints the summary block shown above.

Leaving both gamma and ridge to cross-validation costs a 10x10 grid of k-fold refits, which dominates the runtime. Pass both once you know good values; the fitted values are reported on the result and in the verbose block.

Confidence intervals

n_bootstrap=0 is the default and is the right setting for a routine run. Every replicate refits the kernel surrogate and re-runs the nested conditional integration, so n_bootstrap=100 costs roughly a hundred analyses.

What the interval measures is total uncertainty in the reported index given this training sample. The training rows are resampled with replacement, a new surrogate is fitted to each resample, and a fresh quasi-random stream integrates against it. It folds the sampling error of (X, Y) together with the Monte-Carlo error of the integration, which is what you want when the question is "how much would this index move if I had drawn a different training set". Fitting once and resampling only the integration would instead measure what n_outer/n_inner cost you, which is the part you can shrink for free by raising them.

Two things stay fixed, so the interval is a conditional statement. gamma and ridge are the values cross-validated on the full sample and are not re-selected per replicate. And the interval is about the surrogate's indices: it cannot see the error the surrogate makes against the true model. cv_rmse is what reports that.

Intervals cover the five index fields only. variance, rmse, cv_rmse and n_centers describe the reported fit itself, so an interval over other fits would not be about the thing they name.

conf_level, ci_method and keep_replicates behave as on every other method.

VKOGAResult

Indices: S_TC, S_TU, S_U, S_C, S_IU, each (..., D), with the matching *_conf fields of shape (2, ..., D) when n_bootstrap > 0.

Diagnostics:

Field	Meaning
`S1`, `ST`	Read-only aliases of `S_TC` and `S_TU`, for comparison with Kucherenko. Under independent inputs they are the first-order and total Sobol' indices.
`variance`	Output variance under the correlated input measure, per output slice.
`n_centers`	Kernel centres in the fitted surrogate.
`gamma`, `ridge`	The hyperparameters the fit used, cross-validated or passed.
`rmse`	Training-fit RMSE, per output slice.
`cv_rmse`	Cross-validated RMSE, per output slice. This is the one that tells you whether to trust the indices at all. `None` when you fixed both `gamma` and `ridge`.
`explained_variance`	`1 - rmse**2 / var(Y)` per output slice, in-sample.
`dependence`	The dependence the indices were computed under: `"independent"`, `"gaussian"` or `"custom"`. It always equals `problem.dependence`.
`invalid`, `ci`	The non-finite report and the interval provenance.

Operations: result.predict(X_new, batch_size=None) evaluates the fitted surrogate, and result.to_dataset(time_coords=None) gives the labeled xarray view.

VKOGA has no shapley() method. Shapley effects need a variance decomposition indexed by parameter subsets. A kernel expansion is a sum over centres and every centre involves every parameter, so there is no membership matrix to allocate from. Use jaxgsa.hdmr or jaxgsa.pce for Shapley effects.

Enable float64

The coefficient solve forms $A^{⊤} A$ , which squares the condition number of the cross kernel. float32 cannot carry that for a small gamma. Call jax.config.update("jax_enable_x64", True) before fitting. analyze emits a JaxgsaWarning when x64 is off.

What it refuses

analyze raises ValueError when X/Y violate the shape contract, when gamma or ridge is not finite and positive, when a size argument is out of range, when on_invalid is unknown or refuses the sample, and in two method-specific cases:

Fewer than two parameters. Conditioning on the other parameters means nothing at D = 1.
Any categorical parameter. The isotropic RBF needs a continuous CDF map per coordinate, and a step-CDF coordinate breaks both the kernel metric and the copula conditionals. Use jaxgsa.optimal_transport, jaxgsa.borgonovo, jaxgsa.pawn, or the Saltelli-based Sobol pipeline.

It raises RuntimeError when every cross-validation score is non-finite. Since the non-finite check now runs first, that is a genuine solver failure.

It warns in four cases: a zero-variance output slice; float32; a cross-validated error that is a large fraction of the output standard deviation, which makes every index untrustworthy; and S_U clipped to S_TU by a wide margin, which says the additive component functions cannot represent the model's interactions.

References

Hilhorst, P. L. J., Quicken, S., van de Vosse, F. N. & Huberts, W. (2024). Efficient sensitivity analysis for biomechanical models with correlated inputs. International Journal for Numerical Methods in Biomedical Engineering, 40(2), e3797. Surrogate-based correlated sensitivity analysis.
Li, G., Rabitz, H., Yelvington, P. E., Oluwole, O. O., Bacon, F., Kolb, C. E. & Schoendorf, J. (2010). Global sensitivity analysis for systems with independent and/or correlated inputs. Journal of Physical Chemistry A, 114(19), 6022–6032. Structural and correlative variance decomposition.
De Marchi, S., Schaback, R. & Wendland, H. (2005). Near-optimal data-independent point locations for radial basis function interpolation. Advances in Computational Mathematics, 23, 317–330. P-greedy centre selection implemented here.
Santin, G. & Haasdonk, B. (2021). Kernel methods for surrogate modeling. In Model Order Reduction, Volume 1: System- and Data-Driven Methods and Algorithms, 311–354. De Gruyter. Newton basis and regularised kernel approximation.
Wirtz, D. & Haasdonk, B. (2013). A vectorial kernel orthogonal greedy algorithm. Dolomites Research Notes on Approximation, 6, 83–100. Source of the VKOGA name and a target-dependent selection variant.

See the VKOGA example, Methods, and the API overview.

VKOGA (correlated-input indices) ​

Training data must cover the analysis distribution ​

A run ​

Index reference ​

Dependency structure ​

Custom input dependence ​

Fit and estimator controls ​

Confidence intervals ​

VKOGAResult ​

Enable float64 ​

What it refuses ​

References ​

VKOGA (correlated-input indices)

Training data must cover the analysis distribution

A run

Index reference

Dependency structure

Custom input dependence

Fit and estimator controls

Confidence intervals

VKOGAResult

Enable float64

What it refuses

References