CVMatrix

The cvmatrix package implements the fast cross-validation algorithms by Engstrøm and Jensen [1] for computation of training set $\mathbf{X}^{\mathbf{T}}\mathbf{X}$ and $\mathbf{X}^{\mathbf{T}}\mathbf{Y}$ in a cross-validation setting. In addition to correctly handling arbitrary row-wise pre-processing, the algorithms allow for and efficiently and correctly handle any combination of column-wise centering and scaling of X and Y based on training set statistical moments.

For an implementation of the fast cross-validation algorithms combined with Improved Kernel Partial Least Squares [2], see the Python package ikpls by Engstrøm et al. [3].

The cvmatrix software package now also features weigthed matrix produts $\mathbf{X}^{\mathbf{T}}\mathbf{W}\mathbf{Y}$ without increasing time or space complexity compared to the unweighted case. This is due to a generalization of the algorithms by Engstrøm and Jensen [1]. A new article formally describing the generalization is to be announced.

Installation

• Install the package for Python3 using the following command:

  pip3 install cvmatrix

• Now you can import the class implementing all the algorithms with:

  from cvmatrix import CVMatrix

• You can also install the optional JAX dependency to enable the JAX backend of CVMatrix:

  pip3 install "cvmatrix[jax]"

• Now, you can select the JAX backend by passing backend="jax" to the constructor:

  from cvmatrix import CVMatrix
  
  cvm = CVMatrix(
      center_X=True, center_Y=True, scale_X=True, scale_Y=True, backend="jax"
  )

  With backend="jax", the per-fold training-matrix computations (training_XTX, training_XTY, training_XTX_XTY, and training_statistics) use jax.numpy and can be traced by jax.jit and batched over folds with jax.vmap on CPUs, GPUs, and TPUs. The default backend="numpy" is unchanged and remains byte-for-byte identical to previous releases.

Prerequisites for the JAX backend

The JAX backend supports running on CPU, GPU, and TPU.

• To enable NVIDIA GPU execution, install JAX and CUDA with:

  pip3 install -U "jax[cuda13]"

• To enable Google Cloud TPU execution, install JAX with:

  pip3 install -U "jax[tpu]"

These are typical installation instructions that will be what most users are looking for. For customized installations, follow the instructions from the JAX Installation Guide.

To ensure that the JAX backend uses float64, set the environment variable JAX_ENABLE_X64=True as per the Common Gotchas. CVMatrix also enables this automatically when constructed with backend="jax" and a 64-bit dtype (the default). Alternatively, float64 can be enabled with the following function call:

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

Quick Start

Use the cvmatrix package for fast computation of training set kernel matrices

import numpy as np
from cvmatrix import CVMatrix
from cvmatrix import Partitioner

N = 100  # Number of samples.
K = 50  # Number of features.
M = 10  # Number of targets.

X = np.random.uniform(size=(N, K)) # Random X data
Y = np.random.uniform(size=(N, M)) # Random Y data
folds = np.arange(100) % 5 # 5-fold cross-validation

# Weights must be non-negative. If centering or scaling is used, the sum of weights
# for any training partition must be greater than zero. If scaling is used, the
# number of non-negative weights for any training partition must be greater than
# the ddof provided in the constructor.
weights = np.random.uniform(size=(N,)) + 0.1

# Instantiate CVMatrix
cvm = CVMatrix(
    center_X=True, # Cemter around the weighted mean of X.
    center_Y=True, # Cemter around the weighted mean of Y.
    scale_X=True, # Scale by the weighted standard deviation of X.
    scale_Y=True, # Scale by the weighted standard deviation of Y.
)
# Fit on X, Y, and weights
cvm.fit(X=X, Y=Y, weights=weights)

# Instantiate Partitioner
p = Partitioner(folds=folds)

# Compute training set XTWX and/or XTWY for each fold
for fold in p.folds_dict:
    val_indices = p.get_validation_indices(fold)
    # Get both XTWX, XTWY, and weighted statistics
    result = cvm.training_XTX_XTY(val_indices)
    (training_XTWX, training_XTWY) = result[0]
    (training_X_mean, training_X_std, training_Y_mean, training_Y_std) = result[1]
    
    # Get only XTWX and weighted statistics for X.
    # Weighted statistics for Y are returned as None as they are not computed when
    # only XTWX is requested.
    result = cvm.training_XTX(val_indices)
    training_XTWX = result[0]
    (training_X_mean, training_X_std, training_Y_mean, training_Y_std) = result[1]
    
    # Get only XTWY and weighted statistics
    result = cvm.training_XTY(val_indices)
    training_XTWY = result[0]
    (training_X_mean, training_X_std, training_Y_mean, training_Y_std) = result[1]

Examples

In examples, you will find:

• Compute training matrices with CVMatrix
• Compute training matrices with the JAX backend, batched over folds with jax.vmap on CPU/GPU/TPU

Benchmarks

In benchmarks, we have benchmarked cross-validation of the fast algorithms in cvmatrix against the baseline algorithms implemented in NaiveCVMatrix.

Left: Benchmarking cross-validation with the CVMatrix implementation versus the baseline implementation using three common combinations of (column-wise) centering and scaling. Right: Benchmarking cross-validation with the CVMatrix implementation for all possible combinations of (column-wise) centering and scaling. Here, most of the graphs lie on top of eachother. In general, no preprocessing is faster than centering which, in turn, is faster than scaling.

The figures below add the optional JAX backend to the comparison (NumPy vs. JAX, and the JAX execution modes). They are measured multi-threaded (no thread limit; GPU on an NVIDIA RTX 3090 Ti) and exclude NaiveCVMatrix:

Left: CVMatrix with the NumPy backend versus the JAX backend (warm, JIT-compiled jax.vmap over folds) on CPU and GPU, for the three common preprocessing combinations. Right: The JAX backend under no-JIT (eager vmap), cold-JIT (compilation + run) and warm-JIT (run only), on CPU and GPU. JIT compilation is a large fixed cost: on the GPU, no-JIT beats cold-JIT for a single run, whereas on the CPU, JIT (even cold) is far faster than eager vmap at a large number of folds.

Contribute

To contribute, please read the Contribution Guidelines.

References

1. Engstrøm, O.-C. G. and Jensen, M. H. (2025). Fast partition-based cross-validation with centering and scaling for $\mathbf{X}^\mathbf{T}\mathbf{X}$ and $\mathbf{X}^\mathbf{T}\mathbf{Y}$. Journal of Chemometrics, 39(3).
2. Dayal, B. S. and MacGregor, J. F. (1997). Improved PLS algorithms. Journal of Chemometrics, 11(1), 73-85.
3. Engstrøm, O.-C. G. and Dreier, E. S. and Jespersen, B. M. and Pedersen, K. S. (2024). IKPLS: Improved Kernel Partial Least Squares and Fast Cross-Validation Algorithms for Python with CPU and GPU Implementations Using NumPy and JAX. Journal of Open Source Software, 9(99).

Funding

• Up until May 31st 2025, this work has been carried out as part of an industrial Ph. D. project receiving funding from FOSS Analytical A/S and The Innovation Fund Denmark. Grant number 1044-00108B.
• From June 1st 2025 and onward, this work is sponsored by FOSS Analytical A/S.