Add ss_opm

scripts/create_datasets/test_resources.sh

@@ -138,6 +138,19 @@ for name in bmmc_cite/normal bmmc_cite/swap bmmc_multiome/normal bmmc_multiome/s
     fi
   fi
 
+  echo "pre-train ss_opm on $name"
+  if up_to_date $DATASET_DIR/$name/models/ss_opm/ $STATE; then
+    echo "  already up to date, skipping"
+  else
+    rm -rf $DATASET_DIR/$name/models/ss_opm/
+    mkdir -p $DATASET_DIR/$name/models/ss_opm/
+    viash run src/methods/ss_opm/ss_opm_train/config.vsh.yaml -- \
+      --input_train_mod1 $DATASET_DIR/$name/train_mod1.h5ad \
+      --input_train_mod2 $DATASET_DIR/$name/train_mod2.h5ad \
+      --input_test_mod1 $DATASET_DIR/$name/test_mod1.h5ad \
+      --output $DATASET_DIR/$name/models/ss_opm
+  fi
+
   # scbutterfly only does multiome (GEX<->ATAC)
   if [[ "$name" == bmmc_multiome/* ]]; then
     echo "pre-train scbutterfly on $name"

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

@@ -0,0 +1,26 @@
+__merge__: ../../../api/comp_method.yaml
+name: ss_opm
+label: SS-OPM
+summary: 1st place solution of the Kaggle Open Problems Multimodal Single-Cell Integration challenge.
+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.
+references:
+  doi:
+    - 10.1101/2022.04.11.487796
+links:
+  repository: https://github.com/shu65/open-problems-multimodal
+info:
+  preferred_normalization: log_cp10k
+resources:
+  - path: main.nf
+    type: nextflow_script
+    entrypoint: run_wf
+dependencies:
+  - name: methods/ss_opm_train
+  - name: methods/ss_opm_predict
+runners:
+  - type: nextflow

src/methods/ss_opm/ss_opm/main.nf

@@ -0,0 +1,18 @@
+workflow run_wf {
+  take: input_ch
+  main:
+  output_ch = input_ch
+    | ss_opm_train.run(
+      fromState: ["input_train_mod1", "input_train_mod2", "input_test_mod1"],
+      toState: ["input_model": "output"]
+    )
+    | ss_opm_predict.run(
+      fromState: ["input_test_mod1", "input_model"],
+      toState: ["output": "output"]
+    )
+    | map { tup ->
+      [tup[0], [output: tup[1].output]]
+    }
+
+  emit: output_ch
+}

src/methods/ss_opm/ss_opm_common.py

@@ -0,0 +1,109 @@
+"""Helpers shared by ss_opm_train and ss_opm_predict."""
+
+import numpy as np
+import pandas as pd
+import scipy.sparse
+
+# the cell types the original ss_opm model was trained against. only used to name the
+# cell_ratio_* columns it expects; the ratios themselves are derived from the data when
+# cell type labels are available.
+CITE_CELL_TYPES = ["HSC", "EryP", "NeuP", "MasP", "MkP", "BP", "MoP"]
+
+# number of batch singular-vector columns the cite model expects
+N_BATCH_SV = 8
+
+
+def to_sparse_csr(X):
+    if scipy.sparse.issparse(X):
+        return X.tocsr()
+    return scipy.sparse.csr_matrix(X)
+
+
+def extract_day(batch, pattern=r"d(\d+)"):
+    """Pull the day out of a batch label.
+
+    The NeurIPS 2021 batches are named `s{site}d{day}`, e.g. `s1d2`. Datasets that
+    label their batches differently yield NaN, which the caller fills with 0 -- the
+    model then sees a single constant day rather than failing.
+    """
+    return batch.astype(str).str.extract(pattern, expand=False).astype(float)
+
+
+def build_metadata(
+    adata,
+    task_type,
+    cell_type_col=None,
+    group_by_batch=True,
+    day_pattern=r"d(\d+)",
+):
+    """Build the metadata frame ss_opm expects from an AnnData.
+
+    ss_opm was written against the Kaggle competition tables, which carry columns this
+    task's API does not: `file_train_mod1.yaml` and `file_test_mod1.yaml` guarantee only
+    `batch`. Everything else is either derived from `batch`, computed from the expression
+    matrix, or filled with a neutral constant.
+
+    Parameters
+    ----------
+    adata
+        Input modality, with a `normalized` layer and `obs["batch"]`.
+    task_type
+        Either `"cite"` or `"multi"`; the cite model expects extra columns.
+    cell_type_col
+        Column in `adata.obs` holding cell type labels. When given, it drives both
+        `cell_type` and the `cell_ratio_*` columns. When None -- the case for every
+        dataset this task currently ships -- cell types are `"hidden"` and the ratios
+        are uniform.
+    group_by_batch
+        Assign one group per batch. Set False to put every cell in group 0, which is
+        what the predict path wants, since targets are absent and the group IDs are
+        only used to look up target statistics.
+    day_pattern
+        Regex whose first capture group is the day within a batch label.
+    """
+    obs = pd.DataFrame(index=adata.obs_names)
+
+    obs["batch"] = adata.obs["batch"].values
+    obs["day"] = extract_day(adata.obs["batch"], day_pattern).fillna(0).values
+
+    # per-cell statistics from the normalized expression layer
+    X = adata.layers["normalized"]
+    X_dense = X.toarray() if scipy.sparse.issparse(X) else np.asarray(X, dtype=float)
+
+    obs["nonzero_ratio"] = (X_dense != 0).mean(axis=1)
+    obs["nonzero_q25"] = np.percentile(X_dense, 25, axis=1)
+    obs["nonzero_q50"] = np.percentile(X_dense, 50, axis=1)
+    obs["nonzero_q75"] = np.percentile(X_dense, 75, axis=1)
+    obs["mean"] = X_dense.mean(axis=1)
+    obs["std"] = X_dense.std(axis=1)
+
+    if group_by_batch:
+        batches = adata.obs["batch"].unique().tolist()
+        obs["group"] = adata.obs["batch"].map({b: i for i, b in enumerate(batches)}).astype(int).values
+    else:
+        obs["group"] = 0
+
+    # cell type labels, when the caller can supply them
+    if cell_type_col is not None and cell_type_col in adata.obs:
+        obs["cell_type"] = adata.obs[cell_type_col].astype(str).values
+    else:
+        obs["cell_type"] = "hidden"
+
+    # donor and technology are not in this task's file format; gender_id defaults to 0
+    obs["donor"] = 0
+    obs["technology"] = "unknown"
+
+    if task_type == "cite":
+        ratios = obs["cell_type"].value_counts(normalize=True)
+        for cell_type in CITE_CELL_TYPES:
+            obs[f"cell_ratio_{cell_type}"] = ratios.get(cell_type, 1.0 / len(CITE_CELL_TYPES))
+
+        batch_counts = adata.obs["batch"].value_counts()
+        obs["cell_count"] = adata.obs["batch"].map(batch_counts).astype(float).values
+
+        # the originals are singular vectors of the full Kaggle batch matrix, which we
+        # cannot reconstruct from a single dataset
+        for i in range(N_BATCH_SV):
+            obs[f"batch_sv{i}"] = 0.0
+
+    return obs

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

@@ -0,0 +1,42 @@
+__merge__: ../../../api/comp_method_predict.yaml
+name: ss_opm_predict
+
+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
+  - path: ../ss_opm_common.py
+engines:
+  - type: docker
+    image: openproblems/base_pytorch_nvidia:1
+    setup:
+      - type: docker
+        run: pip install --no-cache-dir --no-deps git+https://github.com/shu65/open-problems-multimodal.git
+      - type: python
+        packages:
+          - pyarrow
+          - fastparquet
+runners:
+  - type: executable
+  - type: nextflow
+    directives:
+      label: [highmem, hightime, midcpu, highsharedmem, gpu]

src/methods/ss_opm/ss_opm_predict/script.py

@@ -0,0 +1,102 @@
+import sys
+import os
+import gc
+import pickle
+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',
+}
+## VIASH END
+
+sys.path.append(meta['resources_dir'])
+from ss_opm_common import build_metadata, to_sparse_csr
+
+# ---- 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)
+
+# ---- 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'],
+)
+
+# ---- 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)
+
+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
+
+mod2_var = pd.read_parquet(os.path.join(par['input_model'], 'mod2_var.parquet'))
+
+# ---- 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,
+)
+
+# ---- Predict ----
+print('Predicting...', flush=True)
+y_pred = model.predict(
+    x=test_inputs,
+    preprocessed_x=preprocessed_test_inputs,
+    metadata=test_metadata,
+)
+gc.collect()
+
+# ---- 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)
+
+output = ad.AnnData(
+    layers={"normalized": y_pred},
+    obs=input_test_mod1.obs,
+    var=mod2_var,
+    uns={
+        "dataset_id": dataset_id,
+        "method_id": "ss_opm",
+    },
+)
+output.write_h5ad(par['output'], compression="gzip")
+print('Done!', flush=True)

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

@@ -0,0 +1,51 @@
+__merge__: ../../../api/comp_method_train.yaml
+name: ss_opm_train
+arguments:
+  - name: "--n_epochs"
+    type: integer
+    default: 40
+    description: Number of training epochs.
+    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.
+    info:
+      test_default: 0
+  - 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
+  - path: ../ss_opm_common.py
+engines:
+  - type: docker
+    image: openproblems/base_pytorch_nvidia:1
+    setup:
+      - type: docker
+        run: pip install --no-cache-dir --no-deps git+https://github.com/shu65/open-problems-multimodal.git
+      - type: python
+        packages:
+          - pyarrow
+          - fastparquet
+runners:
+  - type: executable
+  - type: nextflow
+    directives:
+      label: [highmem, hightime, midcpu, highsharedmem, gpu]