From 568eff111f5a004df96abc655a8741b7e6a6f30d Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Sun, 20 Sep 2026 14:11:05 +0200 Subject: [PATCH 1/3] ss_opm: rebuild the original's derived inputs and make the wrapper faithful The wrapper fed the network raw per-cell statistics and the raw batch cell count (thousands in training, tens for the subsampled test batches), took the donor of the NeurIPS 2021 batch labels for the day, appended all genes instead of the original's ~70 HGNC/Reactome-selected genes as raw CITE features, kept cells without any protein counts (NaN loss on the 2022 CITE data), upcast the dense multiome targets to float64 (OOM on 2022 ATAC->GEX) and returned the model's per-cell z-scores as predictions. - per-cell statistics as in make_additional_files.py (quartiles of the non-zero values, log1p ratio for multiome), standardized over train+test - CITE batch singular vectors from per-batch gene medians, standardized over the training batches, test-only batches projected; cell ratios and cell count constant (no cell types in the task files, subsampled test set) - day/donor regexes default to the {day}_{donor} labels of NeurIPS 2022, constant otherwise - CITE gene masks rebuilt with the original thresholds from HGNC (Ensembl id -> symbol) and Reactome, downloaded into the image - training cells with a constant target vector dropped - dtype-preserving, zero-row-safe patches of the ss_opm normalizers - predictions mapped to the target scale by a global affine fit on training cells (correlations unchanged) Co-Authored-By: Claude Fable 5.1 --- src/methods/ss_opm/ss_opm/config.vsh.yaml | 8 +- src/methods/ss_opm/ss_opm_common.py | 602 ++++++++++++++---- .../ss_opm/ss_opm_predict/config.vsh.yaml | 18 +- src/methods/ss_opm/ss_opm_predict/script.py | 142 ++--- .../ss_opm/ss_opm_train/config.vsh.yaml | 58 +- src/methods/ss_opm/ss_opm_train/script.py | 346 ++++++---- 6 files changed, 809 insertions(+), 365 deletions(-) diff --git a/src/methods/ss_opm/ss_opm/config.vsh.yaml b/src/methods/ss_opm/ss_opm/config.vsh.yaml index beec3faf..018c91ca 100644 --- a/src/methods/ss_opm/ss_opm/config.vsh.yaml +++ b/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 diff --git a/src/methods/ss_opm/ss_opm_common.py b/src/methods/ss_opm/ss_opm_common.py index 3a55ed9a..4e41e02d 100644 --- a/src/methods/ss_opm/ss_opm_common.py +++ b/src/methods/ss_opm/ss_opm_common.py @@ -1,19 +1,47 @@ -"""Helpers shared by ss_opm_train and ss_opm_predict.""" +"""Helpers shared by ss_opm_train and ss_opm_predict. + +ss_opm (https://github.com/shu65/open-problems-multimodal) was written against the Kaggle competition tables and a +set of derived files produced by its `script/make_additional_files.py` and `script/make_cite_input_mask.py`: +z-scored per-cell statistics, z-scored per-batch statistics, a numeric `day`, a donor id (turned into a sex +embedding) and, for CITE-seq, a mask of a few dozen genes whose raw expression is appended to the SVD components. +This module rebuilds all of that from the task's h5ad files, which only guarantee `obs["batch"]` and a +`normalized` layer. +""" + +import json +import os +import re +import zipfile +from urllib.request import urlretrieve import numpy as np import pandas as pd import scipy.sparse +import scipy.stats +from sklearn.decomposition import TruncatedSVD -# cells per block when densifying the expression matrix for per-cell statistics -ROW_BLOCK = 1000 +from ss_opm.model.torch_dataset.citeseq_dataset import METADATA_KEYS as CITE_METADATA_KEYS +from ss_opm.model.torch_dataset.multiome_dataset import METADATA_KEYS as MULTI_METADATA_KEYS +from ss_opm.utility.metadata_utility import CELL_TYPES as KAGGLE_CELL_TYPES -# 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"] +# cells per block when densifying a sparse matrix for per-cell or per-batch statistics +ROW_BLOCK = 2000 -# number of batch singular-vector columns the cite model expects -N_BATCH_SV = 8 +# number of batch singular-vector columns the CITE model expects (see CITE_METADATA_KEYS) +N_BATCH_SV = sum(key.startswith("batch_sv") for key in CITE_METADATA_KEYS) + +CELL_STATISTIC_KEYS = ["nonzero_ratio", "nonzero_q25", "nonzero_q50", "nonzero_q75", "mean", "std"] + +# Batch labels of the OpenProblems NeurIPS 2022 datasets are `{day}_{donor}`; the NeurIPS 2021 ones are +# `s{site}d{donor}` and carry no day. Anything that does not match gets a constant. +DEFAULT_DAY_PATTERN = r"^(\d+)_\d+$" +DEFAULT_DONOR_PATTERN = r"^\d+_(\d+)$" + +# Reference files the CITE gene masks are built from (same sources as the original `make_cite_input_mask.py`) +HGNC_URL = "https://storage.googleapis.com/public-download-files/hgnc/archive/archive/monthly/tsv/hgnc_complete_set_2023-01-01.txt" +REACTOME_URL = "https://reactome.org/download/current/ReactomePathways.gmt.zip" + +ENSEMBL_ID_PATTERN = re.compile(r"^ENSG\d+") def to_sparse_csr(X): @@ -22,142 +50,472 @@ def to_sparse_csr(X): return scipy.sparse.csr_matrix(X) -def extract_day(batch, pattern=r"d(\d+)"): - """Pull the day out of a batch label. +def to_dense(X, dtype=np.float32): + dense = X.toarray() if scipy.sparse.issparse(X) else np.asarray(X) + return dense.astype(dtype, copy=False) + + +# --------------------------------------------------------------------------------------------------------------- +# Reference files +# --------------------------------------------------------------------------------------------------------------- +def download_reference_files(directory): + """Download the HGNC complete set and the Reactome gene sets into `directory` (used at image build time).""" + os.makedirs(directory, exist_ok=True) + hgnc_path = os.path.join(directory, "hgnc_complete_set.txt") + reactome_path = os.path.join(directory, "ReactomePathways.gmt") + if not os.path.exists(hgnc_path): + urlretrieve(HGNC_URL, hgnc_path) + if not os.path.exists(reactome_path): + archive_path = os.path.join(directory, "ReactomePathways.gmt.zip") + urlretrieve(REACTOME_URL, archive_path) + with zipfile.ZipFile(archive_path) as archive: + archive.extractall(directory) + os.remove(archive_path) + return hgnc_path, reactome_path + + +def read_reactome_gmt(file_path): + """Pathway (Reactome stable id) -> gene symbols, from a GMT file (name, id, genes...).""" + pathway_genes = {} + with open(file_path) as handle: + for line in handle: + fields = line.rstrip("\n").split("\t") + if len(fields) < 3: + continue + pathway_genes[fields[1]] = [gene.upper() for gene in fields[2:] if gene] + return pathway_genes + + +def read_hgnc(file_path): + columns = ["symbol", "alias_symbol", "prev_symbol", "ensembl_gene_id"] + return pd.read_table(file_path, usecols=columns, dtype=str, low_memory=False) + + +def gene_symbols_from_var_names(var_names, hgnc): + """Upper-case gene symbols for feature names that are Ensembl ids (`ENSG...` or Kaggle's `ENSG..._SYMBOL`); + other names (already symbols, e.g. protein names) are kept.""" + ensembl_to_symbol = hgnc.dropna(subset=["ensembl_gene_id"]).set_index("ensembl_gene_id")["symbol"].to_dict() + symbols = [] + for name in var_names: + name = str(name) + if ENSEMBL_ID_PATTERN.match(name): + if "_" in name: # Kaggle style: ENSG00000121410_A1BG + symbols.append(name.split("_", 1)[1].upper()) + else: + symbols.append(str(ensembl_to_symbol.get(name.split(".")[0], name)).upper()) + else: + symbols.append(name.upper()) + return np.array(symbols) + + +def make_targets_gene2idx(target_symbols, hgnc): + """Port of the original `make_targets_gene2idx`: map every protein name, and every HGNC symbol for which that + protein name is an alias or previous symbol, to the protein's column index.""" + alias_symbols = {} + for symbol, aliases, previous in hgnc[["symbol", "alias_symbol", "prev_symbol"]].itertuples(index=False): + for value in (aliases, previous): + if pd.isnull(value): + continue + for alias in value.upper().split("|"): + alias_symbols.setdefault(alias, []).append(symbol.upper()) + + alias_symbols["CD3"] = ["CD3D", "CD3E", "CD3G"] + alias_symbols["HLA-A-B-C"] = ["HLA-A", "HLA-B", "HLA-C"] + alias_symbols["CD45RA"] = ["PTPRC"] + alias_symbols["CD45RO"] = ["PTPRC"] + alias_symbols["PODOPLANIN"] = ["PDPN"] + alias_symbols["HLA-DR"] = [f"HLA-DRB{i}" for i in range(1, 10)] + ["HLA-DRA"] + alias_symbols["INTEGRINB7"] = ["ITGB7"] + alias_symbols["CD158"] = ["CD158A"] + alias_symbols["CD158B"] = ["CD158B1", "CD158B2"] + + targets_gene2idx = {} + for target_index, target_symbol in enumerate(target_symbols): + target_symbol = target_symbol.upper() + targets_gene2idx[target_symbol] = target_index + for symbol in alias_symbols.get(target_symbol, []): + targets_gene2idx[symbol] = target_index + return targets_gene2idx + + +def _group_spearman(inputs_column, targets_columns): + """Spearman correlation and p-value of one gene with several targets, over the cells where the gene is + expressed (as in the original scripts). Returns arrays of NaN when there are too few such cells.""" + expressed = inputs_column > 0 + n_targets = targets_columns.shape[1] + if expressed.sum() < 3: + return np.full(n_targets, np.nan), np.full(n_targets, np.nan) + result = scipy.stats.spearmanr(inputs_column[expressed], targets_columns[expressed]) + correlations = np.atleast_2d(result.statistic)[0, 1:] + p_values = np.atleast_2d(result.pvalue)[0, 1:] + return correlations, p_values + + +def _select_robust_pairs(correlations, min_abs_corr, max_p_value, n_groups): + """`correlations`: dict (gene index, target index) -> list of (|corr|, p) per group. Keep the pairs that pass + the thresholds in more than 60 % of the groups and return their median |corr| over the groups that passed, + mirroring the original code (which stored 0 for failing groups and took the median over all groups).""" + robust = {} + for pair, values in correlations.items(): + scores = np.zeros(n_groups) + for group_index, (corr, p_value) in enumerate(values): + if np.isfinite(corr) and abs(corr) > min_abs_corr and p_value < max_p_value: + scores[group_index] = abs(corr) + if (scores > 0).sum() > 0.6 * n_groups: + robust[pair] = np.median(scores) + return robust + + +def make_cite_input_masks(inputs_lognorm, targets_rownorm, gene_symbols, targets_gene2idx, groups, pathway_genes): + """Rebuild `cite_inputs_targets_pair3g.npz` and `cite_inputs_mask2.npz` of the original solution. - 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. + Parameters + ---------- + inputs_lognorm : np.ndarray or scipy.sparse.csr_matrix + Training RNA, log1p of the median-normalized expression (the original's input transform). + targets_rownorm : np.ndarray + Training proteins, row-normalized. + gene_symbols : np.ndarray + Upper-case symbol of every input gene. + targets_gene2idx : dict + Gene symbol -> protein column index (see `make_targets_gene2idx`). + groups : np.ndarray + Group (batch) label of every training cell. + pathway_genes : dict + Reactome pathway id -> gene symbols. + + Returns + ------- + pair_mask : np.ndarray of bool, shape (n_genes, n_proteins) + For every protein, the gene encoding it (or an alias) whose expression correlates most robustly with it. + pathway_mask : np.ndarray of bool, shape (n_genes,) + Up to three genes per protein sharing a Reactome pathway with the protein's gene and correlating with it. """ - return batch.astype(str).str.extract(pattern, expand=False).astype(float) + inputs_lognorm = to_sparse_csr(inputs_lognorm).tocsc() + n_genes, n_proteins = inputs_lognorm.shape[1], targets_rownorm.shape[1] + unique_groups = np.unique(groups) + group_masks = [groups == group for group in unique_groups] + + # candidate (gene, targets) pairs: genes named after a protein, and genes sharing a pathway with such a gene + direct_candidates = {} + for gene_index, symbol in enumerate(gene_symbols): + if symbol in targets_gene2idx: + direct_candidates[gene_index] = [targets_gene2idx[symbol]] + + pathway_targets = {} + for genes in pathway_genes.values(): + targets_in_pathway = {targets_gene2idx[gene] for gene in genes if gene in targets_gene2idx} + if not targets_in_pathway: + continue + for gene in genes: + pathway_targets.setdefault(gene, set()).update(targets_in_pathway) + pathway_candidates = { + gene_index: sorted(pathway_targets[symbol]) + for gene_index, symbol in enumerate(gene_symbols) + if symbol in pathway_targets + } + + def _correlate(candidates): + correlations = {} + for gene_index, target_indices in candidates.items(): + gene_values = inputs_lognorm[:, gene_index].toarray().ravel() + for group_index, group_mask in enumerate(group_masks): + corrs, p_values = _group_spearman(gene_values[group_mask], targets_rownorm[group_mask][:, target_indices]) + for target_index, corr, p_value in zip(target_indices, corrs, p_values): + correlations.setdefault((gene_index, target_index), []).append((corr, p_value)) + return correlations + + direct_pairs = _select_robust_pairs(_correlate(direct_candidates), 0.10, 1e-2, len(unique_groups)) + pair_mask = np.zeros((n_genes, n_proteins), dtype=bool) + for target_index in range(n_proteins): + scored = [(score, gene_index) for (gene_index, t), score in direct_pairs.items() if t == target_index] + if scored: + pair_mask[max(scored)[1], target_index] = True + + pathway_pairs = _select_robust_pairs(_correlate(pathway_candidates), 0.20, 1e-3, len(unique_groups)) + pathway_mask = np.zeros(n_genes, dtype=bool) + for target_index in range(n_proteins): + scored = sorted(((score, gene_index) for (gene_index, t), score in pathway_pairs.items() if t == target_index), reverse=True) + for _, gene_index in scored[:3]: + pathway_mask[gene_index] = True + + return pair_mask, pathway_mask + + +# --------------------------------------------------------------------------------------------------------------- +# Metadata +# --------------------------------------------------------------------------------------------------------------- +def extract_integer_field(batch, pattern): + """First capture group of `pattern` in every batch label as float, NaN where it does not match.""" + if pattern is None: + return np.full(len(batch), np.nan) + return pd.Series(batch).astype(str).str.extract(pattern, expand=False).astype(float).values + + +def compute_cell_statistics(X, task_type): + """Per-cell statistics of the normalized expression as in the original `make_*_cell_statistics`: fraction of + expressed features (log1p of it for multiome), quartiles of the non-zero values, mean and std of the full row.""" + X = to_sparse_csr(X) + n_features = X.shape[1] + nonzero_counts = np.diff(X.indptr) + row_sums = np.asarray(X.sum(axis=1)).ravel() + row_square_sums = np.asarray(X.multiply(X).sum(axis=1)).ravel() + means = row_sums / n_features + stds = np.sqrt(np.maximum(row_square_sums / n_features - means**2, 0)) + + quartiles = np.zeros((X.shape[0], 3)) + for row, values in enumerate(np.split(X.data, X.indptr[1:-1])): + if len(values): + quartiles[row] = np.quantile(values, [0.25, 0.5, 0.75]) + + nonzero_ratio = nonzero_counts / n_features + if task_type == "multi": + nonzero_ratio = np.log1p(nonzero_ratio) + return pd.DataFrame( + { + "nonzero_ratio": nonzero_ratio, + "nonzero_q25": quartiles[:, 0], + "nonzero_q50": quartiles[:, 1], + "nonzero_q75": quartiles[:, 2], + "mean": means, + "std": stds, + } + ) + + +def fit_standardization(frame): + """Mean and std of every column; constant columns get std 1 so that they standardize to 0.""" + means = frame.mean(axis=0) + stds = frame.std(axis=0, ddof=1).replace(0, 1.0).fillna(1.0) + return {"mean": means.to_dict(), "std": stds.to_dict()} + + +def apply_standardization(frame, standardization): + standardized = frame.copy() + for column in frame.columns: + standardized[column] = (frame[column] - standardization["mean"][column]) / standardization["std"][column] + return standardized + + +def median_normalized_log_expression(X_block): + """The original CITE input transform: log1p of every cell divided by its median non-zero expression.""" + dense = np.expm1(to_dense(X_block)) + with np.errstate(invalid="ignore"): + for_median = np.where(dense == 0, np.nan, dense) + medians = np.nanmedian(for_median, axis=1) + medians = np.where(np.isfinite(medians) & (medians > 0), medians, 1.0) + return np.log1p(dense / medians[:, None]) + + +def compute_batch_input_medians(X, batches): + """Per batch, the median over its cells of the (median-normalized, log1p) expression of every gene, ignoring + zeros, as in the original `make_cite_batch_inputs_median`. Returns a DataFrame indexed by batch.""" + X = to_sparse_csr(X) + batches = np.asarray(batches).astype(str) + rows = {} + for batch in np.unique(batches): + cell_indices = np.flatnonzero(batches == batch) + values = np.vstack( + [median_normalized_log_expression(X[cell_indices[start : start + ROW_BLOCK]]) for start in range(0, len(cell_indices), ROW_BLOCK)] + ) + with np.errstate(invalid="ignore"): + values[values == 0] = np.nan + gene_medians = np.nanmedian(values, axis=0) + gene_medians[~np.isfinite(gene_medians)] = 0.0 + rows[batch] = gene_medians + return pd.DataFrame.from_dict(rows, orient="index") + + +class BatchSingularVectors: + """`batch_sv0..7` of the original CITE metadata: TruncatedSVD of the per-batch gene medians, standardized + across the batches the model is trained on. Batches first seen at prediction time are projected with the + fitted components (unknown batches get 0).""" + + def __init__(self, n_components=N_BATCH_SV, random_state=42): + self.n_components = n_components + self.random_state = random_state + self.decomposer = None + self.standardization = None + self.table = None + + def fit(self, batch_medians): + n_components = min(self.n_components, min(batch_medians.shape) - 1) if min(batch_medians.shape) > 1 else 0 + if n_components < 1: + self.decomposer = None + self.table = pd.DataFrame(0.0, index=batch_medians.index, columns=self.columns) + return self + self.decomposer = TruncatedSVD(n_components=n_components, random_state=self.random_state) + transformed = self.decomposer.fit_transform(batch_medians.values) + frame = self._to_frame(transformed, batch_medians.index) + self.standardization = fit_standardization(frame) + self.table = apply_standardization(frame, self.standardization) + return self + + def transform(self, batch_medians): + if self.decomposer is None: + return pd.DataFrame(0.0, index=batch_medians.index, columns=self.columns) + frame = self._to_frame(self.decomposer.transform(batch_medians.values), batch_medians.index) + return apply_standardization(frame, self.standardization) + + @property + def columns(self): + return [f"batch_sv{i}" for i in range(self.n_components)] + + def _to_frame(self, transformed, index): + frame = pd.DataFrame(0.0, index=index, columns=self.columns) + frame.iloc[:, : transformed.shape[1]] = transformed + return frame + + def lookup(self, batches): + """Rows of the batch table for every cell; unknown batches get 0.""" + return self.table.reindex(np.asarray(batches).astype(str)).fillna(0.0).reset_index(drop=True) def build_metadata( - adata, + batch, + cell_statistics, task_type, - cell_type_col=None, + day_pattern=DEFAULT_DAY_PATTERN, + donor_pattern=DEFAULT_DONOR_PATTERN, + batch_sv=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. + """The metadata frame ss_opm's datasets and pre-processing expect. 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. + batch : array-like + Batch label of every cell. + cell_statistics : pd.DataFrame + Standardized per-cell statistics (see `compute_cell_statistics`, `apply_standardization`). + task_type : str + "cite" or "multi"; the CITE model expects the batch-level columns as well. + day_pattern, donor_pattern : str or None + Regex whose first capture group is the day / donor id in a batch label. Cells whose label does not match + get day 0 and donor -1. The donor only feeds the original's sex embedding, which knows the four Kaggle + donors; any other donor is embedded like an unknown one. + batch_sv : pd.DataFrame or None + Standardized `batch_sv*` columns per cell (see `BatchSingularVectors.lookup`); zeros when None. + group_by_batch : bool + Group id used for the per-batch target medians. The predict path has no targets, so it can use one group. """ - 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, densified one row - # block at a time -- the whole matrix is 222 GiB on the multiome datasets - X = adata.layers["normalized"] - names = ("nonzero_ratio", "nonzero_q25", "nonzero_q50", "nonzero_q75", "mean", "std") - stats = {name: np.empty(adata.n_obs, dtype=float) for name in names} - - for start in range(0, adata.n_obs, ROW_BLOCK): - end = min(start + ROW_BLOCK, adata.n_obs) - block = X[start:end] - block = block.toarray() if scipy.sparse.issparse(block) else np.asarray(block, dtype=float) - - stats["nonzero_ratio"][start:end] = (block != 0).mean(axis=1) - stats["nonzero_q25"][start:end] = np.percentile(block, 25, axis=1) - stats["nonzero_q50"][start:end] = np.percentile(block, 50, axis=1) - stats["nonzero_q75"][start:end] = np.percentile(block, 75, axis=1) - stats["mean"][start:end] = block.mean(axis=1) - stats["std"][start:end] = block.std(axis=1) - - for name in names: - obs[name] = stats[name] - + batch = pd.Series(np.asarray(batch).astype(str)) + metadata = pd.DataFrame(index=range(len(batch))) + metadata["batch"] = batch.values + metadata["day"] = np.nan_to_num(extract_integer_field(batch, day_pattern), nan=0.0) + donor = extract_integer_field(batch, donor_pattern) + metadata["donor"] = np.where(np.isfinite(donor), donor, -1).astype(int) + metadata["technology"] = "citeseq" if task_type == "cite" else "multiome" + # cell types are not part of this task's file format; "hidden" is the label the original used for test cells + metadata["cell_type"] = "hidden" 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 + metadata["group"] = pd.factorize(batch)[0] else: - obs["cell_type"] = "hidden" + metadata["group"] = 0 - # donor and technology are not in this task's file format; gender_id defaults to 0 - obs["donor"] = 0 - obs["technology"] = "unknown" + for key in CELL_STATISTIC_KEYS: + metadata[key] = cell_statistics[key].values 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 - - + # Cell-type ratios per batch need cell type labels, which the task does not provide, and the batch cell + # count is meaningless for a subsampled test set: both are the standardized value of the mean, 0. + for cell_type in KAGGLE_CELL_TYPES: + if cell_type != "hidden": + metadata[f"cell_ratio_{cell_type}"] = 0.0 + metadata["cell_count"] = 0.0 + for column in [f"batch_sv{i}" for i in range(N_BATCH_SV)]: + metadata[column] = 0.0 if batch_sv is None else batch_sv[column].values + + expected_keys = CITE_METADATA_KEYS if task_type == "cite" else MULTI_METADATA_KEYS + missing = [key for key in expected_keys if key not in metadata.columns] + assert not missing, f"metadata is missing {missing}" + return metadata + + +# --------------------------------------------------------------------------------------------------------------- +# Training targets and prediction scale +# --------------------------------------------------------------------------------------------------------------- +def informative_cells(targets): + """Cells whose target vector is not constant. The model is trained with a per-cell correlation loss and + row-normalizes the targets, both undefined for a constant row (e.g. a cell without any protein counts).""" + targets = to_sparse_csr(targets) + n_features = targets.shape[1] + row_sums = np.asarray(targets.sum(axis=1)).ravel() + row_square_sums = np.asarray(targets.multiply(targets).sum(axis=1)).ravel() + variances = row_square_sums / n_features - (row_sums / n_features) ** 2 + return variances > 1e-12 + + +def fit_prediction_rescaling(predictions, targets): + """The model predicts per-cell z-scores. One global affine map to the target scale keeps every per-cell and + per-feature correlation unchanged and makes the RMSE/MAE metrics meaningful.""" + slope, intercept = np.polyfit(np.asarray(predictions, dtype=np.float64).ravel(), to_dense(targets, np.float64).ravel(), deg=1) + return {"slope": float(slope), "intercept": float(intercept)} + + +def save_json(path, payload): + with open(path, "w") as handle: + json.dump(payload, handle, indent=2) + + +def load_json(path): + with open(path) as handle: + return json.load(handle) + + +# --------------------------------------------------------------------------------------------------------------- +# Runtime patches of the ss_opm package +# --------------------------------------------------------------------------------------------------------------- def _safe_median_normalize(values, ignore_zero=True, log=False): - """Median-normalize rows, substituting 1 (identity) when the median is 0 or NaN.""" - arr = np.asarray(values.toarray() if hasattr(values, 'toarray') else values, dtype=float).copy() - for_median = arr.copy() - if ignore_zero: - for_median[for_median == 0] = np.nan - med = np.nanquantile(for_median, q=0.5, axis=1) - # Use 1 as fallback so rows with zero/undefined median are left unchanged - med = np.where((med == 0) | ~np.isfinite(med), 1.0, med) + """`median_normalize` of ss_opm, keeping the input dtype (the original upcasts float32 data to float64, which + doubles the memory of the dense multiome targets) and leaving rows with an undefined or zero median unchanged.""" + dense = to_dense(values, dtype=values.dtype if isinstance(values, np.ndarray) else np.float32) + medians = np.empty(dense.shape[0], dtype=np.float64) + for start in range(0, dense.shape[0], ROW_BLOCK): + block = dense[start : start + ROW_BLOCK].astype(np.float64) + if ignore_zero: + block[block == 0] = np.nan + with np.errstate(invalid="ignore"): + medians[start : start + ROW_BLOCK] = np.nanquantile(block, q=0.5, axis=1) + medians = np.where(np.isfinite(medians) & (medians != 0), medians, 1.0).astype(dense.dtype) if log: - return arr - med[:, None] - else: - return arr / med[:, None] + return dense - medians[:, None] + return dense / medians[:, None] + + +def _safe_row_normalize(values): + """`row_normalize` of ss_opm; constant rows become zeros instead of NaN.""" + means = np.mean(values, axis=1, keepdims=True) + stds = np.std(values, axis=1, keepdims=True) + stds[stds == 0] = 1.0 + return (values - means) / stds -def _safe_row_normalize(v): - """Row-standardize; rows with std=0 are mean-subtracted only (result is zeros).""" - mu = np.mean(v, axis=1) - sigma = np.std(v, axis=1) - sigma = np.where(sigma == 0, 1.0, sigma) - return (v - mu[:, None]) / sigma[:, None] +def _safe_row_quantile_normalize(values, q=0.5): + """`row_quantile_normalize` of ss_opm without mutating its input and skipping rows without non-zero values.""" + values = values.tocsr() + normalized_data = values.data.astype(np.float32, copy=True) + for row, (start, end) in enumerate(zip(values.indptr[:-1], values.indptr[1:])): + if end > start: + quantile = np.quantile(normalized_data[start:end], q=q) + if quantile > 0: + normalized_data[start:end] /= quantile + return scipy.sparse.csr_matrix((normalized_data, values.indices, values.indptr), values.shape) def apply_runtime_patches(): - """Swap in all-zero-row-safe normalizers for every caller inside ss_opm. + """Swap in the dtype-preserving, zero-row-safe normalizers for every caller inside ss_opm. - Train and predict run the same preprocessing chain, so both have to call this - before `PrePostProcessing.preprocess()`. + Train and predict run the same preprocessing chain, so both have to call this before `PrePostProcessing`. """ - import ss_opm.utility.nonzero_median_normalize as _mnm_module - import ss_opm.utility.row_normalize as _rn_module - import ss_opm.pre_post_processing.pre_post_processing as _pp_module - - # patch the source modules so every 'from X import Y' binding stays in sync - _mnm_module.median_normalize = _safe_median_normalize - _rn_module.row_normalize = _safe_row_normalize - # and the names already bound inside pre_post_processing's namespace - _pp_module.median_normalize = _safe_median_normalize - _pp_module.row_normalize = _safe_row_normalize + import ss_opm.pre_post_processing.pre_post_processing as pre_post_processing_module + import ss_opm.utility.nonzero_median_normalize as median_normalize_module + import ss_opm.utility.row_normalize as row_normalize_module + + median_normalize_module.median_normalize = _safe_median_normalize + median_normalize_module.row_quantile_normalize = _safe_row_quantile_normalize + row_normalize_module.row_normalize = _safe_row_normalize + # names already bound inside pre_post_processing's namespace by `from X import Y` + pre_post_processing_module.median_normalize = _safe_median_normalize + pre_post_processing_module.row_quantile_normalize = _safe_row_quantile_normalize + pre_post_processing_module.row_normalize = _safe_row_normalize diff --git a/src/methods/ss_opm/ss_opm_predict/config.vsh.yaml b/src/methods/ss_opm/ss_opm_predict/config.vsh.yaml index e9224c77..112c6fd4 100644 --- a/src/methods/ss_opm/ss_opm_predict/config.vsh.yaml +++ b/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] diff --git a/src/methods/ss_opm/ss_opm_predict/script.py b/src/methods/ss_opm/ss_opm_predict/script.py index ab7b02d3..df6cc886 100644 --- a/src/methods/ss_opm/ss_opm_predict/script.py +++ b/src/methods/ss_opm/ss_opm_predict/script.py @@ -1,104 +1,96 @@ -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_json, + 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 = load_json(os.path.join(par["input_model"], "task_info.json")) +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) +with open(os.path.join(par["input_model"], "batch_singular_vectors.pickle"), "rb") as handle: + batch_singular_vectors = 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) diff --git a/src/methods/ss_opm/ss_opm_train/config.vsh.yaml b/src/methods/ss_opm/ss_opm_train/config.vsh.yaml index 50bccceb..bcb1615e 100644 --- a/src/methods/ss_opm/ss_opm_train/config.vsh.yaml +++ b/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,18 @@ engines: packages: - pyarrow - fastparquet + # reference files for the CITE input gene masks (same sources as the original solution) + - type: docker + run: | + mkdir -p /opt/ss_opm_reference && cd /opt/ss_opm_reference && \ + python -c "import urllib.request; urllib.request.urlretrieve('https://storage.googleapis.com/public-download-files/hgnc/archive/archive/monthly/tsv/hgnc_complete_set_2023-01-01.txt', 'hgnc_complete_set.txt')" && \ + python -c "import urllib.request, zipfile; urllib.request.urlretrieve('https://reactome.org/download/current/ReactomePathways.gmt.zip', 'ReactomePathways.gmt.zip'); 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] diff --git a/src/methods/ss_opm/ss_opm_train/script.py b/src/methods/ss_opm/ss_opm_train/script.py index b5f2d16c..69883840 100644 --- a/src/methods/ss_opm/ss_opm_train/script.py +++ b/src/methods/ss_opm/ss_opm_train/script.py @@ -1,167 +1,223 @@ -import sys -import os import gc +import os import pickle +import sys import tempfile + +import anndata as ad import numpy as np import pandas as pd -import scipy.sparse -import anndata as ad - -from ss_opm.pre_post_processing.pre_post_processing import PrePostProcessing -from ss_opm.model.encoder_decoder.encoder_decoder import EncoderDecoder -from ss_opm.utility.set_seed import set_seed - import torch -device = 'cuda' if torch.cuda.is_available() else 'cpu' -print(f'Using device: {device}', flush=True) ## 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', - 'output': 'output/models/ss_opm', - 'cell_type_col': None, - 'day_pattern': r'd(\d+)', - 'n_epochs': 40, - 'burnin_length_epoch': 10, -} -meta = { - 'name': 'ss_opm_train', - 'resources_dir': 'src/methods/ss_opm', + "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", + "output": "output/models/ss_opm", + "day_pattern": r"^(\d+)_\d+$", + "donor_pattern": r"^\d+_(\d+)$", + "hgnc_complete_set": None, + "reactome_pathways": None, + "n_epochs": 40, + "burnin_length_epoch": 10, + "n_rescaling_cells": 5000, } +meta = {"name": "ss_opm_train", "resources_dir": "src/methods/ss_opm", "cpus": None} ## 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 + ROW_BLOCK, + BatchSingularVectors, + apply_runtime_patches, + apply_standardization, + build_metadata, + compute_batch_input_medians, + compute_cell_statistics, + download_reference_files, + fit_prediction_rescaling, + fit_standardization, + gene_symbols_from_var_names, + informative_cells, + make_cite_input_masks, + make_targets_gene2idx, + median_normalized_log_expression, + read_hgnc, + read_reactome_gmt, + save_json, + to_dense, + to_sparse_csr, +) apply_runtime_patches() -# The SVD decomposer components are stored as float64 tensors inside -# MultiEncoderDecoderModule, but the neural-network outputs are float32. -# Patch _train_step_forward to convert the whole sub-model to float32 -# immediately before any forward pass, so all tensors share the same dtype. -import ss_opm.model.encoder_decoder.encoder_decoder as _ed_module +from ss_opm.model.encoder_decoder.encoder_decoder import EncoderDecoder # noqa: E402 +from ss_opm.pre_post_processing.pre_post_processing import PrePostProcessing # noqa: E402 +from ss_opm.utility.row_normalize import row_normalize # noqa: E402 +from ss_opm.utility.set_seed import set_seed # noqa: E402 -_orig_train_step_fwd = _ed_module.EncoderDecoder._train_step_forward +device = "cuda" if torch.cuda.is_available() else "cpu" +print(f"Using device: {device}", flush=True) -def _patched_train_step_fwd(self, batch, training_length_ratio): - if hasattr(self, 'model') and self.model is not None: - self.model.float() - return _orig_train_step_fwd(self, batch, training_length_ratio) - -_ed_module.EncoderDecoder._train_step_forward = _patched_train_step_fwd +# The SVD components are stored as float64 tensors inside the torch module while the network is float32: +# cast the module before every forward pass (a no-op after the first one). +import ss_opm.model.encoder_decoder.encoder_decoder as encoder_decoder_module # noqa: E402 -SEED = 42 -set_seed(SEED) +_original_train_step_forward = encoder_decoder_module.EncoderDecoder._train_step_forward -# ---- Load data ---- -print('Loading data...', flush=True) -input_train_mod1 = ad.read_h5ad(par['input_train_mod1']) -input_train_mod2 = ad.read_h5ad(par['input_train_mod2']) +def _float32_train_step_forward(self, batch, training_length_ratio): + if getattr(self, "model", None) is not None: + self.model.float() + return _original_train_step_forward(self, batch, training_length_ratio) -mod1 = input_train_mod1.uns['modality'] -mod2 = input_train_mod2.uns['modality'] -dataset_id = input_train_mod1.uns['dataset_id'] -print(f'Modalities: {mod1} -> {mod2}', flush=True) -# Determine task type: 'cite' when ADT is involved, 'multi' for ATAC/GEX -task_type = 'cite' if 'ADT' in (mod1, mod2) else 'multi' -print(f'Task type: {task_type}', flush=True) +encoder_decoder_module.EncoderDecoder._train_step_forward = _float32_train_step_forward -train_inputs = to_sparse_csr(input_train_mod1.layers['normalized']) -train_targets = to_sparse_csr(input_train_mod2.layers['normalized']) -n_vars_mod1 = train_inputs.shape[1] -n_vars_mod2 = train_targets.shape[1] +# The DataLoader spawns one worker per OMP_NUM_THREADS (the original's convention); keep it within the allocation +if meta.get("cpus"): + os.environ.setdefault("OMP_NUM_THREADS", str(max(1, min(int(meta["cpus"]), 8)))) -train_metadata = build_metadata( - input_train_mod1, - task_type, - cell_type_col=par['cell_type_col'], - day_pattern=par['day_pattern'], -) +SEED = 42 +set_seed(SEED) -# Store mod2 var for the predict step +# ---- Load data ---- +print("Loading data...", flush=True) +input_train_mod1 = ad.read_h5ad(par["input_train_mod1"]) +input_train_mod2 = ad.read_h5ad(par["input_train_mod2"]) +input_test_mod1 = ad.read_h5ad(par["input_test_mod1"]) if par.get("input_test_mod1") else None + +mod1 = input_train_mod1.uns["modality"] +mod2 = input_train_mod2.uns["modality"] +dataset_id = input_train_mod1.uns["dataset_id"] +print(f"Modalities: {mod1} -> {mod2}", flush=True) + +# 'cite' when ADT is involved (the CITE model expects protein targets or inputs), 'multi' for ATAC/GEX +task_type = "cite" if "ADT" in (mod1, mod2) else "multi" +print(f"Task type: {task_type}", flush=True) + +train_inputs = to_sparse_csr(input_train_mod1.layers["normalized"]).astype(np.float32) +train_targets = to_sparse_csr(input_train_mod2.layers["normalized"]).astype(np.float32) +train_batches = input_train_mod1.obs["batch"].astype(str).values +test_inputs = to_sparse_csr(input_test_mod1.layers["normalized"]).astype(np.float32) if input_test_mod1 is not None else None +test_batches = input_test_mod1.obs["batch"].astype(str).values if input_test_mod1 is not None else None +mod1_var_names = input_train_mod1.var_names.to_numpy() mod2_var = input_train_mod2.var.copy() - -del input_train_mod1, input_train_mod2 +del input_train_mod1, input_train_mod2, input_test_mod1 gc.collect() -# ---- Load test inputs for SVD fitting (optional but improves preprocessing) ---- -test_inputs = None +# ---- Cells without target signal ---- +# The loss is a per-cell correlation and the targets are row-normalized: both are undefined for a constant target +# vector (e.g. a cell without any protein counts), and a single such cell turns the whole training into NaN. +keep = informative_cells(train_targets) +if not keep.all(): + print(f"Dropping {(~keep).sum()} training cells with a constant target vector", flush=True) + train_inputs, train_targets, train_batches = train_inputs[keep], train_targets[keep], train_batches[keep] + +# ---- Metadata: the derived files of the original pipeline, rebuilt from the h5ad files ---- +print("Computing cell statistics...", flush=True) +train_cell_statistics = compute_cell_statistics(train_inputs, task_type) +test_cell_statistics = compute_cell_statistics(test_inputs, task_type) if test_inputs is not None else None +# the original standardized the statistics over train and test cells together +all_cell_statistics = pd.concat([train_cell_statistics, test_cell_statistics], ignore_index=True) +cell_statistics_standardization = fit_standardization(all_cell_statistics) +train_cell_statistics = apply_standardization(train_cell_statistics, cell_statistics_standardization) + +batch_singular_vectors = None +if task_type == "cite": + print("Computing batch singular vectors...", flush=True) + train_batch_medians = compute_batch_input_medians(train_inputs, train_batches) + batch_singular_vectors = BatchSingularVectors().fit(train_batch_medians) + if test_inputs is not None: + # batches of the test set that were not seen in training are projected with the fitted components; for + # batches present in both, the training cells (many more of them) define the statistics + new_batches = np.setdiff1d(np.unique(test_batches), train_batch_medians.index) + if len(new_batches): + in_new_batch = np.isin(test_batches, new_batches) + test_batch_medians = compute_batch_input_medians(test_inputs[in_new_batch], test_batches[in_new_batch]) + batch_singular_vectors.table = pd.concat([batch_singular_vectors.table, batch_singular_vectors.transform(test_batch_medians)]) + +metadata_kwargs = {"task_type": task_type, "day_pattern": par["day_pattern"], "donor_pattern": par["donor_pattern"]} +train_metadata = build_metadata( + train_batches, + train_cell_statistics, + batch_sv=batch_singular_vectors.lookup(train_batches) if batch_singular_vectors is not None else None, + **metadata_kwargs, +) test_metadata = None -if par.get('input_test_mod1'): - print('Loading test data for SVD fitting...', flush=True) - input_test_mod1 = ad.read_h5ad(par['input_test_mod1']) - test_inputs = to_sparse_csr(input_test_mod1.layers['normalized']) +if test_inputs is not None: test_metadata = build_metadata( - input_test_mod1, - task_type, - cell_type_col=par['cell_type_col'], - day_pattern=par['day_pattern'], + test_batches, + apply_standardization(test_cell_statistics, cell_statistics_standardization), + batch_sv=batch_singular_vectors.lookup(test_batches) if batch_singular_vectors is not None else None, + **metadata_kwargs, ) - del input_test_mod1 - gc.collect() -# ---- Create data_dir with CITE-specific mask files ---- -# The original PrePostProcessing loads pre-computed feature-target correlation -# masks from data_dir. We replace them with all-True masks so all input features -# are retained as supplementary raw features alongside the SVD components. +# ---- CITE input masks: the genes whose raw expression is appended to the SVD components ---- data_dir = tempfile.mkdtemp() -if task_type == 'cite': - mask_pair = np.ones((n_vars_mod1, n_vars_mod2), dtype=bool) - np.savez(os.path.join(data_dir, 'cite_inputs_targets_pair3g.npz'), mask=mask_pair) - mask2 = np.zeros((n_vars_mod1,), dtype=bool) - np.savez(os.path.join(data_dir, 'cite_inputs_mask2.npz'), mask=mask2) - -# ---- Get parameters ---- -pre_post_process_params = PrePostProcessing.get_params( - task_type=task_type, - data_dir=data_dir, - device=device, - seed=SEED, -) -model_params = EncoderDecoder.get_params( - task_type=task_type, - device=device, -) -model_params['epoch'] = par['n_epochs'] -model_params['burnin_length_epoch'] = par['burnin_length_epoch'] - -# ---- Fit preprocessing ---- -print('Fitting preprocessing...', flush=True) -pre_post_process = PrePostProcessing(pre_post_process_params) +if task_type == "cite": + print("Building the CITE input gene masks...", flush=True) + hgnc_path, reactome_path = par.get("hgnc_complete_set"), par.get("reactome_pathways") + if not (hgnc_path and reactome_path): + # the docker image ships the files in SS_OPM_REFERENCE_DIR; download them when running elsewhere + hgnc_path, reactome_path = download_reference_files(os.environ.get("SS_OPM_REFERENCE_DIR", os.path.join(data_dir, "reference"))) + hgnc = read_hgnc(hgnc_path) + if mod1 == "GEX": + gene_symbols = gene_symbols_from_var_names(mod1_var_names, hgnc) + protein_symbols = gene_symbols_from_var_names(mod2_var.index.to_numpy(), hgnc) + targets_gene2idx = make_targets_gene2idx(protein_symbols, hgnc) + inputs_lognorm = to_sparse_csr( + np.vstack( + [median_normalized_log_expression(train_inputs[start : start + ROW_BLOCK]) for start in range(0, train_inputs.shape[0], ROW_BLOCK)] + ) + ) + pair_mask, pathway_mask = make_cite_input_masks( + inputs_lognorm, + row_normalize(to_dense(train_targets)), + gene_symbols, + targets_gene2idx, + train_batches, + read_reactome_gmt(reactome_path), + ) + del inputs_lognorm + print(f"{pair_mask.any(axis=1).sum()} genes paired with a protein, {pathway_mask.sum()} pathway genes selected", flush=True) + else: + # ADT -> GEX: the inputs are proteins, there is nothing to pair them with; keep the SVD components only + pair_mask = np.zeros((train_inputs.shape[1], train_targets.shape[1]), dtype=bool) + pathway_mask = np.zeros(train_inputs.shape[1], dtype=bool) + np.savez(os.path.join(data_dir, "cite_inputs_targets_pair3g.npz"), mask=pair_mask) + np.savez(os.path.join(data_dir, "cite_inputs_mask2.npz"), mask=pathway_mask) + gc.collect() -# Use test inputs alongside train inputs for fitting SVD (improves coverage) -_test_inputs_for_svd = test_inputs if test_inputs is not None else train_inputs -_test_metadata_for_svd = test_metadata if test_metadata is not None else train_metadata +# ---- Parameters: the authors' defaults ---- +pre_post_process_params = PrePostProcessing.get_params(task_type=task_type, data_dir=data_dir, device=device, seed=SEED) +model_params = EncoderDecoder.get_params(task_type=task_type, device=device) +model_params["epoch"] = par["n_epochs"] +model_params["burnin_length_epoch"] = par["burnin_length_epoch"] +# ---- Fit preprocessing (the SVDs are fit on train and test inputs together, as in the original) ---- +print("Fitting preprocessing...", flush=True) +pre_post_process = PrePostProcessing(pre_post_process_params) pre_post_process.fit_preprocess( inputs_values=train_inputs, targets_values=train_targets, metadata=train_metadata, - test_inputs_values=_test_inputs_for_svd, - test_metadata=_test_metadata_for_svd, + test_inputs_values=test_inputs if test_inputs is not None else train_inputs, + test_metadata=test_metadata if test_metadata is not None else train_metadata, ) -# ---- Preprocess training data ---- -print('Preprocessing training data...', flush=True) +print("Preprocessing training data...", flush=True) preprocessed_inputs, preprocessed_targets = pre_post_process.preprocess( - inputs_values=train_inputs, - targets_values=train_targets, - metadata=train_metadata, + inputs_values=train_inputs, targets_values=train_targets, metadata=train_metadata ) +preprocessed_inputs = np.asarray(preprocessed_inputs, dtype=np.float32) +preprocessed_targets = np.asarray(preprocessed_targets, dtype=np.float32) +print(f"model input shape X:{preprocessed_inputs.shape} Y:{preprocessed_targets.shape}", flush=True) +gc.collect() -# ---- Train model ---- -# Cast preprocessed arrays to float32 to match what PyTorch expects. -if isinstance(preprocessed_inputs, np.ndarray): - preprocessed_inputs = preprocessed_inputs.astype(np.float32) -if isinstance(preprocessed_targets, np.ndarray): - preprocessed_targets = preprocessed_targets.astype(np.float32) - -print('Training model...', flush=True) +# ---- Train ---- +print("Training model...", flush=True) model = EncoderDecoder(model_params) model.fit( x=train_inputs, @@ -173,20 +229,42 @@ def _patched_train_step_fwd(self, batch, training_length_ratio): ) gc.collect() -# ---- Save model and preprocessing artifacts ---- -print('Saving model...', flush=True) -os.makedirs(par['output'], exist_ok=True) +# ---- Prediction scale ---- +# The network outputs per-cell z-scores (the competition only scored per-cell correlations). Fit one global affine +# map from those to the normalized targets on a subsample of training cells; it leaves every correlation unchanged. +rng = np.random.default_rng(SEED) +n_rescaling_cells = min(par["n_rescaling_cells"], train_inputs.shape[0]) +rescaling_cells = np.sort(rng.choice(train_inputs.shape[0], size=n_rescaling_cells, replace=False)) +train_predictions = model.predict( + x=train_inputs[rescaling_cells], + preprocessed_x=preprocessed_inputs[rescaling_cells], + metadata=train_metadata.iloc[rescaling_cells].reset_index(drop=True), +) +rescaling = fit_prediction_rescaling(train_predictions, train_targets[rescaling_cells]) +print(f"Prediction rescaling: slope {rescaling['slope']:.4f}, intercept {rescaling['intercept']:.4f}", flush=True) -model_dir = os.path.join(par['output'], 'model') +# ---- Save ---- +print("Saving model...", flush=True) +os.makedirs(par["output"], exist_ok=True) +model_dir = os.path.join(par["output"], "model") os.makedirs(model_dir, exist_ok=True) model.save(model_dir) - -with open(os.path.join(par['output'], 'pre_post_process.pickle'), 'wb') as f: - pickle.dump(pre_post_process, f) - -mod2_var.to_parquet(os.path.join(par['output'], 'mod2_var.parquet')) - -with open(os.path.join(par['output'], 'task_info.pickle'), 'wb') as f: - pickle.dump({'task_type': task_type, 'mod2': mod2, 'dataset_id': dataset_id}, f) - -print('Done!', flush=True) +with open(os.path.join(par["output"], "pre_post_process.pickle"), "wb") as handle: + pickle.dump(pre_post_process, handle) +with open(os.path.join(par["output"], "batch_singular_vectors.pickle"), "wb") as handle: + pickle.dump(batch_singular_vectors, handle) +mod2_var.to_parquet(os.path.join(par["output"], "mod2_var.parquet")) +save_json( + os.path.join(par["output"], "task_info.json"), + { + "task_type": task_type, + "mod1": mod1, + "mod2": mod2, + "dataset_id": dataset_id, + "day_pattern": par["day_pattern"], + "donor_pattern": par["donor_pattern"], + "cell_statistics_standardization": cell_statistics_standardization, + "prediction_rescaling": rescaling, + }, +) +print("Done!", flush=True) From 464d8c73a73174b6db59390a6c3f0a24e57351c9 Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Sun, 20 Sep 2026 14:21:34 +0200 Subject: [PATCH 2/3] ss_opm: rank-based Spearman for the gene masks, read old model bundles The mask construction correlates every candidate gene with all proteins at once via rankdata instead of a full spearmanr matrix per gene (identical values, minutes instead of hours on 14k genes). Bundles written before the metadata rebuild are read with neutral defaults so the pre-trained test resources keep working until they are regenerated. Co-Authored-By: Claude Fable 5.1 --- src/methods/ss_opm/ss_opm_common.py | 50 ++++++++++++++++++--- src/methods/ss_opm/ss_opm_predict/script.py | 6 +-- 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/src/methods/ss_opm/ss_opm_common.py b/src/methods/ss_opm/ss_opm_common.py index 4e41e02d..eff01575 100644 --- a/src/methods/ss_opm/ss_opm_common.py +++ b/src/methods/ss_opm/ss_opm_common.py @@ -140,14 +140,24 @@ def make_targets_gene2idx(target_symbols, hgnc): def _group_spearman(inputs_column, targets_columns): """Spearman correlation and p-value of one gene with several targets, over the cells where the gene is - expressed (as in the original scripts). Returns arrays of NaN when there are too few such cells.""" + expressed (as in the original scripts, which called `scipy.stats.spearmanr` on that subset). Only the gene's + row of the correlation matrix is computed; the p-value is spearmanr's two-sided t-test with n - 2 degrees + of freedom. NaN when there are too few expressing cells.""" expressed = inputs_column > 0 + n_cells = int(expressed.sum()) n_targets = targets_columns.shape[1] - if expressed.sum() < 3: + if n_cells < 3: return np.full(n_targets, np.nan), np.full(n_targets, np.nan) - result = scipy.stats.spearmanr(inputs_column[expressed], targets_columns[expressed]) - correlations = np.atleast_2d(result.statistic)[0, 1:] - p_values = np.atleast_2d(result.pvalue)[0, 1:] + input_ranks = scipy.stats.rankdata(inputs_column[expressed]) + target_ranks = scipy.stats.rankdata(targets_columns[expressed], axis=0) + input_centered = input_ranks - input_ranks.mean() + target_centered = target_ranks - target_ranks.mean(axis=0) + with np.errstate(invalid="ignore", divide="ignore"): + correlations = (input_centered @ target_centered) / np.sqrt( + (input_centered**2).sum() * (target_centered**2).sum(axis=0) + ) + t_statistics = correlations * np.sqrt((n_cells - 2) / (1 - correlations**2)) + p_values = 2 * scipy.stats.t.sf(np.abs(t_statistics), n_cells - 2) return correlations, p_values @@ -463,6 +473,36 @@ def load_json(path): return json.load(handle) +def identity_standardization(): + return {"mean": {key: 0.0 for key in CELL_STATISTIC_KEYS}, "std": {key: 1.0 for key in CELL_STATISTIC_KEYS}} + + +def load_model_bundle(model_dir): + """Task info and batch singular vectors of a trained bundle. + + Bundles written before the metadata rebuild (`task_info.pickle`, no standardization, no rescaling, no batch + singular vectors) are still readable, with neutral defaults, so that models pre-trained for the component tests + keep working until they are regenerated. + """ + import pickle + + json_path = os.path.join(model_dir, "task_info.json") + if os.path.exists(json_path): + task_info = load_json(json_path) + with open(os.path.join(model_dir, "batch_singular_vectors.pickle"), "rb") as handle: + batch_singular_vectors = pickle.load(handle) + return task_info, batch_singular_vectors + + with open(os.path.join(model_dir, "task_info.pickle"), "rb") as handle: + task_info = pickle.load(handle) + task_info.setdefault("mod1", "unknown") + task_info.setdefault("day_pattern", DEFAULT_DAY_PATTERN) + task_info.setdefault("donor_pattern", DEFAULT_DONOR_PATTERN) + task_info.setdefault("cell_statistics_standardization", identity_standardization()) + task_info.setdefault("prediction_rescaling", {"slope": 1.0, "intercept": 0.0}) + return task_info, None + + # --------------------------------------------------------------------------------------------------------------- # Runtime patches of the ss_opm package # --------------------------------------------------------------------------------------------------------------- diff --git a/src/methods/ss_opm/ss_opm_predict/script.py b/src/methods/ss_opm/ss_opm_predict/script.py index df6cc886..5c42909e 100644 --- a/src/methods/ss_opm/ss_opm_predict/script.py +++ b/src/methods/ss_opm/ss_opm_predict/script.py @@ -23,7 +23,7 @@ apply_standardization, build_metadata, compute_cell_statistics, - load_json, + load_model_bundle, to_sparse_csr, ) @@ -35,14 +35,12 @@ print(f"Using device: {device}", flush=True) # ---- Load the training bundle ---- -task_info = load_json(os.path.join(par["input_model"], "task_info.json")) +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) with open(os.path.join(par["input_model"], "pre_post_process.pickle"), "rb") as handle: pre_post_process = pickle.load(handle) -with open(os.path.join(par["input_model"], "batch_singular_vectors.pickle"), "rb") as handle: - batch_singular_vectors = pickle.load(handle) mod2_var = pd.read_parquet(os.path.join(par["input_model"], "mod2_var.parquet")) model = EncoderDecoder(params=None) From 3225cadaea599f4ba3e500b3994c7edf96281e43 Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Wed, 23 Sep 2026 17:37:25 +0200 Subject: [PATCH 3/3] ss_opm: download Reactome with a browser user agent; fit the input SVD on train only without a test set reactome.org answers HTTP 403 to Python's default urllib user agent, which failed the docker image build of ss_opm_train in CI (PR #69). Both the image-build step and the runtime fallback now send `User-Agent: Mozilla/5.0`. `fit_preprocess` no longer receives the training set again in place of a missing test set: the original's `use_test_inputs` switch is set from whether `--input_test_mod1` was given (review comment on #69). CHANGELOG entry added. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 2 ++ src/methods/ss_opm/ss_opm_common.py | 18 ++++++++++++++---- .../ss_opm/ss_opm_train/config.vsh.yaml | 10 +++++++--- src/methods/ss_opm/ss_opm_train/script.py | 9 ++++++--- 4 files changed, 29 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fad70d15..3a23c2d4 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/methods/ss_opm/ss_opm_common.py b/src/methods/ss_opm/ss_opm_common.py index eff01575..822b74c6 100644 --- a/src/methods/ss_opm/ss_opm_common.py +++ b/src/methods/ss_opm/ss_opm_common.py @@ -11,8 +11,9 @@ import json import os import re +import shutil +import urllib.request import zipfile -from urllib.request import urlretrieve import numpy as np import pandas as pd @@ -40,6 +41,8 @@ # Reference files the CITE gene masks are built from (same sources as the original `make_cite_input_mask.py`) HGNC_URL = "https://storage.googleapis.com/public-download-files/hgnc/archive/archive/monthly/tsv/hgnc_complete_set_2023-01-01.txt" REACTOME_URL = "https://reactome.org/download/current/ReactomePathways.gmt.zip" +# reactome.org answers HTTP 403 to Python's default user agent +DOWNLOAD_USER_AGENT = "Mozilla/5.0" ENSEMBL_ID_PATTERN = re.compile(r"^ENSG\d+") @@ -58,16 +61,23 @@ def to_dense(X, dtype=np.float32): # --------------------------------------------------------------------------------------------------------------- # Reference files # --------------------------------------------------------------------------------------------------------------- +def download_file(url, path): + request = urllib.request.Request(url, headers={"User-Agent": DOWNLOAD_USER_AGENT}) + with urllib.request.urlopen(request) as response, open(path, "wb") as handle: + shutil.copyfileobj(response, handle) + + def download_reference_files(directory): - """Download the HGNC complete set and the Reactome gene sets into `directory` (used at image build time).""" + """Download the HGNC complete set and the Reactome gene sets into `directory` unless they are already there + (the docker image ships them in SS_OPM_REFERENCE_DIR, see the train component's setup).""" os.makedirs(directory, exist_ok=True) hgnc_path = os.path.join(directory, "hgnc_complete_set.txt") reactome_path = os.path.join(directory, "ReactomePathways.gmt") if not os.path.exists(hgnc_path): - urlretrieve(HGNC_URL, hgnc_path) + download_file(HGNC_URL, hgnc_path) if not os.path.exists(reactome_path): archive_path = os.path.join(directory, "ReactomePathways.gmt.zip") - urlretrieve(REACTOME_URL, archive_path) + download_file(REACTOME_URL, archive_path) with zipfile.ZipFile(archive_path) as archive: archive.extractall(directory) os.remove(archive_path) diff --git a/src/methods/ss_opm/ss_opm_train/config.vsh.yaml b/src/methods/ss_opm/ss_opm_train/config.vsh.yaml index bcb1615e..c7b15269 100644 --- a/src/methods/ss_opm/ss_opm_train/config.vsh.yaml +++ b/src/methods/ss_opm/ss_opm_train/config.vsh.yaml @@ -62,12 +62,16 @@ engines: packages: - pyarrow - fastparquet - # reference files for the CITE input gene masks (same sources as the original solution) + # 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 urllib.request; urllib.request.urlretrieve('https://storage.googleapis.com/public-download-files/hgnc/archive/archive/monthly/tsv/hgnc_complete_set_2023-01-01.txt', 'hgnc_complete_set.txt')" && \ - python -c "import urllib.request, zipfile; urllib.request.urlretrieve('https://reactome.org/download/current/ReactomePathways.gmt.zip', 'ReactomePathways.gmt.zip'); zipfile.ZipFile('ReactomePathways.gmt.zip').extractall('.')" && \ + 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: diff --git a/src/methods/ss_opm/ss_opm_train/script.py b/src/methods/ss_opm/ss_opm_train/script.py index 69883840..7aa8e3f8 100644 --- a/src/methods/ss_opm/ss_opm_train/script.py +++ b/src/methods/ss_opm/ss_opm_train/script.py @@ -192,19 +192,22 @@ def _float32_train_step_forward(self, batch, training_length_ratio): # ---- Parameters: the authors' defaults ---- pre_post_process_params = PrePostProcessing.get_params(task_type=task_type, data_dir=data_dir, device=device, seed=SEED) +# The original fits the input SVD on the training and test inputs together (transductive). The test set is an optional +# input of this component: without it, the original's own switch makes the SVD fit on the training cells only. +pre_post_process_params["use_test_inputs"] = test_inputs is not None model_params = EncoderDecoder.get_params(task_type=task_type, device=device) model_params["epoch"] = par["n_epochs"] model_params["burnin_length_epoch"] = par["burnin_length_epoch"] -# ---- Fit preprocessing (the SVDs are fit on train and test inputs together, as in the original) ---- +# ---- Fit preprocessing ---- print("Fitting preprocessing...", flush=True) pre_post_process = PrePostProcessing(pre_post_process_params) pre_post_process.fit_preprocess( inputs_values=train_inputs, targets_values=train_targets, metadata=train_metadata, - test_inputs_values=test_inputs if test_inputs is not None else train_inputs, - test_metadata=test_metadata if test_metadata is not None else train_metadata, + test_inputs_values=test_inputs, + test_metadata=test_metadata, ) print("Preprocessing training data...", flush=True)