Fix CSR Spearman scores and pbmc_multiome failures, add a seed to all methods

CHANGELOG.md

@@ -1,3 +1,17 @@
+# task_predict_modality 0.2.1
+
+## MAJOR CHANGES
+
+* Add `--seed` to all methods, controls and `run_benchmark`, so that reruns reproduce the predictions (PR #83).
+
+## BUG FIXES
+
+* `correlation`: Convert matrices to CSC, since dynutils < 1.0.13 gave wrong Spearman scores on CSR (PR #83).
+
+* `guanlab_dengkw_pm`: Replace `scipy.linalg.cho_factor()`, which failed on `pbmc_multiome`, with `numpy.linalg.solve()`, and ask for more memory and time (PR #83).
+
+* `simple_mlp`: Densify one batch at a time, which fixes running out of memory on `pbmc_multiome` (PR #83).
+
 # task_predict_modality 0.2.0
 
 Adds eight methods, selects the ATAC target peaks by `hvg_score`, and fixes the bugs surfaced by the first full benchmark runs.

README.md

@@ -415,6 +415,7 @@ Arguments:
 | `--input_train_mod2` | `file` | The mod2 expression values of the train cells. |
 | `--input_test_mod1` | `file` | The mod1 expression values of the test cells. |
 | `--output` | `file` | (*Output*) A prediction of the mod2 expression values of the test cells. |
+| `--seed` | `integer` | (*Optional*) Seed for the random number generators of the method. Default: `0`. |
 
 </div>
 
@@ -432,6 +433,7 @@ Arguments:
 | `--input_train_mod2` | `file` | The mod2 expression values of the train cells. |
 | `--input_test_mod1` | `file` | (*Optional*) The mod1 expression values of the test cells. |
 | `--output` | `file` | (*Output*) A pretrained model for predicting the expression of one modality from another. |
+| `--seed` | `integer` | (*Optional*) Seed for the random number generators used for training the model. Default: `0`. |
 
 </div>
 
@@ -450,6 +452,7 @@ Arguments:
 | `--input_test_mod1` | `file` | The mod1 expression values of the test cells. |
 | `--input_test_mod2` | `file` | The ground-truth mod2 expression values of the test cells. |
 | `--output` | `file` | (*Output*) A prediction of the mod2 expression values of the test cells. |
+| `--seed` | `integer` | (*Optional*) Seed for the random number generators of the control method, if it uses any. Default: `0`. |
 
 </div>
 

src/api/comp_control_method.yaml

@@ -32,6 +32,10 @@ arguments:
     __merge__: file_prediction.yaml
     direction: output
     required: true
+  - name: "--seed"
+    type: integer
+    default: 0
+    description: Seed for the random number generators of the control method, if it uses any.
 test_resources:
   - type: python_script
     path: /common/component_tests/check_config.py

src/api/comp_method.yaml

@@ -23,6 +23,10 @@ arguments:
     __merge__: file_prediction.yaml
     direction: output
     required: true
+  - name: "--seed"
+    type: integer
+    default: 0
+    description: Seed for the random number generators of the method.
 test_resources:
   - type: python_script
     path: /common/component_tests/check_config.py

src/api/comp_method_train.yaml

@@ -23,6 +23,10 @@ arguments:
     __merge__: file_pretrained_model.yaml
     direction: output
     required: true
+  - name: "--seed"
+    type: integer
+    default: 0
+    description: Seed for the random number generators used for training the model.
 test_resources:
   - type: python_script
     path: /common/component_tests/run_and_check_output.py

src/api/wf_method.yaml

@@ -25,6 +25,10 @@ arguments:
     __merge__: file_prediction.yaml
     direction: output
     required: true
+  - name: "--seed"
+    type: integer
+    default: 0
+    description: Seed for the random number generators of the method.
 test_resources:
   - type: python_script
     path: /common/component_tests/check_config.py

src/control_methods/random_predict/script.R

@@ -7,11 +7,14 @@ par <- list(
   input_train_mod1 = "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/train_mod1.h5ad",
   input_test_mod1 = "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/test_mod1.h5ad",
   input_train_mod2 = "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/train_mod2.h5ad",
-  output = "output.h5ad"
+  output = "output.h5ad",
+  seed = 0L
 )
 meta <- list(name = "foo")
 ## VIASH END
 
+set.seed(par$seed)
+
 cat("Reading h5ad files\n")
 input_train_mod2 <- anndata::read_h5ad(par$input_train_mod2)
 input_test_mod1 <- anndata::read_h5ad(par$input_test_mod1)

src/methods/babel/babel/main.nf

@@ -3,7 +3,7 @@ workflow run_wf {
   main:
   output_ch = input_ch
     | babel_train.run(
-      fromState: ["input_train_mod1", "input_train_mod2", "input_test_mod1"],
+      fromState: ["input_train_mod1", "input_train_mod2", "input_test_mod1", "seed"],
       toState: ["input_model": "output"]
     )
     | babel_predict.run(

src/methods/babel/babel_train/config.vsh.yaml

@@ -17,6 +17,7 @@ resources:
   - path: ../chrom_utils.py
     dest: chrom_utils.py
   - path: /src/utils/exit_codes.py
+  - path: /src/utils/seed.py
 test_resources:
   - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap
     dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap

src/methods/babel/babel_train/script.py

@@ -26,6 +26,7 @@
     "batch_size": 64,
     "loss2_weight": 3.0,
     "early_stopping_patience": 20,
+    "seed": 0,
 }
 meta = {"name": "babel", "resources_dir": "src/methods/babel"}
 ## VIASH END
@@ -36,6 +37,7 @@
 from exit_codes import exit_non_applicable
 from model import AssymSplicedAutoEncoder
 from losses import QuadLoss
+from seed import seed_everything
 
 
 def _as_csr(X):
@@ -131,6 +133,8 @@ def get_loss(self, y_pred, y_true, X=None, training=False):
         return self.criterion_(preds11, preds12, preds21, preds22, target1, target2_bin)
 
 
+seed_everything(par["seed"])
+
 logger.info("Reading input files...")
 adata_mod1_train = ad.read_h5ad(par["input_train_mod1"])
 adata_mod2_train = ad.read_h5ad(par["input_train_mod2"])

src/methods/cellmapper_linear/config.vsh.yaml

@@ -60,6 +60,7 @@ arguments:
 resources:
   - type: python_script
     path: script.py
+  - path: /src/utils/seed.py
 engines:
   - type: docker
     image: openproblems/base_python:1

src/methods/cellmapper_linear/script.py

@@ -1,3 +1,5 @@
+import sys
+
 import anndata as ad
 import cellmapper as cm
 from scipy.sparse import csc_matrix
@@ -13,13 +15,20 @@
     'fallback_representation': 'joint_pca',  # or None for fallback_representation
     'n_neighbors': 30,
     'kernel_method': 'gauss',
-    'mask_var': "hvg"  # variable to mask for fallback representation
+    'mask_var': "hvg",  # variable to mask for fallback representation
+    'seed': 0,
 }
 meta = {
   'name': 'cellmapper_linear',
+  'resources_dir': 'src/utils',
 }
 ## VIASH END
 
+sys.path.append(meta['resources_dir'])
+from seed import seed_everything
+
+seed_everything(par['seed'])
+
 print('Reading input files', flush=True)
 input_train_mod1 = ad.read_h5ad(par['input_train_mod1'])
 input_train_mod2 = ad.read_h5ad(par['input_train_mod2'])

src/methods/cellmapper_scvi/config.vsh.yaml

@@ -67,6 +67,7 @@ resources:
     path: script.py
   - path: utils.py
     dest: utils.py 
+  - path: /src/utils/seed.py
 engines:
   - type: docker
     image: openproblems/base_pytorch_nvidia:1

src/methods/cellmapper_scvi/script.py

@@ -16,7 +16,7 @@
     'use_hvg': False,
     'adt_normalization': 'clr',  # Normalization method for ADT data
     'plot_umap': True,
-
+    'seed': 0,
 }
 meta = {
   'name': 'cellmapper_scvi',
@@ -25,8 +25,13 @@
 ## VIASH END
 
 sys.path.append(meta['resources_dir'])
+import scvi
+from seed import seed_everything
 from utils import get_representation
 
+seed_everything(par['seed'])
+scvi.settings.seed = par['seed']
+
 print('Reading input files', flush=True)
 input_train_mod1 = ad.read_h5ad(par['input_train_mod1'])
 input_train_mod2 = ad.read_h5ad(par['input_train_mod2'])

src/methods/guanlab_dengkw_pm/config.vsh.yaml

@@ -26,10 +26,6 @@ arguments:
       kernel ridge model, so the consensus averages `2 * n_repeats` predictions.
     info:
       test_default: 1
-  - name: "--seed"
-    type: "integer"
-    default: 1000
-    description: Seed for the batch reshuffling.
 resources:
   - type: python_script
     path: script.py
@@ -48,4 +44,4 @@ runners:
   - type: executable
   - type: nextflow
     directives:
-      label: [hightime, highmem, highcpu]
+      label: [veryhightime, veryhighmem, highcpu]

src/methods/guanlab_dengkw_pm/script.py

@@ -1,7 +1,6 @@
 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
@@ -14,7 +13,7 @@
     'input_test_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/normal/test_mod1.h5ad',
     'output': 'output.h5ad',
     'n_repeats': 5,
-    'seed': 1000
+    'seed': 0
 }
 meta = {
     'name': 'guanlab_dengkw_pm',
@@ -32,22 +31,23 @@
 # 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.
+# KernelRidge solves (K + alpha * I) w = y with scipy.linalg.solve(assume_a="pos"), which fails on large systems.
+# Tried on the halves of the batches of pbmc_multiome (56k-74k cells), or at the same size:
+# - scipy.linalg.solve(assume_a="pos"): raises a MemoryError or segfaults, already from ~30-40k cells
+# - scipy.linalg.cho_factor(): reports K + alpha * I as not positive definite ("potrf return info")
+# - numpy.linalg.cholesky(): segfaults
+# - numpy.linalg.solve(): works, 270 s for 66k cells on 30 threads. Where scipy's Cholesky works, the coefficients
+#   agree to 6e-13.
 #
-# 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.
+# The Cholesky failures depend on the CPU: on the Cascade Lake workers of de.NBI, OpenBLAS picks its AVX-512 kernel
+# (SkylakeX). cho_factor() succeeds at 66k cells with the Haswell kernel, on AMD CPUs or with
+# OPENBLAS_CORETYPE=Haswell, but numpy.linalg.solve() doesn't need that set on every backend.
 #
-# Switch back to KernelRidge once scipy.linalg.solve() handles matrices of this size again.
+# Switch back to KernelRidge once OpenBLAS's Cholesky handles these systems on every CPU.
 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)
+    dual_coef = np.linalg.solve(K, y)
     return kernel(X_test, X) @ dual_coef
 
 
@@ -92,7 +92,7 @@ def kernel_ridge_predict(kernel, alpha, X, y, X_test):
 print('Models using the Truncated SVD to reduce the dimension', flush=True)
 
 if n_mod1 is not None and n_mod1 < input_train.n_vars:
-    embedder_mod1 = TruncatedSVD(n_components=n_mod1)
+    embedder_mod1 = TruncatedSVD(n_components=n_mod1, random_state=par['seed'])
     mod1_pca = embedder_mod1.fit_transform(input_train.layers["normalized"]).astype(np.float32)
     train_matrix = mod1_pca[input_train.obs['group'] == 'train']
     test_matrix = mod1_pca[input_train.obs['group'] == 'test']
@@ -101,7 +101,7 @@ def kernel_ridge_predict(kernel, alpha, X, y, X_test):
     test_matrix = input_test_mod1.to_df(layer="normalized").values.astype(np.float32)
 
 if n_mod2 is not None and n_mod2 < input_train_mod2.n_vars:
-    embedder_mod2 = TruncatedSVD(n_components=n_mod2)
+    embedder_mod2 = TruncatedSVD(n_components=n_mod2, random_state=par['seed'])
     train_gs = embedder_mod2.fit_transform(input_train_mod2.layers["normalized"]).astype(np.float32)
 else:
     # no dimensionality reduction, so the KRR predicts the mod2 features directly

src/methods/knnr_py/script.py

@@ -12,6 +12,7 @@
     'output': 'output.h5ad',
     'n_pcs': 4,
     'n_neighbors': 5,
+    'seed': 0,
 }
 meta = { 'name': 'foo' }
 ## VIASH END
@@ -31,7 +32,7 @@
 )
 
 print('Performing dimensionality reduction on modality 1 values...', flush=True)
-embedder = TruncatedSVD(n_components=par['n_pcs'])
+embedder = TruncatedSVD(n_components=par['n_pcs'], random_state=par['seed'])
 X = embedder.fit_transform(input_train.layers["normalized"])
 
 # split dimred back up

src/methods/knnr_r/script.R

@@ -11,11 +11,14 @@ par <- list(
   output = "output.h5ad",
   n_pcs = 4L,
   n_neighbors = 3,
-  distance_method = "pearson"
+  distance_method = "pearson",
+  seed = 0L
 )
 meta <- list(name = "foo")
 ## VIASH END
 
+set.seed(par$seed)
+
 cat("Reading mod1 h5ad files\n")
 input_train_mod1 <- anndata::read_h5ad(par$input_train_mod1)
 dataset_id <- input_train_mod1$uns[["dataset_id"]]

src/methods/lm/script.R

@@ -9,11 +9,14 @@ par <- list(
   input_test_mod1 = paste0(path, "test_mod1.h5ad"),
   input_train_mod2 = paste0(path, "train_mod2.h5ad"),
   output = "output.h5ad",
-  n_pcs = 4L
+  n_pcs = 4L,
+  seed = 0L
 )
 meta <- list(name = "foo")
 ## VIASH END
 
+set.seed(par$seed)
+
 cat("Reading mod1 files\n")
 input_train_mod1 <- anndata::read_h5ad(par$input_train_mod1)
 input_test_mod1 <- anndata::read_h5ad(par$input_test_mod1)

src/methods/novel/novel/main.nf

@@ -3,7 +3,7 @@ workflow run_wf {
   main:
   output_ch = input_ch
     | novel_train.run(
-      fromState: ["input_train_mod1", "input_train_mod2"],
+      fromState: ["input_train_mod1", "input_train_mod2", "seed"],
       toState: ["input_model": "output"]
     )
     | novel_predict.run(

src/methods/novel/novel_train/config.vsh.yaml

@@ -7,14 +7,11 @@ arguments:
     description: Number of training epochs.
     info:
       test_default: 2
-  - name: "--seed"
-    type: "integer"
-    default: 1
-    description: "Seed for choosing the internal validation batches, when the dataset does not contain the batches used in the competition."
 resources:
   - path: script.py
     type: python_script
   - path: ../helper_functions.py
+  - path: /src/utils/seed.py
 engines:
   - type: docker
     image: openproblems/base_pytorch_nvidia:1

src/methods/novel/novel_train/script.py

@@ -25,7 +25,7 @@
   'input_train_mod2': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/normal/train_mod2.h5ad',
   'output': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/normal/models/novel',
   'n_epochs': 100,
-  'seed': 1
+  'seed': 0
 }
 meta = {
   'resources_dir': 'src/methods/novel',
@@ -35,6 +35,9 @@
 sys.path.append(meta['resources_dir'])
 from helper_functions import train_and_valid, lsiTransformer, ModalityMatchingDataset
 from helper_functions import ModelRegressionAtac2Gex, ModelRegressionAdt2Gex, ModelRegressionGex2Adt, ModelRegressionGex2Atac
+from seed import seed_everything
+
+seed_everything(par['seed'])
 
 print('Load data', flush=True)
 input_train_mod1 = ad.read_h5ad(par['input_train_mod1'])