From 9638f641ec70972793e88dec772f0c1f0e32f121 Mon Sep 17 00:00:00 2001 From: Robrecht Cannoodt Date: Wed, 30 Sep 2026 09:48:51 +0200 Subject: [PATCH 1/2] Solve guanlab_dengkw_pm KRR by Cholesky --- CHANGELOG.md | 2 ++ src/methods/guanlab_dengkw_pm/script.py | 20 +++++++++++++++----- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6dde6c89..92e3274b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -124,6 +124,8 @@ * `scipenn`: Exit 99 (non-applicable) instead of raising a `ValueError` on anything other than GEX -> ADT (PR #77). +* `guanlab_dengkw_pm`: Solve the kernel ridge regression with `scipy.linalg.cho_factor()`/`cho_solve()` instead of `KernelRidge`, whose `scipy.linalg.solve()` raised a `MemoryError` or segfaulted beyond ~30k cells with the OpenBLAS in this image. Predictions are unchanged; the method failed on every dataset except `bmmc_multiome`. + # task_predict_modality 0.1.1 ## NEW FUNCTIONALITY diff --git a/src/methods/guanlab_dengkw_pm/script.py b/src/methods/guanlab_dengkw_pm/script.py index 9e57faa4..1313e45f 100644 --- a/src/methods/guanlab_dengkw_pm/script.py +++ b/src/methods/guanlab_dengkw_pm/script.py @@ -1,10 +1,10 @@ import sys import anndata as ad import numpy as np +from scipy.linalg import cho_factor, cho_solve from scipy.sparse import csc_matrix from sklearn.decomposition import TruncatedSVD from sklearn.gaussian_process.kernels import RBF -from sklearn.kernel_ridge import KernelRidge from threadpoolctl import threadpool_limits ## VIASH START @@ -29,6 +29,15 @@ if meta.get('cpus'): threadpool_limits(meta['cpus']) +# same as KernelRidge(alpha=alpha, kernel=kernel).fit(X, y).predict(X_test), solved by Cholesky: +# the scipy.linalg.solve() that KernelRidge uses fails beyond ~30k cells with this image's OpenBLAS +def kernel_ridge_predict(kernel, alpha, X, y, X_test): + K = kernel(X) + K.flat[::K.shape[0] + 1] += alpha + # K is symmetric, so its transpose is a Fortran-ordered view LAPACK can factorise in place + dual_coef = cho_solve(cho_factor(K.T, overwrite_a=True), y) + return kernel(X_test, X) @ dual_coef + ## Removed PCA and normalization steps, as they arr already performed with the input data print('Reading input files', flush=True) @@ -118,13 +127,14 @@ print(batch, flush=True) kernel = RBF(length_scale = scale) - krr = KernelRidge(alpha=alpha, kernel=kernel) print('Fitting KRR ... ', flush=True) - krr.fit( + y_pred_batch = kernel_ridge_predict( + kernel, + alpha, train_norm[input_train_mod1.obs.batch.isin(batch)], - train_gs[input_train_mod2.obs.batch.isin(batch)] + train_gs[input_train_mod2.obs.batch.isin(batch)], + test_norm ) - y_pred_batch = krr.predict(test_norm) if embedder_mod2 is not None: # map the predicted components back to the mod2 feature space y_pred_batch = y_pred_batch @ embedder_mod2.components_ From b64186e655810c5d60e39e3a644290478201f500 Mon Sep 17 00:00:00 2001 From: Robrecht Cannoodt Date: Wed, 30 Sep 2026 09:51:55 +0200 Subject: [PATCH 2/2] Explain guanlab_dengkw_pm KRR workaround --- src/methods/guanlab_dengkw_pm/script.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/methods/guanlab_dengkw_pm/script.py b/src/methods/guanlab_dengkw_pm/script.py index 1313e45f..1a0bdd81 100644 --- a/src/methods/guanlab_dengkw_pm/script.py +++ b/src/methods/guanlab_dengkw_pm/script.py @@ -29,8 +29,20 @@ if meta.get('cpus'): threadpool_limits(meta['cpus']) -# same as KernelRidge(alpha=alpha, kernel=kernel).fit(X, y).predict(X_test), solved by Cholesky: -# the scipy.linalg.solve() that KernelRidge uses fails beyond ~30k cells with this image's OpenBLAS +# Workaround for sklearn's KernelRidge: this computes +# KernelRidge(alpha=alpha, kernel=kernel).fit(X, y).predict(X_test), but solves the system itself. +# +# KernelRidge solves (K + alpha * I) w = y with scipy.linalg.solve(assume_a="pos"). In this image +# (scipy 1.17 and 1.18 with OpenBLAS) that call raises a MemoryError or segfaults once a half of the +# batches holds more than ~30-40k cells, no matter how much memory the job gets (200 GB was not enough +# for 40k cells, which need ~13 GB). So the method failed on every dataset except bmmc_multiome. +# +# K + alpha * I is symmetric positive definite, so solve() does a Cholesky factorisation anyway. +# cho_factor()/cho_solve() run that same factorisation directly and handle 46k cells in about a minute; +# the predictions are identical to KernelRidge's. They also factorise K in place, whereas KernelRidge +# copies it first, which saves one n x n float64 matrix of peak memory. +# +# Switch back to KernelRidge once scipy.linalg.solve() handles matrices of this size again. def kernel_ridge_predict(kernel, alpha, X, y, X_test): K = kernel(X) K.flat[::K.shape[0] + 1] += alpha