ikpls.jax

This file exposes the JAX Improved Kernel PLS implementation as a single PLS class parameterized by algorithm (1 or 2), mirroring ikpls.numpy.PLS. It dispatches to the algorithm-specific implementations in ikpls._impl; those classes are private and should be used through this entry point.

Requires the jax extra (pip install ikpls[jax]).

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

Classes

PLS([algorithm])

JAX Improved Kernel PLS (Algorithm #1 or #2).

class ikpls.jax.PLS(algorithm: int = 1, *args, **kwargs)

Bases: PLSBase

JAX Improved Kernel PLS (Algorithm #1 or #2).

ikpls.jax.PLS(algorithm=1) and ikpls.jax.PLS(algorithm=2) construct the Algorithm #1 and #2 implementations, respectively. Both share the same public API (fit, predict, transform, fit_transform, inverse_transform, cross_validate, and the stateless_* methods). Algorithm #1 is faster when the number of features K exceeds the number of samples N; Algorithm #2 is faster when N > K. Both yield identical results.

ikpls.jax.PLS subclasses the shared JAX PLS base, so it carries and documents the full public API listed above. Constructing it is a dispatcher that returns an instance of the selected algorithm-specific implementation (registered as a virtual subclass of this class), so:

isinstance(model, ikpls.jax.PLS) (and issubclass) is True;
model.algorithm is 1 or 2;
model.stateless_fit / jax.vmap operate directly on the implementation with no wrapping indirection or performance cost;
the algorithm-specific classes stay private in ikpls._impl.

Parameters:

algorithm (int, default=1) – Which Improved Kernel PLS algorithm to use (1 or 2).
**kwargs – Forwarded verbatim to the algorithm-specific implementation (center_X, center_Y, scale_X, scale_Y, ddof, dtype, verbose, …). See the implementation for the full signature.

Performs cross-validation for the Partial Least-Squares (PLS) model on given data. preprocessing_function will be applied before any potential centering and scaling as determined by self.center_X, self.center_Y, self.scale_X, and self.scale_Y. Any such potential centering and scaling is applied for each split using training set statistics to avoid data leakage from the validation set.

Parameters:

X (Array of shape (N, K)) – Predictor variables.
Y (Array of shape (N, M) or (N,)) – Response variables.
A (int) – Number of components in the PLS model.
folds (Array of Int of shape (N,)) – An array defining cross-validation splits. Each unique Int in folds corresponds to a different fold.
metric_function (Callable receiving arrays Y_val (N_val, M), Y_pred (A, N_val, M), and, if sample_weight is not None, also, weights_val (N_val,), and returning Any.) – Computes a metric based on true values Y_val and predicted values Y_pred. Y_pred contains a prediction for all A components.
metric_names (list of str) – A list of names for the metrics used for evaluation.
preprocessing_function (Callable or None, optional, default=None,) – A function that preprocesses the training and validation data for each fold. It should return preprocessed arrays for X_train, Y_train, X_val, and Y_val. If Callable, it should receive arrays X_train, Y_train, X_val, Y_val, and, if sample_weight is not None, also weights_train, and weights_val, and returning a Tuple of preprocessed X_train, Y_train, X_val, and Y_val.
show_progress (bool, default=True) – If True, displays a progress bar for the cross-validation.
sample_weight (Array of shape (N,) or None, optional, default=None) – Weights for each observation. If None, then all observations are weighted equally.
batch_size (int or None, optional, default=None) – The folds are batched together with jax.vmap (grouped by validation-set size so shapes are fixed). batch_size caps how many folds are vmapped at once; lower it to bound peak memory when there are many folds and/or a large number of features. None vmaps all folds of a given size together.

Returns:

metrics – A dictionary containing evaluation metrics for each metric specified in metric_names. The keys are metric names, and the values are lists of metric values for each cross-validation fold.

Return type:

dict[str, Any]

Raises:

ValueError – If sample_weight are provided and not all weights are non-negative, or if batch_size is not None and less than 1.

See also

_inner_cv: Performs cross-validation for a single fold and computes evaluation

metrics.

_update_metric_value_lists: Updates lists of metric values for each metric and

fold.

_finalize_metric_values: Organizes and finalizes the metric values into a

dictionary

stateless_fit_predict_eval: Fits the PLS model, makes predictions, and

evaluates

Notes

This method is used to perform cross-validation on the PLS model with different data splits and evaluate its performance using user-defined metrics.

Note that, because jax.vmap is used, metric_function and preprocessing_function must be JAX-traceable. The JAX implementation emits no underflow warning in any path (single fit or cross-validation).

Fits Improved Kernel PLS Algorithm #1 on X and Y using A components.

Parameters:

X (Array of shape (N, K)) – Predictor variables.
Y (Array of shape (N, M) or (N,)) – Response variables.
A (int) – Number of components in the PLS model.
sample_weight (Array of shape (N,) or None, optional, default=None) – Weights for each observation. If None, then all observations are weighted equally.

A

Number of components in the PLS model.

Type:: int

max_stable_components

The number of leading, numerically stable components – the index of the first component whose X-weight norm underflows below machine epsilon (X-Y cross-covariance exhausted) or whose score norm t^T t collapses relative to the largest score norm seen so far (a null-space direction past the numerical rank of X). Equals A when neither occurs. Computed on-device with no host callback, so it is also returned by stateless_fit and is therefore available per fit under jax.vmap.

Type:: int

B

PLS regression coefficients tensor.

Type:: Array of shape (A, K, M)

W

PLS weights matrix for X.

Type:: Array of shape (K, A)

P

PLS loadings matrix for X.

Type:: Array of shape (K, A)

Q

PLS Loadings matrix for Y.

Type:: Array of shape (M, A)

R

PLS weights matrix to compute scores T directly from original X.

Type:: Array of shape (K, A)

R_Y

Mapping from number of components to PLS weights matrix to compute scores U directly from original Y. Keys range from 1 to A. Values are arrays of shape (M, n_components) where n_components is the key. Values are computed lazily and cached upon first access. See Notes for more information.

Type:: Mapping[int, Array]

C

PLS Y-weights. In Improved Kernel PLS these equal the Y-loadings Q (matching scikit-learn’s PLSRegression, whose Y-weights equal its Y-loadings). Exposed as an alias of Q under the conventional Y-weights name; never referenced inside a jitted region.

Type:: Array of shape (M, A)

T

PLS scores matrix of X. Only assigned for Improved Kernel PLS Algorithm #1. IMPORTANT: If weights are provided, these are NOT the scores of X but instead weighted scores. In this case, scores can be computerd using transform.

Type:: Array of shape (N, A)

Returns:: self – Fitted model.
Return type:: PLS
Raises:: ValueError – If sample_weight are provided and not all weights are non-negative.

See also

stateless_fit: Performs the same operation but returns the output matrices

instead

Fits Improved Kernel PLS Algorithm #1 on X and Y using A components and returns T and U which are the scores of X and Y, respectively.

Parameters:

X (Array of shape (N, K)) – Predictor variables.
Y (Array of shape (N, M) or (N,)) – Response variables.
A (int) – Number of components in the PLS model.
sample_weight (Array of shape (N,) or None, optional, default=None) – Weights for each observation. If None, then all observations are weighted equally.

Returns:

T (Array of shape (N, A)) – PLS scores matrix of X.
U (Array of shape (N, A)) – PLS scores matrix of Y.

See also

transform: Transforms X and Y to their respective scores.
inverse_transform: Reconstructs X and Y from their respective scores.

Reconstructs X and Y from their respective scores.

Parameters:

T (Array of shape (N, n_X_components) or (N, A) or None, optional, default=None) – Scores of predictor variables.
U (Array of shape (N, n_Y_components) or (N, A) or None, optional, default=None) – Scores of response variables.

Returns:

X_reconstructed (Array of shape (N, K)) – If T is not None, returns the reconstructed X.
Y_reconstructed (Array of shape (N, M)) – If U is not None, returns the reconstructed Y.

Raises:

NotFittedError – If the model has not been fitted before calling inverse_transform().

See also

transform: Transforms X and Y to their respective scores.
fit_transform: Fits the model and returns the scores of X and Y.

Predicts with Improved Kernel PLS Algorithm #1 on X with B using n_components components. If n_components is None, then predictions are returned for all number of components.

Parameters:

X (Array of shape (N, K)) – Predictor variables.
n_components (int or None, optional) – Number of components in the PLS model. If None, then all number of components are used.

Returns:

Y_pred – If n_components is an int, then an array of shape (N, M) with the predictions for that specific number of components is used. If n_components is None, returns a prediction for each number of components up to A.

Return type:

Array of shape (N, M) or (A, N, M)

See also

stateless_predict: Performs the same operation but uses inputs B, X_mean,

None, None, and, instance.

Fits Improved Kernel PLS Algorithm #1 on X and Y using A components. Returns the internal matrices instead of storing them in the class instance.

Parameters:

X (Array of shape (N, K)) – Predictor variables.
Y (Array of shape (N, M) or (N,)) – Response variables.
A (int) – Number of components in the PLS model.
sample_weight (Array of shape (N,) or None, optional, default=None) – Weights for each observation. If None, then all observations are weighted equally.

Returns:

B (Array of shape (A, K, M)) – PLS regression coefficients tensor.
W (Array of shape (A, K)) – PLS weights matrix for X.
P (Array of shape (A, K)) – PLS loadings matrix for X.
Q (Array of shape (A, M)) – PLS Loadings matrix for Y.
R (Array of shape (A, K)) – PLS weights matrix to compute scores T directly from original X.
R_Y (Mapping[int, Array]) – Mapping from number of components to PLS weights matrix to compute scores U directly from original Y. Keys range from 1 to A. Values are arrays of shape (M, n_components) where n_components is the key. Values are computed lazily and cached upon first access. See Notes for more information.
T (Array of shape (A, N)) – PLS scores matrix of X. Only Returned for Improved Kernel PLS Algorithm #1.
max_stable_components (int) – The number of leading, numerically stable components – the index of the first component whose X-weight norm underflows below machine epsilon (X-Y cross-covariance exhausted) or whose score norm t^T t collapses relative to the largest score norm seen so far (a null-space direction past the numerical rank of X). Equals A when neither occurs.
X_mean (Array of shape (1, K) or None) – Mean of the predictor variables center_X is True, otherwise None.
Y_mean (Array of shape (1, M) or None) – Mean of the response variables center_Y is True, otherwise None.
X_std (Array of shape (1, K) or None) – Sample standard deviation of the predictor variables scale_X is True, otherwise None.
Y_std (Array of shape (1, M) or None) – Sample standard deviation of the response variables scale_Y is True, otherwise None.

See also

fit: Performs the same operation but stores the output matrices in the class

instance

Notes

For optimization purposes, the internal representation of all matrices (except B) is transposed from the usual representation.

Computes B with stateless_fit. Then computes Y_pred with stateless_predict. Y_pred is an array of shape (A, N, M). Then evaluates and returns the result of metric_function(Y_val, Y_pred).

Parameters:

X_train (Array of shape (N_train, K)) – Predictor variables.
Y_train (Array of shape (N_train, M) or (N_train,)) – Response variables.
A (int) – Number of components in the PLS model.
weights_train (Array of shape (N_train,) or None) – Weights for each observation. If None, then all observations are weighted equally.
X_val (Array of shape (N_val, K)) – Predictor variables.
Y_val (Array of shape (N_val, M) or (N_val,)) – Response variables.
weights_val (Array of shape (N_val,) or None) – Weights for each observation. If None, then all observations are weighted equally.
metric_function (Callable receiving arrays Y_val (N_val, M), Y_pred (A, N_val, M), and, if sample_weight is not None, also, weights_val (N_val,), and returning Any.) – Computes a metric based on true values Y_val and predicted values Y_pred. Y_pred contains a prediction for all A components.

Returns:

metric_function(Y_val, Y_pred, weights_val)

Return type:

Any.

See also

stateless_fit: Fits on X_train and Y_train using A components while

optionally, instead

stateless_predict: Computes Y_pred given predictor variables X and

regression

Predicts with Improved Kernel PLS Algorithm #1 on X with B using n_components components. If n_components is None, then predictions are returned for all number of components.

Parameters:

X (Array of shape (N, K)) – Predictor variables.
B (Array of shape (A, K, M)) – PLS regression coefficients tensor.
n_components (int or None, optional) – Number of components in the PLS model. If None, then all number of components are used.
X_mean (Array of shape (1, K) or None, optional, default=None) – Mean of the predictor variables. If None, then no mean is subtracted from X.
X_std (Array of shape (1, K) or None, optional, default=None) – Sample standard deviation of the predictor variables. If None, then no scaling is applied to X.
Y_mean (Array of shape (1, M) or None, optional, default=None) – Mean of the response variables. If None, then no mean is subtracted from Y.
Y_std (Array of shape (1, M) or None, optional, default=None) – Sample standard deviation of the response variables. If None, then no scaling is applied to Y.

Returns:

Y_pred – If n_components is an int, then an array of shape (N, M) with the predictions for that specific number of components is used. If n_components is None, returns a prediction for each number of components up to A.

Return type:

Array of shape (N, M) or (A, N, M)

See also

predict: Performs the same operation but uses the class instances of B,

None, None, None, and, arguments.

Transforms X and Y to their respective scores using n_components components. If n_components is None, then scores for all components up to A are returned.

Parameters:

X (Array of shape (N, K) or None, optional, default=None) – Predictor variables.
Y (Array of shape (N, M) or None, optional, default=None) – Response variables.
n_components (int or None, optional, default=None.) – Number of components in the PLS model. If None, then scores for all components up to A are returned.
copy (bool, default=True) – Whether to copy X and Y before potentially applying centering and scaling. If True, then the data is copied beforehand. If False, and dtype matches the type of X and Y, then centering and scaling is done inplace, modifying both arrays.

Returns:

T (Array of shape (N, n_components) or (N, A) or None) – X scores of X. If n_components is an int, then the scores up to n_components is used. If n_components is None, returns scores for all components up to A.
U (Array of shape (N, n_components) or (N, A) or None) – Y scores of Y. If n_components is an int, then the scores up to n_components is used. If n_components is None, returns scores for all components up to A.

Raises:

NotFittedError – If the model has not been fitted before calling transform().

See also

fit_transform: Fits the model and returns the scores of X and Y.
inverse_transform: Reconstructs X and Y from their respective scores.

Notes

If multiple calls to transform are made with Y provided and a previously seen value of n_components, then self.R_Y[n_components] is only computed once and stored for future use.

Any centering and scaling of X`and `Y is carried out before computation of the scores and is undone afterwards.