Fix/senkin fit walltime

src/methods/senkin_tmp/senkin_tmp_predict/config.vsh.yaml

@@ -18,12 +18,13 @@ 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@be8bee65c146316e579ae092c43247284ba801fe
       - type: python
         packages:
           - lightgbm>=4.0
           - tensorflow>=2.12
           - scikit-learn>=1.1
+          - joblib>=1.4
           - mudata>=0.2
           - muon>=0.1
           - fast-array-utils

src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml

@@ -17,16 +17,27 @@ arguments:
   - name: "--n_folds"
     type: integer
     default: 5
-    description: Number of cross-validation folds for LightGBM and neural network training.
+    description: Number of cross-validation folds for the neural networks (one network per fold, predictions averaged).
+  - name: "--lgbm_n_folds"
+    type: integer
+    default: 3
+    description: |
+      Number of cross-validation folds for the LightGBM models, whose out-of-fold predictions are
+      features of the neural networks. Fewer folds cut the LightGBM time proportionally at a small
+      cost in the quality of these features.
   - name: "--lgbm_boost_rounds"
     type: integer
-    default: 10000
-    description: Maximum LightGBM boosting rounds (early stopping applies).
+    default: 100
+    description: |
+      Maximum LightGBM boosting rounds (early stopping applies). The original solution used up to
+      10000 rounds at learning rate 0.01, which takes days on the benchmark datasets. At learning
+      rate 0.1, 100 rounds keep the LightGBM per-target quality within 0.015 (NeurIPS 2021) / 0.003
+      (NeurIPS 2022) of 300 rounds, and the final predictions of the networks are unchanged.
     info:
       test_default: 50
   - name: "--lgbm_early_stopping"
     type: integer
-    default: 100
+    default: 20
     description: LightGBM early stopping patience (rounds without improvement).
     info:
       test_default: 10
@@ -40,6 +51,26 @@ arguments:
     type: integer
     default: 100
     description: TSVD components for reducing LightGBM predictions before NN input.
+  - name: "--lgbm_learning_rate"
+    type: double
+    default: 0.1
+    description: |
+      LightGBM learning rate. The original solution used 0.01 with up to 10000 rounds; 0.1 with a few
+      hundred rounds reaches the same LightGBM quality in a fraction of the time.
+  - name: "--lgbm_max_bin"
+    type: integer
+    default: 63
+    description: |
+      Number of histogram bins per feature in LightGBM (library default 255). 63 halves the memory
+      of the binned whole-transcriptome inputs held by every worker process (3 instead of 6.6 GB per
+      fold on the NeurIPS 2022 CITE data) at identical validation loss and slightly faster rounds.
+  - name: "--lgbm_n_jobs"
+    type: integer
+    default: -1
+    description: |
+      Number of worker processes training different protein targets in parallel; the allocated CPUs
+      are divided among them. -1 uses one process per allocated CPU (each with one thread), which is
+      several times faster than one target at a time with all threads.
 engines:
   - type: docker
     image: nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04
@@ -60,13 +91,14 @@ engines:
           - jsonschema
           - lightgbm>=4.0
           - scikit-learn>=1.1
+          - joblib>=1.4
           - mudata>=0.2
           - muon>=0.1
           - fast-array-utils
         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@be8bee65c146316e579ae092c43247284ba801fe
 runners:
   - type: executable
   - type: nextflow

src/methods/senkin_tmp/senkin_tmp_train/script.py

@@ -29,12 +29,16 @@
     "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,
-    "lgbm_early_stopping": 100,
+    "lgbm_n_folds": 3,
+    "lgbm_boost_rounds": 100,
+    "lgbm_early_stopping": 20,
     "nn_epochs": 100,
     "n_tsvd_components": 100,
+    "lgbm_learning_rate": 0.1,
+    "lgbm_max_bin": 63,
+    "lgbm_n_jobs": -1,
 }
-meta = {"name": "senkin_tmp", "resources_dir": "src/methods/senkin_tmp/senkin_tmp_train", "cpus": None}
+meta = {"name": "senkin_tmp", "resources_dir": "src/methods/senkin_tmp/senkin_tmp_train", "cpus": None, "memory_gb": None}
 ## VIASH END
 
 sys.path.append(meta["resources_dir"])
@@ -148,6 +152,7 @@ def _feature_names(adata):
 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)
+lgbm_folds = KFold(n_splits=par["lgbm_n_folds"], shuffle=True, random_state=666)
 n_tsvd = par["n_tsvd_components"]
 boost_rounds = par["lgbm_boost_rounds"]
 early_stop = par["lgbm_early_stopping"]
@@ -157,9 +162,14 @@ def _feature_names(adata):
 # 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):
+for _p in (lgbm_params_1, lgbm_params_2, lgbm_params_3, lgbm_params_4):
+    if _n_threads:
         _p["num_threads"] = _n_threads
+    _p["learning_rate"] = par["lgbm_learning_rate"]
+    _p["max_bin"] = par["lgbm_max_bin"]
+lgbm_n_jobs = par["lgbm_n_jobs"]
+# Let the library keep the LightGBM worker processes within the allocated memory (a tenth is left for the rest).
+lgbm_memory_budget_gb = meta["memory_gb"] * 0.9 if meta.get("memory_gb") else None
 
 # ---------------------------------------------------------------------------
 # LightGBM — 4 models, train+test passed together (original design)
@@ -168,19 +178,22 @@ def _feature_names(adata):
 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,
+        X_all[train_idx], Y, X_all[test_idx], lgbm_folds, params,
         n_tsvd_components=n_tsvd, num_boost_round=boost_rounds, early_stopping_rounds=early_stop,
+        n_jobs=lgbm_n_jobs, memory_budget_gb=lgbm_memory_budget_gb,
     )
 
 lgbm1_svd_all = _lgbm(X_lognorm_all, Y_prot_train, lgbm_params_1, "model 1 (log-normalized RNA -> proteins)")
+del X_lognorm_all  # not needed any more; keeps the parent's footprint (and hence the memory left for workers) small
+gc.collect()
 
 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
 
 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
+del X_counts_all
 gc.collect()
 
 # ---------------------------------------------------------------------------