Fix/ss opm faithful

CHANGELOG.md

@@ -100,6 +100,8 @@
 
 * `scbutterfly_train`, `scbutterfly_predict`: Upgrade pip before installing torch. The pip shipped in `python:3.9` rejects the `typing_extensions` wheel over the underscore in its metadata name, falls back to the sdist, and cannot build it because `flit_core` is not on the PyTorch index -- so the image stopped building the moment `typing_extensions` started requiring `flit_core>=3.11` (PR #64).
 
+* `ss_opm`: Rebuild the derived inputs the original solution was written against instead of feeding it the raw h5ad columns: standardized per-cell statistics, the day and donor parsed from the `{day}_{donor}` batch labels (the 2021 donor was taken as the day), batch singular vectors from per-batch gene medians, and the HGNC/Reactome-selected CITE input genes rather than all 14k-22k genes (a 96 GB peak and a different model). Drop training cells with a constant target vector, which made the correlation loss `NaN` from epoch 0 on the 2022 CITE datasets; keep the multiome targets in float32, whose dense float64 copies OOM-killed 2022 ATAC->GEX; and map the per-cell z-scored network output back to the target scale, so RMSE and Spearman are meaningful. The image build and the runtime fallback download the Reactome gene sets with a browser-like user agent, because reactome.org answers HTTP 403 to Python's default one (PR #69).
+
 # task_predict_modality 0.1.1
 
 ## NEW FUNCTIONALITY

src/methods/ss_opm/ss_opm/config.vsh.yaml

@@ -6,8 +6,12 @@ description: |
   Encoder-decoder MLP method using SVD-based dimensionality reduction for both inputs and
   targets, followed by batch-median correction. The encoder maps (optionally augmented)
   cell embeddings to a latent space; multiple decoder blocks predict target expression in
-  the SVD-compressed space. The method was the winning solution of the NeurIPS 2021
-  Open Problems Multimodal Single-Cell Integration Kaggle competition.
+  the SVD-compressed space. The method was the winning solution of the NeurIPS 2022
+  Open Problems Multimodal Single-Cell Integration Kaggle competition. The inputs the
+  original derived from the competition tables (standardized per-cell and per-batch
+  statistics, the CITE gene masks built from HGNC and Reactome) are rebuilt from the
+  task's files, and the per-cell z-scored output is mapped to the target scale with a
+  global affine transform.
 references:
   doi:
     - 10.1101/2022.04.11.487796

src/methods/ss_opm/ss_opm_predict/config.vsh.yaml

@@ -5,22 +5,6 @@ info:
   test_setup:
     with_model:
       input_model: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/models/ss_opm
-arguments:
-  - name: "--cell_type_col"
-    type: string
-    required: false
-    description: |
-      Column in `obs` holding cell type labels. ss_opm uses them for the `cell_ratio_*`
-      features. This task's file format does not carry cell types, so it is unset by
-      default and the ratios fall back to uniform; set it when plugging in a dataset
-      that does have them.
-  - name: "--day_pattern"
-    type: string
-    default: 'd(\d+)'
-    description: |
-      Regex whose first capture group is the collection day within a batch label. The
-      default matches the NeurIPS 2021 `s{site}d{day}` naming. Batches that do not match
-      get day 0, i.e. the model sees a single constant day rather than failing.
 resources:
   - type: python_script
     path: script.py
@@ -39,4 +23,4 @@ runners:
   - type: executable
   - type: nextflow
     directives:
-      label: [highmem, hightime, midcpu, highsharedmem, midgpu]
+      label: [midmem, midtime, lowcpu]

src/methods/ss_opm/ss_opm_predict/script.py

@@ -1,104 +1,94 @@
-import sys
 import os
-import gc
 import pickle
+import sys
+
+import anndata as ad
 import numpy as np
 import pandas as pd
 import scipy.sparse
-import anndata as ad
-from ss_opm.model.encoder_decoder.encoder_decoder import EncoderDecoder
-
 import torch
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
-print(f'Using device: {device}', flush=True)
 
 ## VIASH START
 par = {
-    'input_test_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/test_mod1.h5ad',
-    'input_model': 'output/models/ss_opm',
-    'output': 'output/prediction.h5ad',
-    'cell_type_col': None,
-    'day_pattern': r'd(\d+)',
-}
-meta = {
-    'name': 'ss_opm_predict',
-    'resources_dir': 'src/methods/ss_opm',
+    "input_test_mod1": "resources_test/task_predict_modality/openproblems_neurips2021/bmmc_cite/swap/test_mod1.h5ad",
+    "input_model": "output/models/ss_opm",
+    "output": "output/prediction.h5ad",
 }
+meta = {"name": "ss_opm_predict", "resources_dir": "src/methods/ss_opm"}
 ## VIASH END
 
-sys.path.append(meta['resources_dir'])
-from ss_opm_common import apply_runtime_patches, build_metadata, to_sparse_csr
+sys.path.append(meta["resources_dir"])
+from ss_opm_common import (  # noqa: E402
+    apply_runtime_patches,
+    apply_standardization,
+    build_metadata,
+    compute_cell_statistics,
+    load_model_bundle,
+    to_sparse_csr,
+)
 
 apply_runtime_patches()
 
-# ---- Load task info ----
-with open(os.path.join(par['input_model'], 'task_info.pickle'), 'rb') as f:
-    task_info = pickle.load(f)
-task_type = task_info['task_type']
-mod2 = task_info['mod2']
-dataset_id = task_info['dataset_id']
-print(f'Task type: {task_type}, mod2: {mod2}', flush=True)
+from ss_opm.model.encoder_decoder.encoder_decoder import EncoderDecoder  # noqa: E402
 
-# ---- Load test data ----
-print('Loading test data...', flush=True)
-input_test_mod1 = ad.read_h5ad(par['input_test_mod1'])
-test_inputs = to_sparse_csr(input_test_mod1.layers['normalized'])
-test_metadata = build_metadata(
-    input_test_mod1,
-    task_type,
-    cell_type_col=par['cell_type_col'],
-    group_by_batch=False,
-    day_pattern=par['day_pattern'],
-)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}", flush=True)
 
-# ---- Load model and preprocessing artifacts ----
-print('Loading model...', flush=True)
-with open(os.path.join(par['input_model'], 'pre_post_process.pickle'), 'rb') as f:
-    pre_post_process = pickle.load(f)
+# ---- Load the training bundle ----
+task_info, batch_singular_vectors = load_model_bundle(par["input_model"])
+task_type = task_info["task_type"]
+print(f"Task type: {task_type}, {task_info['mod1']} -> {task_info['mod2']}", flush=True)
 
-model = EncoderDecoder(params=None)
-# PyTorch >=2.6 defaults weights_only=True, which blocks custom classes.
-# Patch torch.load to use weights_only=False for trusted local model files.
-import torch as _torch
-_orig_torch_load = _torch.load
-_torch.load = lambda *a, **kw: _orig_torch_load(*a, **{**kw, 'weights_only': False})
-model.load(os.path.join(par['input_model'], 'model'))
-_torch.load = _orig_torch_load
-model.params['device'] = device
+with open(os.path.join(par["input_model"], "pre_post_process.pickle"), "rb") as handle:
+    pre_post_process = pickle.load(handle)
+mod2_var = pd.read_parquet(os.path.join(par["input_model"], "mod2_var.parquet"))
 
-mod2_var = pd.read_parquet(os.path.join(par['input_model'], 'mod2_var.parquet'))
+model = EncoderDecoder(params=None)
+# PyTorch >= 2.6 defaults to weights_only=True, which rejects the pickled module classes of this trusted local file
+original_torch_load = torch.load
+torch.load = lambda *args, **kwargs: original_torch_load(*args, **{**kwargs, "weights_only": False})
+model.load(os.path.join(par["input_model"], "model"))
+torch.load = original_torch_load
+model.params["device"] = device
+model.model.float()
 
-# ---- Preprocess test inputs ----
-print('Preprocessing test data...', flush=True)
-preprocessed_test_inputs, _ = pre_post_process.preprocess(
-    inputs_values=test_inputs,
-    targets_values=None,
-    metadata=test_metadata,
+# ---- Test data and its metadata, standardized with the training statistics ----
+print("Loading test data...", flush=True)
+input_test_mod1 = ad.read_h5ad(par["input_test_mod1"])
+test_inputs = to_sparse_csr(input_test_mod1.layers["normalized"]).astype(np.float32)
+test_batches = input_test_mod1.obs["batch"].astype(str).values
+test_cell_statistics = apply_standardization(
+    compute_cell_statistics(test_inputs, task_type), task_info["cell_statistics_standardization"]
+)
+test_metadata = build_metadata(
+    test_batches,
+    test_cell_statistics,
+    task_type=task_type,
+    day_pattern=task_info["day_pattern"],
+    donor_pattern=task_info["donor_pattern"],
+    batch_sv=batch_singular_vectors.lookup(test_batches) if batch_singular_vectors is not None else None,
+    group_by_batch=False,
 )
 
 # ---- Predict ----
-print('Predicting...', flush=True)
-y_pred = model.predict(
-    x=test_inputs,
-    preprocessed_x=preprocessed_test_inputs,
-    metadata=test_metadata,
-)
-gc.collect()
+print("Preprocessing test data...", flush=True)
+preprocessed_test_inputs, _ = pre_post_process.preprocess(inputs_values=test_inputs, targets_values=None, metadata=test_metadata)
+preprocessed_test_inputs = np.asarray(preprocessed_test_inputs, dtype=np.float32)
 
-# ---- Write output ----
-print('Writing output...', flush=True)
-# Prediction must be a sparse matrix to be compatible with all metrics.
-if not scipy.sparse.issparse(y_pred):
-    y_pred = scipy.sparse.csr_matrix(y_pred)
+print("Predicting...", flush=True)
+predictions = model.predict(x=test_inputs, preprocessed_x=preprocessed_test_inputs, metadata=test_metadata)
+rescaling = task_info["prediction_rescaling"]
+predictions = predictions * rescaling["slope"] + rescaling["intercept"]
+predictions = np.nan_to_num(predictions, nan=rescaling["intercept"], posinf=rescaling["intercept"], neginf=rescaling["intercept"])
+assert predictions.shape == (input_test_mod1.n_obs, mod2_var.shape[0])
 
+# ---- Write ----
+print("Writing output...", flush=True)
 output = ad.AnnData(
-    layers={"normalized": y_pred},
+    layers={"normalized": scipy.sparse.csr_matrix(predictions.astype(np.float32))},
     obs=input_test_mod1.obs,
     var=mod2_var,
-    uns={
-        "dataset_id": dataset_id,
-        "method_id": "ss_opm",
-    },
+    uns={"dataset_id": task_info["dataset_id"], "method_id": "ss_opm"},
 )
-output.write_h5ad(par['output'], compression="gzip")
-print('Done!', flush=True)
+output.write_h5ad(par["output"], compression="gzip")
+print("Done!", flush=True)

src/methods/ss_opm/ss_opm_train/config.vsh.yaml

@@ -4,32 +4,50 @@ arguments:
   - name: "--n_epochs"
     type: integer
     default: 40
-    description: Number of training epochs.
+    description: Number of training epochs (the authors' setting).
     info:
       test_default: 2
   - name: "--burnin_length_epoch"
     type: integer
     default: 10
     description: |
-      Epochs before the training-length ratio starts ramping up. Must be below
-      `--n_epochs`, otherwise every epoch stays at ratio 0.
+      Epochs before the training-length ratio starts ramping up (the authors' setting). Must be
+      below `--n_epochs`, otherwise every epoch stays at ratio 0.
     info:
       test_default: 0
-  - name: "--cell_type_col"
+  - name: "--day_pattern"
     type: string
-    required: false
+    default: '^(\d+)_\d+$'
     description: |
-      Column in `obs` holding cell type labels. ss_opm uses them for the `cell_ratio_*`
-      features. This task's file format does not carry cell types, so it is unset by
-      default and the ratios fall back to uniform; set it when plugging in a dataset
-      that does have them.
-  - name: "--day_pattern"
+      Regex whose first capture group is the collection day in a batch label; the model uses the day
+      as a numeric covariate. The default matches the `{day}_{donor}` labels of the NeurIPS 2022
+      datasets. Batches that do not match (e.g. the `s{site}d{donor}` labels of the NeurIPS 2021
+      datasets, which have no day) get a constant day.
+  - name: "--donor_pattern"
     type: string
-    default: 'd(\d+)'
+    default: '^\d+_(\d+)$'
     description: |
-      Regex whose first capture group is the collection day within a batch label. The
-      default matches the NeurIPS 2021 `s{site}d{day}` naming. Batches that do not match
-      get day 0, i.e. the model sees a single constant day rather than failing.
+      Regex whose first capture group is the donor id in a batch label. It only feeds the original's
+      donor-sex embedding, which knows the four Kaggle donors; other or unmatched donors are embedded
+      like an unknown donor.
+  - name: "--hgnc_complete_set"
+    type: file
+    required: false
+    description: |
+      HGNC complete set (tab separated), used to name the genes and proteins when building the CITE
+      input gene masks. Downloaded into the image when not given.
+  - name: "--reactome_pathways"
+    type: file
+    required: false
+    description: |
+      Reactome gene sets (GMT), used to build the CITE input gene masks. Downloaded into the image
+      when not given.
+  - name: "--n_rescaling_cells"
+    type: integer
+    default: 5000
+    description: |
+      Number of training cells used to fit the global affine map from the model's per-cell z-scored
+      output to the target scale.
 resources:
   - type: python_script
     path: script.py
@@ -44,8 +62,22 @@ engines:
         packages:
           - pyarrow
           - fastparquet
+      # Reference files for the CITE input gene masks (same sources as the original solution). reactome.org rejects
+      # Python's default user agent with HTTP 403, hence the browser-like header.
+      - type: docker
+        run: |
+          mkdir -p /opt/ss_opm_reference && cd /opt/ss_opm_reference && \
+          python -c "import shutil, urllib.request; \
+            [shutil.copyfileobj(urllib.request.urlopen(urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})), open(file_name, 'wb')) \
+             for url, file_name in [('https://storage.googleapis.com/public-download-files/hgnc/archive/archive/monthly/tsv/hgnc_complete_set_2023-01-01.txt', 'hgnc_complete_set.txt'), \
+                                    ('https://reactome.org/download/current/ReactomePathways.gmt.zip', 'ReactomePathways.gmt.zip')]]" && \
+          python -c "import zipfile; zipfile.ZipFile('ReactomePathways.gmt.zip').extractall('.')" && \
+          rm ReactomePathways.gmt.zip
+      - type: docker
+        env:
+          - SS_OPM_REFERENCE_DIR=/opt/ss_opm_reference
 runners:
   - type: executable
   - type: nextflow
     directives:
-      label: [highmem, hightime, midcpu, highsharedmem, gpu]
+      label: [veryhighmem, hightime, midcpu, highsharedmem, gpu]