From fa22a2a3e07c55a967660f5cfa55f87324ed77e4 Mon Sep 17 00:00:00 2001
From: benjaminfreyuu <benjamin.frey.de@gmail.com>
Date: Tue, 18 Aug 2026 13:24:15 +0200
Subject: [PATCH 1/5] senkin: pin LightGBM threads and skip non-applicable
 datasets

The 2026-08-14 run had senkin_tmp_train hit the 8h walltime on the CITE datasets
and OOM/crash on the rest. Two causes:

1. LightGBM set no num_threads, so it defaulted to one thread per core the
   container *sees* (the whole node) while the job is cgroup-throttled to
   meta["cpus"] -- the threads oversubscribe and thrash. The original solution
   ran on Kaggle with dedicated cores so never hit this. Pin num_threads to
   meta["cpus"] (results-preserving; same class of fix as guanlab in #59).

2. senkin has no applicability guard, so it ran on all 8 datasets including
   Multiome and the ADT->GEX swap, where it wastes hours and OOMs. Add an
   exit_non_applicable guard (via src/utils/exit_codes.py) so it only runs on the
   GEX->ADT CITE direction, and point the component test at bmmc_cite/normal.

Verified: viash test passes (2/2) on the normal (GEX->ADT) direction; the guard
correctly skips the swap direction.

Note: pinning threads removes the oversubscription, but senkin runs 4 LightGBM
models (134 targets x 5 folds) + 2 NNs in one job, which the original spread
across separate multi-hour Kaggle sessions - it may still need a walltime above
8h or a lighter config on the full data.
---
 .../senkin_tmp_train/config.vsh.yaml          | 16 ++++++++++++
 .../senkin_tmp/senkin_tmp_train/script.py     | 26 ++++++++++++++++++-
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
index 6f549b24..46642bb3 100644
--- a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
@@ -1,8 +1,18 @@
 __merge__: ../../../api/comp_method_train.yaml
 name: senkin_tmp_train
+info:
+  # The api default test inputs are bmmc_cite/swap (ADT->GEX). senkin predicts protein
+  # (ADT) from RNA (GEX), so point the component test at the 'normal' (GEX->ADT)
+  # direction, otherwise the new modality guard correctly exits non-applicable.
+  test_setup:
+    normal_direction:
+      input_train_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/train_mod1.h5ad
+      input_train_mod2: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/train_mod2.h5ad
+      input_test_mod1: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/test_mod1.h5ad
 resources:
   - path: script.py
     type: python_script
+  - path: /src/utils/exit_codes.py
 arguments:
   - name: "--n_folds"
     type: integer
@@ -62,3 +72,9 @@ runners:
   - type: nextflow
     directives:
       label: [highmem, hightime, midcpu, gpu]
+# Override the api default (bmmc_cite/swap = ADT->GEX). senkin only supports GEX->ADT.
+test_resources:
+  - type: python_script
+    path: /common/component_tests/run_and_check_output.py
+  - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal
+    dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal
diff --git a/src/methods/senkin_tmp/senkin_tmp_train/script.py b/src/methods/senkin_tmp/senkin_tmp_train/script.py
index 11760cb5..ef3aec54 100644
--- a/src/methods/senkin_tmp/senkin_tmp_train/script.py
+++ b/src/methods/senkin_tmp/senkin_tmp_train/script.py
@@ -1,6 +1,7 @@
 import gc
 import logging
 import pickle
+import sys
 
 import anndata as ad
 import numpy as np
@@ -29,9 +30,12 @@
     "nn_epochs": 100,
     "n_tsvd_components": 100,
 }
-meta = {"name": "senkin_tmp"}
+meta = {"name": "senkin_tmp", "resources_dir": "src/methods/senkin_tmp/senkin_tmp_train"}
 ## VIASH END
 
+sys.path.append(meta["resources_dir"])
+from exit_codes import exit_non_applicable
+
 
 def _to_dense(X):
     return X.toarray() if issparse(X) else np.array(X)
@@ -60,6 +64,17 @@ def _split(b):
 adata_prot_train = ad.read_h5ad(par["input_train_mod2"])
 adata_rna_test   = ad.read_h5ad(par["input_test_mod1"])
 
+# senkin is a CITE-seq method that predicts protein (ADT) from RNA (GEX); it treats
+# mod1 as RNA and mod2 as protein. Skip the datasets/directions it cannot handle
+# (e.g. Multiome, or the ADT->GEX swap) instead of running for hours and OOMing.
+_mod1 = adata_rna_train.uns.get("modality")
+_mod2 = adata_prot_train.uns.get("modality")
+if _mod1 != "GEX" or _mod2 != "ADT":
+    exit_non_applicable(
+        f"senkin only supports predicting protein (ADT) from RNA (GEX); "
+        f"got mod1={_mod1!r}, mod2={_mod2!r}."
+    )
+
 adata_rna_train.obs  = _parse_batch(adata_rna_train.obs)
 adata_prot_train.obs = _parse_batch(adata_prot_train.obs)
 adata_rna_test.obs   = _parse_batch(adata_rna_test.obs)
@@ -133,6 +148,15 @@ def _split(b):
 boost_rounds = par["lgbm_boost_rounds"]
 early_stop   = par["lgbm_early_stopping"]
 
+# Pin LightGBM to the allocated cores. Its default (num_threads=0) spawns one thread
+# per core the container *sees* (the whole node) while the job is cgroup-throttled to
+# meta["cpus"], so the threads oversubscribe and thrash -- the same class of slowdown
+# fixed for guanlab in #59. Leave the library default when cpus is unknown (local runs).
+_n_threads = meta.get("cpus")
+if _n_threads:
+    for _p in (lgbm_params_1, lgbm_params_2, lgbm_params_3, lgbm_params_4):
+        _p["num_threads"] = _n_threads
+
 # Protein targets
 Y_prot_raw = _to_dense(adata_prot_train.layers.get("counts", adata_prot_train.X)).astype(np.float64)
 

From f6b9a3f8b682879cdf9cc6b169df594213eaf7a8 Mon Sep 17 00:00:00 2001
From: Vladimir Shitov <vladimirshitov98@gmail.com>
Date: Tue, 8 Sep 2026 00:51:38 +0200
Subject: [PATCH 2/5] senkin_tmp: match the original pipeline and make the
 wrapper dataset agnostic

Uses the fixed senkin-tmp-cite-pred (fix/match-original-pipeline) and
restores the original feature blocks: log1p(CPM) inputs, CLR-TSVD 200,
correlated genes selected on the log-normalized RNA per batch (plus genes
whose name matches a protein), TSVD-100 and PCA-64 of the sqrt/z-score/
batch-median normalization, and the four LightGBM prediction blocks, all
z-scored per cell before the two neural networks.

The batch column is used as is (or a single batch when absent) instead of
parsing day/donor out of it, which only made sense for the NeurIPS 2022
naming. Predictions are rescaled to the target scale with one global affine
transform fitted on the out-of-fold training predictions, so RMSE/MAE are
meaningful while every per-cell and per-protein correlation is unchanged.
The predict step checks that the test cells match the trained bundle and
uses light resource labels.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 .../senkin_tmp/senkin_tmp/config.vsh.yaml     |  17 +-
 .../senkin_tmp_predict/config.vsh.yaml        |   4 +-
 .../senkin_tmp/senkin_tmp_predict/script.py   |  12 +-
 .../senkin_tmp_train/config.vsh.yaml          |   2 +-
 .../senkin_tmp/senkin_tmp_train/script.py     | 274 +++++++++---------
 5 files changed, 152 insertions(+), 157 deletions(-)

diff --git a/src/methods/senkin_tmp/senkin_tmp/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp/config.vsh.yaml
index 8f8885d9..9375a36d 100644
--- a/src/methods/senkin_tmp/senkin_tmp/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp/config.vsh.yaml
@@ -1,14 +1,17 @@
 __merge__: ../../../api/comp_method.yaml
 name: senkin_tmp
 label: "senkin & tmp"
-summary: "LightGBM + bidirectional GRU ensemble for CITE-seq protein prediction (OpenProblems 2022 2nd place)"
+summary: "LightGBM + bidirectional GRU ensemble for CITE-seq protein prediction (best CITE-seq model of the Open Problems 2022 competition)"
 description: |
-  Two-stage method from the OpenProblems NeurIPS 2021 competition. Stage 1 trains four
-  LightGBM models on different RNA feature representations (log-normalized, CLR-TSVD,
-  custom sqrt-normalized, and raw counts). Stage 2 refines predictions with two neural
-  network architectures: a bidirectional GRU with cosine-similarity loss and a dense
-  bidirectional GRU with MSE loss. Final predictions are a weighted blend (55% cosine,
-  45% MSE) of per-fold averaged outputs.
+  Two-stage method of senkin13 from the Open Problems NeurIPS 2022 multimodal single-cell
+  integration competition (2nd place overall, best CITE-seq submission). Stage 1 trains four
+  LightGBM models on different RNA feature representations (log-normalized counts; CLR-TSVD
+  together with correlated genes and TSVD/PCA of a custom sqrt normalization with per-batch
+  median correction; raw counts; raw counts predicting raw protein counts). Stage 2 refines
+  the TSVD-reduced LightGBM predictions, together with the feature blocks (each z-scored per
+  cell), with two neural networks: a bidirectional GRU with cosine-similarity loss and a dense
+  bidirectional GRU with MSE loss. Final predictions are a weighted blend (55% cosine, 45% MSE)
+  of per-fold averaged outputs, rescaled to the target scale with a global affine transform.
 references:
   doi:
     - 10.64898/2026.02.24.707614
diff --git a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
index e24c556c..7d85117c 100644
--- a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
@@ -18,7 +18,7 @@ engines:
     image: openproblems/base_pytorch_nvidia:1
     setup:
       - type: docker
-        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git
+        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@fix/match-original-pipeline
       - type: python
         packages:
           - lightgbm>=4.0
@@ -31,4 +31,4 @@ runners:
   - type: executable
   - type: nextflow
     directives:
-      label: [highmem, hightime, midcpu]
+      label: [lowmem, lowtime, lowcpu]
diff --git a/src/methods/senkin_tmp/senkin_tmp_predict/script.py b/src/methods/senkin_tmp/senkin_tmp_predict/script.py
index 8529f85d..fcec42ff 100644
--- a/src/methods/senkin_tmp/senkin_tmp_predict/script.py
+++ b/src/methods/senkin_tmp/senkin_tmp_predict/script.py
@@ -10,8 +10,8 @@
 
 ## VIASH START
 par = {
-    "input_test_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/test_mod1.h5ad",
-    "input_train_mod2": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/train_mod2.h5ad",
+    "input_test_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/test_mod1.h5ad",
+    "input_train_mod2": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/train_mod2.h5ad",
     "input_model": "output_model.pkl",
     "output": "output_pred.h5ad",
 }
@@ -26,9 +26,15 @@
 with open(par["input_model"], "rb") as f:
     bundle = pickle.load(f)
 
+predictions = bundle["test_predictions"]
+if "test_obs_names" in bundle:
+    # The train step predicted the test cells in the order of input_test_mod1; make sure nothing changed
+    assert (np.asarray(bundle["test_obs_names"]) == adata_rna_test.obs_names.values).all(), "test cells do not match the trained model"
+assert predictions.shape == (adata_rna_test.n_obs, adata_prot_train.n_vars)
+
 logger.info("Writing predictions...")
 adata_out = ad.AnnData(
-    layers={"normalized": csc_matrix(bundle["test_predictions"])},
+    layers={"normalized": csc_matrix(predictions)},
     obs=adata_rna_test.obs,
     var=adata_prot_train.var,
     uns={
diff --git a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
index 46642bb3..5c370020 100644
--- a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
@@ -66,7 +66,7 @@ engines:
         github:
           - openproblems-bio/core#subdirectory=packages/python/openproblems
       - type: docker
-        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git
+        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@fix/match-original-pipeline
 runners:
   - type: executable
   - type: nextflow
diff --git a/src/methods/senkin_tmp/senkin_tmp_train/script.py b/src/methods/senkin_tmp/senkin_tmp_train/script.py
index ef3aec54..0c6c1eaf 100644
--- a/src/methods/senkin_tmp/senkin_tmp_train/script.py
+++ b/src/methods/senkin_tmp/senkin_tmp_train/script.py
@@ -5,24 +5,28 @@
 
 import anndata as ad
 import numpy as np
-from scipy.sparse import issparse
-from sklearn.decomposition import PCA
+from sklearn.decomposition import PCA, TruncatedSVD
 from sklearn.model_selection import KFold
-import tensorflow as tf
 
-from senkin_tmp_cite_pred.preprocess import remove_constant_vars, senkin_normalize, get_top_correlated_features, log_normalize, clr_tsvd
+from senkin_tmp_cite_pred.preprocess import (
+    clr_tsvd,
+    get_top_correlated_features,
+    log_normalize,
+    remove_constant_vars,
+    senkin_normalize,
+    to_dense,
+)
 from senkin_tmp_cite_pred.lgbm_models import get_lgbm_predictions, lgbm_params_1, lgbm_params_2, lgbm_params_3, lgbm_params_4
-from senkin_tmp_cite_pred.nn_models import cite_cos_sim_model, cite_mse_model, nn_kfold, zscore
-from senkin_tmp_cite_pred.metrics import cosine_similarity_loss
+from senkin_tmp_cite_pred.nn_models import cite_cos_sim_model, cite_mse_model, nn_kfold, prepare_nn_inputs, zscore
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 ## VIASH START
 par = {
-    "input_train_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/train_mod1.h5ad",
-    "input_train_mod2": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/train_mod2.h5ad",
-    "input_test_mod1":  "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/test_mod1.h5ad",
+    "input_train_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/train_mod1.h5ad",
+    "input_train_mod2": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/train_mod2.h5ad",
+    "input_test_mod1":  "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/normal/test_mod1.h5ad",
     "output": "output_model.pkl",
     "n_folds": 5,
     "lgbm_boost_rounds": 10000,
@@ -30,39 +34,25 @@
     "nn_epochs": 100,
     "n_tsvd_components": 100,
 }
-meta = {"name": "senkin_tmp", "resources_dir": "src/methods/senkin_tmp/senkin_tmp_train"}
+meta = {"name": "senkin_tmp", "resources_dir": "src/methods/senkin_tmp/senkin_tmp_train", "cpus": None}
 ## VIASH END
 
 sys.path.append(meta["resources_dir"])
 from exit_codes import exit_non_applicable
 
-
-def _to_dense(X):
-    return X.toarray() if issparse(X) else np.array(X)
-
-
-def _parse_batch(obs, batch_col="batch"):
-    if batch_col not in obs.columns:
-        obs["day"] = "unknown"
-        obs["donor"] = "unknown"
-        return obs
-    def _split(b):
-        b = str(b)
-        d_idx = b.find("d")
-        if d_idx > 0:
-            return b[1:d_idx], b[d_idx + 1:]
-        return b, b
-    obs["day"], obs["donor"] = zip(*obs[batch_col].astype(str).map(_split))
-    return obs
-
+# The original solution corrected batch effects per day and computed gene-protein correlations per donor and day.
+# The benchmark datasets only carry a generic `batch` column (e.g. "s1d1" in NeurIPS 2021, "day_donor" in
+# NeurIPS 2022), which is used for both purposes here. Nothing dataset specific is assumed about its format.
+BATCH_KEY = "batch"
+SEED = 42
 
 # ---------------------------------------------------------------------------
 # Load data
 # ---------------------------------------------------------------------------
 logger.info("Reading input files...")
-adata_rna_train  = ad.read_h5ad(par["input_train_mod1"])
+adata_rna_train = ad.read_h5ad(par["input_train_mod1"])
 adata_prot_train = ad.read_h5ad(par["input_train_mod2"])
-adata_rna_test   = ad.read_h5ad(par["input_test_mod1"])
+adata_rna_test = ad.read_h5ad(par["input_test_mod1"])
 
 # senkin is a CITE-seq method that predicts protein (ADT) from RNA (GEX); it treats
 # mod1 as RNA and mod2 as protein. Skip the datasets/directions it cannot handle
@@ -75,78 +65,92 @@ def _split(b):
         f"got mod1={_mod1!r}, mod2={_mod2!r}."
     )
 
-adata_rna_train.obs  = _parse_batch(adata_rna_train.obs)
-adata_prot_train.obs = _parse_batch(adata_prot_train.obs)
-adata_rna_test.obs   = _parse_batch(adata_rna_test.obs)
+# Align protein cells with RNA cells
+adata_prot_train = adata_prot_train[adata_rna_train.obs_names].copy()
 
-# Mark split membership before concatenating
-adata_rna_train.obs["split"] = "train"
-adata_rna_test.obs["split"]  = "test"
+for adata in (adata_rna_train, adata_rna_test):
+    if BATCH_KEY not in adata.obs.columns:
+        adata.obs[BATCH_KEY] = "all"
+    adata.obs[BATCH_KEY] = adata.obs[BATCH_KEY].astype(str)
+adata_prot_train.obs[BATCH_KEY] = adata_rna_train.obs[BATCH_KEY].values
 
 # Concatenate train + test RNA so the original pipeline sees both at once
-import scanpy as sc
-adata_rna_all = sc.concat([adata_rna_train, adata_rna_test], axis=0)
+# (all unsupervised transformations were fit on train and test cells together)
+adata_rna_train.obs["split"] = "train"
+adata_rna_test.obs["split"] = "test"
+adata_rna_all = ad.concat([adata_rna_train, adata_rna_test], axis=0, join="inner", merge="same")
 adata_rna_all.X = adata_rna_all.layers["counts"]
+adata_rna_all.obs[BATCH_KEY] = adata_rna_all.obs[BATCH_KEY].astype(str)
+del adata_rna_all.layers["counts"]
+if "normalized" in adata_rna_all.layers:
+    del adata_rna_all.layers["normalized"]
 
 # ---------------------------------------------------------------------------
 # Preprocessing on combined train+test RNA
 # ---------------------------------------------------------------------------
 logger.info("Preprocessing RNA (train + test combined)...")
 
-adata_rna_all_filt = remove_constant_vars(adata_rna_all)
-X_counts_all = _to_dense(adata_rna_all_filt.layers.get("counts", adata_rna_all_filt.X))
-
-train_mask = adata_rna_all_filt.obs["split"] == "train"
-test_mask  = adata_rna_all_filt.obs["split"] == "test"
+adata_rna_all = remove_constant_vars(adata_rna_all)
+train_idx = np.flatnonzero((adata_rna_all.obs["split"] == "train").values)
+test_idx = np.flatnonzero((adata_rna_all.obs["split"] == "test").values)
+n_train, n_test = len(train_idx), len(test_idx)
+logger.info(f"{n_train} train cells, {n_test} test cells, {adata_rna_all.n_vars} non-constant genes")
 
-# Log-normalize (CP10K + log1p) via the senkin_tmp_cite_pred library helper.
-# log_normalize preserves the input sparsity, so densify for the array math below.
-X_lognorm_all = _to_dense(log_normalize(adata_rna_all_filt, target_sum=1e4)).astype(np.float64)
+# Log-normalize the way the competition inputs were normalized: log1p(counts per million)
+X_lognorm_all = log_normalize(adata_rna_all, target_sum=1e6).tocsr()
 
-# CLR-TSVD via the library helper (random_state=42 for reproducible components)
 logger.info("Computing CLR-TSVD...")
-X_clr_tsvd_all = clr_tsvd(adata_rna_all_filt, n_components=200, random_state=42)
+X_clr_tsvd_all = clr_tsvd(adata_rna_all, n_components=200, random_state=SEED)
+
+# Custom sqrt normalization with per-batch median correction, reduced with TSVD (100) and PCA (64) as in the original
+logger.info("Computing SenKin normalization, TSVD and PCA...")
+X_sqrt_norm_all = senkin_normalize(adata_rna_all, batch_key=BATCH_KEY)
+n_sqrt_tsvd = min(100, min(X_sqrt_norm_all.shape) - 1)
+X_sqrt_tsvd_all = TruncatedSVD(n_components=n_sqrt_tsvd, algorithm="arpack", random_state=SEED).fit_transform(X_sqrt_norm_all)
+n_sqrt_pca = min(64, min(X_sqrt_norm_all.shape) - 1)
+X_sqrt_pca_all = PCA(n_components=n_sqrt_pca, copy=False, random_state=SEED).fit_transform(X_sqrt_norm_all)
+del X_sqrt_norm_all
+gc.collect()
+
+# Correlated gene selection: computed on train cells only (no label leakage), on log-normalized RNA and
+# normalized proteins, per batch. As in the original, the raw counts of the selected genes are used as features.
+logger.info("Selecting correlated features...")
+adata_rna_train_filt = adata_rna_all[train_idx].copy()
+adata_rna_train_filt.obsm["X_log_normalized"] = X_lognorm_all[train_idx]
+top_corr_genes = get_top_correlated_features(
+    adata_rna_train_filt,
+    adata_prot_train,
+    group_key=BATCH_KEY,
+    quantile_threshold=0.1,
+    top_n=10,
+    rna_key="X_log_normalized",
+    prot_key="normalized",
+)
+del adata_rna_train_filt
 
-# SenKin normalization + PCA — fit on all
-logger.info("Computing SenKin normalization and PCA...")
-X_sqrt_norm_all = np.asarray(senkin_normalize(adata_rna_all_filt, batch_key="day"))
-n_pca = min(100, min(X_sqrt_norm_all.shape) - 1)
-pca_model = PCA(n_components=n_pca, random_state=42)
-X_pca_all = pca_model.fit_transform(X_sqrt_norm_all)
+# The original solution additionally used a hand-curated list of genes encoding the measured proteins.
+# Generic equivalent: genes whose name matches a protein name.
+def _feature_names(adata):
+    names = adata.var["feature_name"] if "feature_name" in adata.var.columns else adata.var_names.to_series()
+    return names.astype(str).str.upper()
 
-# Correlated gene selection — computed on train cells only (no label leakage)
-logger.info("Selecting correlated features...")
-adata_rna_train_filt = adata_rna_all_filt[train_mask]
-Y_prot_train = _to_dense(adata_prot_train.layers.get("normalized", adata_prot_train.X)).astype(np.float64)
-# get_top_correlated_features hardcodes .layers["dsb"]; alias our normalized layer
-_added_dsb = "dsb" not in adata_prot_train.layers
-if _added_dsb:
-    adata_prot_train.layers["dsb"] = adata_prot_train.layers.get("normalized", adata_prot_train.X)
-# Use "day" as group key — benchmark data has no "donor" column
-_group_key = "donor" if "donor" in adata_rna_train_filt.obs.columns else "day"
-top_corr_genes = get_top_correlated_features(adata_rna_train_filt, adata_prot_train, group_key=_group_key)
-if _added_dsb:
-    del adata_prot_train.layers["dsb"]
-all_var_names = list(adata_rna_all_filt.var_names)
-selected_gene_idxs = [all_var_names.index(g) for g in top_corr_genes if g in all_var_names]
-X_raw_selected_all = X_counts_all[:, selected_gene_idxs].astype(np.float64)
-
-# Split back into train / test portions
-X_lognorm_train    = X_lognorm_all[train_mask]
-X_lognorm_test     = X_lognorm_all[test_mask]
-X_clr_tsvd_train   = X_clr_tsvd_all[train_mask]
-X_clr_tsvd_test    = X_clr_tsvd_all[test_mask]
-X_pca_train        = X_pca_all[train_mask]
-X_pca_test         = X_pca_all[test_mask]
-X_raw_sel_train    = X_raw_selected_all[train_mask]
-X_raw_sel_test     = X_raw_selected_all[test_mask]
-X_counts_train     = X_counts_all[train_mask]
-X_counts_test      = X_counts_all[test_mask]
+_gene_names = _feature_names(adata_rna_all)
+_protein_names = set(_feature_names(adata_prot_train))
+known_genes = adata_rna_all.var_names[_gene_names.isin(_protein_names).values].tolist()
+selected_genes = sorted(set(top_corr_genes) | set(known_genes))
+logger.info(f"{len(top_corr_genes)} correlated genes + {len(known_genes)} protein-encoding genes = {len(selected_genes)} selected genes")
+X_raw_selected_all = to_dense(adata_rna_all[:, selected_genes].X, dtype=np.float32)
+
+X_counts_all = adata_rna_all.X.tocsr()
+
+# Protein targets
+Y_prot_train = to_dense(adata_prot_train.layers["normalized"], dtype=np.float64)
+Y_prot_raw = to_dense(adata_prot_train.layers.get("counts", adata_prot_train.X), dtype=np.float64)
 
 folds = KFold(n_splits=par["n_folds"], shuffle=True, random_state=666)
-n_tsvd       = par["n_tsvd_components"]
+n_tsvd = par["n_tsvd_components"]
 boost_rounds = par["lgbm_boost_rounds"]
-early_stop   = par["lgbm_early_stopping"]
+early_stop = par["lgbm_early_stopping"]
 
 # Pin LightGBM to the allocated cores. Its default (num_threads=0) spawns one thread
 # per core the container *sees* (the whole node) while the job is cgroup-throttled to
@@ -157,103 +161,85 @@ def _split(b):
     for _p in (lgbm_params_1, lgbm_params_2, lgbm_params_3, lgbm_params_4):
         _p["num_threads"] = _n_threads
 
-# Protein targets
-Y_prot_raw = _to_dense(adata_prot_train.layers.get("counts", adata_prot_train.X)).astype(np.float64)
-
 # ---------------------------------------------------------------------------
 # LightGBM — 4 models, train+test passed together (original design)
 # get_lgbm_predictions concatenates train+test, fits TSVD on combined array
 # ---------------------------------------------------------------------------
-logger.info("Training LightGBM model 1 (log-norm → proteins)...")
-lgbm1_svd_all = get_lgbm_predictions(
-    X_lognorm_train, Y_prot_train, X_lognorm_test,
-    folds, lgbm_params_1,
-    n_tsvd_components=n_tsvd,
-    num_boost_round=boost_rounds,
-    early_stopping_rounds=early_stop,
-)
-
-logger.info("Training LightGBM model 2 (combined → proteins)...")
-X_comb_train = np.concatenate([X_clr_tsvd_train, X_raw_sel_train, X_pca_train], axis=1)
-X_comb_test  = np.concatenate([X_clr_tsvd_test,  X_raw_sel_test,  X_pca_test],  axis=1)
-lgbm2_svd_all = get_lgbm_predictions(
-    X_comb_train, Y_prot_train, X_comb_test,
-    folds, lgbm_params_2,
-    n_tsvd_components=n_tsvd,
-    num_boost_round=boost_rounds,
-    early_stopping_rounds=early_stop,
-)
+def _lgbm(X_all, Y, params, description):
+    logger.info(f"Training LightGBM {description}...")
+    return get_lgbm_predictions(
+        X_all[train_idx], Y, X_all[test_idx], folds, params,
+        n_tsvd_components=n_tsvd, num_boost_round=boost_rounds, early_stopping_rounds=early_stop,
+    )
 
-logger.info("Training LightGBM model 3 (raw counts → proteins)...")
-lgbm3_svd_all = get_lgbm_predictions(
-    X_counts_train, Y_prot_train, X_counts_test,
-    folds, lgbm_params_3,
-    n_tsvd_components=n_tsvd,
-    num_boost_round=boost_rounds,
-    early_stopping_rounds=early_stop,
-)
+lgbm1_svd_all = _lgbm(X_lognorm_all, Y_prot_train, lgbm_params_1, "model 1 (log-normalized RNA -> proteins)")
 
-logger.info("Training LightGBM model 4 (raw counts → raw proteins)...")
-lgbm4_svd_all = get_lgbm_predictions(
-    X_counts_train, Y_prot_raw, X_counts_test,
-    folds, lgbm_params_4,
-    n_tsvd_components=n_tsvd,
-    num_boost_round=boost_rounds,
-    early_stopping_rounds=early_stop,
-)
+X_comb_all = np.concatenate([X_clr_tsvd_all, X_raw_selected_all, X_sqrt_tsvd_all, X_sqrt_pca_all], axis=1)
+lgbm2_svd_all = _lgbm(X_comb_all, Y_prot_train, lgbm_params_2, "model 2 (CLR-TSVD + selected genes + normalized TSVD/PCA -> proteins)")
+del X_comb_all
 
-# get_lgbm_predictions returns shape (n_train + n_test, n_tsvd)
-n_train = X_lognorm_train.shape[0]
-n_test  = X_lognorm_test.shape[0]
+lgbm3_svd_all = _lgbm(X_counts_all, Y_prot_train, lgbm_params_3, "model 3 (raw counts -> proteins)")
+lgbm4_svd_all = _lgbm(X_counts_all, Y_prot_raw, lgbm_params_4, "model 4 (raw counts -> raw proteins)")
+del X_counts_all, X_lognorm_all
+gc.collect()
 
-lgbm1_tr = lgbm1_svd_all[:n_train]; lgbm1_te = lgbm1_svd_all[n_train:]
-lgbm2_tr = lgbm2_svd_all[:n_train]; lgbm2_te = lgbm2_svd_all[n_train:]
-lgbm3_tr = lgbm3_svd_all[:n_train]; lgbm3_te = lgbm3_svd_all[n_train:]
-lgbm4_tr = lgbm4_svd_all[:n_train]; lgbm4_te = lgbm4_svd_all[n_train:]
+# ---------------------------------------------------------------------------
+# Neural networks. Every feature block is z-scored per cell before concatenation, as in the original.
+# ---------------------------------------------------------------------------
+def _nn_inputs(idx):
+    return prepare_nn_inputs(
+        X_clr_tsvd_all[idx], X_raw_selected_all[idx], X_sqrt_tsvd_all[idx], X_sqrt_pca_all[idx],
+        lgbm1_svd_all[idx], lgbm2_svd_all[idx], lgbm3_svd_all[idx], lgbm4_svd_all[idx],
+    )
 
-# Build NN inputs
-nn_X_train = np.concatenate([X_clr_tsvd_train, X_pca_train, X_raw_sel_train, lgbm1_tr, lgbm2_tr, lgbm3_tr, lgbm4_tr], axis=1).astype(np.float32)
-nn_X_test  = np.concatenate([X_clr_tsvd_test,  X_pca_test,  X_raw_sel_test,  lgbm1_te, lgbm2_te, lgbm3_te, lgbm4_te], axis=1).astype(np.float32)
+# get_lgbm_predictions returns train cells first, then test cells
+nn_X_train = _nn_inputs(np.arange(n_train))
+nn_X_test = _nn_inputs(np.arange(n_train, n_train + n_test))
 nn_y_train = Y_prot_train.astype(np.float32)
 
-train_cell_ids = np.array(adata_rna_all_filt.obs_names[train_mask])
-test_cell_ids  = np.array(adata_rna_all_filt.obs_names[test_mask])
+train_cell_ids = np.array(adata_rna_all.obs_names[train_idx])
+test_cell_ids = np.array(adata_rna_all.obs_names[test_idx])
 
-# ---------------------------------------------------------------------------
-# Neural network — use original nn_kfold which saves checkpoint weights
-# ---------------------------------------------------------------------------
 logger.info("Training neural network (cosine model)...")
-import os
-os.makedirs("models", exist_ok=True)
-
 train_preds_cos, test_preds_cos = nn_kfold(
     train_cell_ids, nn_X_train, nn_y_train,
-    test_cell_ids,  nn_X_test,
+    test_cell_ids, nn_X_test,
     cite_cos_sim_model, folds,
     model_name="cite_cos_model",
     BATCH_SIZE=620, EPOCHS=par["nn_epochs"], LR_FACTOR=0.05,
+    models_dir="models",
 )
 
 logger.info("Training neural network (MSE model)...")
 nn_y_train_z = zscore(nn_y_train)
 train_preds_mse, test_preds_mse = nn_kfold(
     train_cell_ids, nn_X_train, nn_y_train_z,
-    test_cell_ids,  nn_X_test,
+    test_cell_ids, nn_X_test,
     cite_mse_model, folds,
     model_name="cite_mse_model",
     BATCH_SIZE=600, EPOCHS=par["nn_epochs"], LR_FACTOR=0.1,
+    models_dir="models",
 )
 
 # Blend — identical to original train_nn_models
+train_preds = zscore(train_preds_cos) * 0.55 + zscore(train_preds_mse) * 0.45
 test_preds = zscore(test_preds_cos) * 0.55 + zscore(test_preds_mse) * 0.45
 
+# The original solution was scored with a per-cell Pearson correlation only, so its predictions are z-scored per
+# cell. The benchmark also computes RMSE/MAE, so bring the predictions back to the scale of the normalized proteins
+# with a single global affine transform fitted on the out-of-fold training predictions. One (slope, intercept) pair
+# for all cells and proteins leaves every per-cell and per-protein correlation untouched.
+slope, intercept = np.polyfit(train_preds.ravel(), Y_prot_train.ravel(), deg=1)
+logger.info(f"Rescaling z-scored predictions to the target scale: slope {slope:.4f}, intercept {intercept:.4f}")
+test_preds = test_preds * slope + intercept
+
 # ---------------------------------------------------------------------------
 # Save bundle — test predictions stored directly, predict script just reads them
 # ---------------------------------------------------------------------------
 logger.info("Saving model bundle...")
 bundle = {
     "test_predictions": test_preds.astype(np.float32),  # (n_test, n_proteins)
-    "test_obs": adata_rna_test.obs,
+    "test_obs_names": test_cell_ids,
     "prot_var": adata_prot_train.var,
     "dataset_id": adata_rna_train.uns.get("dataset_id", ""),
 }

From aa87dbd5260f254d1d00600fd957bb229b54fbd5 Mon Sep 17 00:00:00 2001
From: Vladimir Shitov <vladimirshitov98@gmail.com>
Date: Tue, 8 Sep 2026 02:04:01 +0200
Subject: [PATCH 3/5] senkin_tmp: pin the library to the verified commit

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml | 2 +-
 src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
index 7d85117c..a03eaeb0 100644
--- a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
@@ -18,7 +18,7 @@ engines:
     image: openproblems/base_pytorch_nvidia:1
     setup:
       - type: docker
-        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@fix/match-original-pipeline
+        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@4372e3d635136bb2de62cf01ccb5e68e08ff1c7e
       - type: python
         packages:
           - lightgbm>=4.0
diff --git a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
index 5c370020..8acb94e2 100644
--- a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
@@ -66,7 +66,7 @@ engines:
         github:
           - openproblems-bio/core#subdirectory=packages/python/openproblems
       - type: docker
-        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@fix/match-original-pipeline
+        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@4372e3d635136bb2de62cf01ccb5e68e08ff1c7e
 runners:
   - type: executable
   - type: nextflow

From 25f851f1a8f5ff5e26e2c8c414d99c13fafee40b Mon Sep 17 00:00:00 2001
From: Vladimir Shitov <vladimirshitov98@gmail.com>
Date: Tue, 8 Sep 2026 10:45:25 +0200
Subject: [PATCH 4/5] senkin_tmp: pin the library to the commit merged with
 main (#4)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml | 2 +-
 src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
index a03eaeb0..d9725649 100644
--- a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
@@ -18,7 +18,7 @@ engines:
     image: openproblems/base_pytorch_nvidia:1
     setup:
       - type: docker
-        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@4372e3d635136bb2de62cf01ccb5e68e08ff1c7e
+        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@4249aa0c81753c1dd97ea4f1aa089a78817e0ee8
       - type: python
         packages:
           - lightgbm>=4.0
diff --git a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
index 8acb94e2..e0a2a4d4 100644
--- a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
@@ -66,7 +66,7 @@ engines:
         github:
           - openproblems-bio/core#subdirectory=packages/python/openproblems
       - type: docker
-        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@4372e3d635136bb2de62cf01ccb5e68e08ff1c7e
+        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@4249aa0c81753c1dd97ea4f1aa089a78817e0ee8
 runners:
   - type: executable
   - type: nextflow

From c8401beba464d47dc8f5b85bf78dcdff0885c6de Mon Sep 17 00:00:00 2001
From: Vladimir Shitov <vladimirshitov98@gmail.com>
Date: Tue, 8 Sep 2026 21:26:34 +0200
Subject: [PATCH 5/5] senkin_tmp: pin the library to a830343 (robust
 correlation_score)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml | 2 +-
 src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml   | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
index d9725649..4c183390 100644
--- a/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml
@@ -18,7 +18,7 @@ engines:
     image: openproblems/base_pytorch_nvidia:1
     setup:
       - type: docker
-        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@4249aa0c81753c1dd97ea4f1aa089a78817e0ee8
+        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@a830343a148c087dd95f60d5c7e14529aaba2425
       - type: python
         packages:
           - lightgbm>=4.0
diff --git a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
index e0a2a4d4..4782679b 100644
--- a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
+++ b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml
@@ -66,7 +66,7 @@ engines:
         github:
           - openproblems-bio/core#subdirectory=packages/python/openproblems
       - type: docker
-        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@4249aa0c81753c1dd97ea4f1aa089a78817e0ee8
+        run: pip install --no-cache-dir --no-deps git+https://github.com/lueckenlab/senkin-tmp-cite-pred.git@a830343a148c087dd95f60d5c7e14529aaba2425
 runners:
   - type: executable
   - type: nextflow