ikpls.fast_cross_validation.jax

Contains the PLS class which implements fast cross-validation with partial least-squares regression using Improved Kernel PLS by Dayal and MacGregor: https://arxiv.org/abs/2401.13185 https://doi.org/10.1002/(SICI)1099-128X(199701)11:1%3C73::AID-CEM435%3E3.0.CO;2-%23

This is the JAX counterpart of ikpls.fast_cross_validation.numpy. Instead of parallelizing folds across CPU processes with joblib, it batches the folds with jax.vmap so the per-fold fits run together on a CPU/GPU/TPU. The per-fold training \(\mathbf{X}^{\mathbf{T}}\mathbf{W}\mathbf{X}\) and/or \(\mathbf{X}^{\mathbf{T}}\mathbf{W}\mathbf{Y}\) (centered/scaled/weighted using training-set statistics) are obtained by rank-update from the dataset-wide matrices using cvmatrix with its JAX backend (cvmatrix[jax]).

Both Improved Kernel PLS Algorithm #1 and #2 are supported, mirroring ikpls.fast_cross_validation.numpy:

Algorithm #2 fits each fold directly from the downdated \(\mathbf{X}^{\mathbf{T}}\mathbf{W}\mathbf{X}\) (K, K) and \(\mathbf{X}^{\mathbf{T}}\mathbf{W}\mathbf{Y}\) (K, M). It is efficient when X has more rows than columns.
Algorithm #1 gathers each fold’s training X (centered/scaled/weighted using the downdated training statistics) and the downdated \(\mathbf{X}^{\mathbf{T}}\mathbf{W}\mathbf{Y}\). It avoids forming \(\mathbf{X}^{\mathbf{T}}\mathbf{W}\mathbf{X}\) and is efficient when X has more columns than rows (but materializes per-fold training rows, so use batch_size to bound memory).

Algorithm #1 and #2 produce identical regression coefficients.

Note that the JAX implementations do not emit any underflow warning, in any path (single fit or cross-validation). Degenerate folds (training sets with too few non-zero weights to compute centering/scaling statistics) are still rejected up-front with a ValueError before the vmapped computation.

Author: Ole-Christian Galbo Engstrøm E-mail: ocge@foss.dk

Classes

PLS(algorithm, center_X, center_Y, scale_X, ...)

Implements fast cross-validation with partial least-squares regression using Improved Kernel PLS by Dayal and MacGregor: https://arxiv.org/abs/2401.13185 https://doi.org/10.1002/(SICI)1099-128X(199701)11:1%3C73::AID-CEM435%3E3.0.CO;2-%23

class ikpls.fast_cross_validation.jax.PLS(algorithm: int = 1, center_X: bool = True, center_Y: bool = True, scale_X: bool = True, scale_Y: bool = True, ddof: int = 0, copy: bool = True, dtype: str | type[Any] | dtype | SupportsDType = <class 'jax.numpy.float64'>)

Bases: object

Implements fast cross-validation with partial least-squares regression using Improved Kernel PLS by Dayal and MacGregor: https://arxiv.org/abs/2401.13185 https://doi.org/10.1002/(SICI)1099-128X(199701)11:1%3C73::AID-CEM435%3E3.0.CO;2-%23

The per-fold training kernel matrices are computed by rank-update with cvmatrix (JAX backend) and the folds are batched with jax.vmap.

Parameters:

algorithm (int, default=1) – Whether to use Improved Kernel PLS Algorithm #1 or #2. Generally, Algorithm #1 is faster if X has fewer rows than columns, while Algorithm #2 is faster if X has more rows than columns. Both yield identical regression coefficients.
center_X (bool, default=True) – Whether to center X before fitting by subtracting its row of column-wise means from each row. The row of column-wise means is computed on the training set for each fold to avoid data leakage.
center_Y (bool, default=True) – Whether to center Y. See center_X.
scale_X (bool, default=True) – Whether to scale X before fitting by dividing each row with the row of X’s column-wise standard deviations, computed on the training set for each fold.
scale_Y (bool, default=True) – Whether to scale Y. See scale_X.
ddof (int, default=0) – The delta degrees of freedom to use when computing the sample standard deviation. A value of 0 corresponds to the biased estimate of the sample standard deviation, while a value of 1 corresponds to Bessel’s correction.
copy (bool, default=True) – Whether to copy X, Y, and sample_weight when cross-validating.
dtype (DTypeLike, default=jnp.float64) – The float datatype to use. Using a lower precision than float64 will yield significantly worse results with an increasing number of components due to propagation of numerical errors.

Raises:

ValueError – If algorithm is not 1 or 2.

Notes

Any centering and scaling is undone before returning predictions so that predictions are on the original scale. If both centering and scaling are True, then the data is first centered and then scaled.