From 51240b43e7e019716ed957638de410cbdd60b70a Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Mon, 5 Oct 2026 03:02:58 +0200 Subject: [PATCH 1/8] Restore the per-cell scale of senkin_tmp and ss_opm predictions Both 2022 competition solutions predict a z-scored profile per cell (the competition scored per-cell Pearson only). Mapping that back with one global slope and intercept gave every predicted cell the same mean and standard deviation, so per-gene correlations, overall correlations and RMSE/MAE scored the missing per-cell level rather than the methods. src/utils/cell_scale.py predicts each cell's target mean and log standard deviation from its input (ridge regression on a truncated SVD of the normalized input plus two sequencing-depth features) and restores them, with the spread shrunk by one least-squares factor fitted on training-cell predictions (out-of-fold for senkin_tmp). Per-cell correlations are unchanged. ss_opm_predict keeps the old global map for model bundles trained before this change. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 + .../senkin_tmp/senkin_tmp/config.vsh.yaml | 4 +- .../senkin_tmp_train/config.vsh.yaml | 1 + .../senkin_tmp/senkin_tmp_train/script.py | 17 ++- src/methods/ss_opm/ss_opm/config.vsh.yaml | 5 +- src/methods/ss_opm/ss_opm_common.py | 7 - .../ss_opm/ss_opm_predict/config.vsh.yaml | 1 + src/methods/ss_opm/ss_opm_predict/script.py | 14 +- .../ss_opm/ss_opm_train/config.vsh.yaml | 5 +- src/methods/ss_opm/ss_opm_train/script.py | 22 +++- src/utils/cell_scale.py | 123 ++++++++++++++++++ 11 files changed, 174 insertions(+), 27 deletions(-) create mode 100644 src/utils/cell_scale.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 92e3274b..831dc115 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -126,6 +126,8 @@ * `guanlab_dengkw_pm`: Solve the kernel ridge regression with `scipy.linalg.cho_factor()`/`cho_solve()` instead of `KernelRidge`, whose `scipy.linalg.solve()` raised a `MemoryError` or segfaulted beyond ~30k cells with the OpenBLAS in this image. Predictions are unchanged; the method failed on every dataset except `bmmc_multiome`. +* `senkin_tmp_train`, `ss_opm_train`, `ss_opm_predict`: Give every predicted cell its own level and spread instead of one global slope and intercept. Both 2022 competition solutions predict a z-scored profile per cell, so every cell came out with the same mean (0.5955 on `bmmc_cite`, where the true per-cell mean ranges from 0.16 to 2.83): per-cell correlations were unaffected, but per-gene correlations, overall correlations and RMSE/MAE scored that missing level, not the methods. The new `src/utils/cell_scale.py` predicts each cell's target mean and log standard deviation from its input (ridge regression on a truncated SVD plus two sequencing-depth features) and shrinks the restored spread by one factor fitted on training-cell predictions. `ss_opm_predict` still applies the old global map to model bundles trained before this change. + # task_predict_modality 0.1.1 ## NEW FUNCTIONALITY diff --git a/src/methods/senkin_tmp/senkin_tmp/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp/config.vsh.yaml index 9375a36d..86832777 100644 --- a/src/methods/senkin_tmp/senkin_tmp/config.vsh.yaml +++ b/src/methods/senkin_tmp/senkin_tmp/config.vsh.yaml @@ -11,7 +11,9 @@ description: | the TSVD-reduced LightGBM predictions, together with the feature blocks (each z-scored per cell), with two neural networks: a bidirectional GRU with cosine-similarity loss and a dense bidirectional GRU with MSE loss. Final predictions are a weighted blend (55% cosine, 45% MSE) - of per-fold averaged outputs, rescaled to the target scale with a global affine transform. + of per-fold averaged outputs. These are z-scored per cell (the competition scored per-cell + correlations only), so each test cell gets its own protein level and spread back, predicted + from its RNA profile by a ridge regression fitted on the training cells. references: doi: - 10.64898/2026.02.24.707614 diff --git a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml index ef7ec762..87f24bbd 100644 --- a/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml +++ b/src/methods/senkin_tmp/senkin_tmp_train/config.vsh.yaml @@ -13,6 +13,7 @@ resources: - path: script.py type: python_script - path: /src/utils/exit_codes.py + - path: /src/utils/cell_scale.py arguments: - name: "--n_folds" type: integer diff --git a/src/methods/senkin_tmp/senkin_tmp_train/script.py b/src/methods/senkin_tmp/senkin_tmp_train/script.py index efed5dc0..8404b4ab 100644 --- a/src/methods/senkin_tmp/senkin_tmp_train/script.py +++ b/src/methods/senkin_tmp/senkin_tmp_train/script.py @@ -43,6 +43,7 @@ ## VIASH END sys.path.append(meta["resources_dir"]) +from cell_scale import CellScale from exit_codes import exit_non_applicable # The original solution corrected batch effects per day and computed gene-protein correlations per donor and day. @@ -240,12 +241,16 @@ def _nn_inputs(idx): test_preds = zscore(test_preds_cos) * 0.55 + zscore(test_preds_mse) * 0.45 # The original solution was scored with a per-cell Pearson correlation only, so its predictions are z-scored per -# cell. The benchmark also computes RMSE/MAE, so bring the predictions back to the scale of the normalized proteins -# with a single global affine transform fitted on the out-of-fold training predictions. One (slope, intercept) pair -# for all cells and proteins leaves every per-cell and per-protein correlation untouched. -slope, intercept = np.polyfit(train_preds.ravel(), Y_prot_train.ravel(), deg=1) -logger.info(f"Rescaling z-scored predictions to the target scale: slope {slope:.4f}, intercept {intercept:.4f}") -test_preds = test_preds * slope + intercept +# cell. Give every test cell its own protein level and spread back, predicted from its RNA profile (see +# src/utils/cell_scale.py); the shrinkage of the spread is fitted on the out-of-fold training predictions. Per-cell +# correlations are unchanged; per-protein correlations and RMSE/MAE now see each cell's level. +def _counts(adata): + return adata.layers["counts"] if "counts" in adata.layers else None + +cell_scale = CellScale().fit(adata_rna_train.layers["normalized"], Y_prot_train, _counts(adata_rna_train)) +cell_scale.fit_spread(train_preds, adata_rna_train.layers["normalized"], Y_prot_train, _counts(adata_rna_train)) +logger.info(f"Restoring the per-cell protein scale (spread shrinkage {cell_scale.spread_:.3f})") +test_preds = cell_scale.apply(test_preds, adata_rna_test.layers["normalized"], _counts(adata_rna_test)) # --------------------------------------------------------------------------- # Save the model: the solution is transductive, so the test predictions themselves are the model. They are stored as an diff --git a/src/methods/ss_opm/ss_opm/config.vsh.yaml b/src/methods/ss_opm/ss_opm/config.vsh.yaml index 018c91ca..74e1e369 100644 --- a/src/methods/ss_opm/ss_opm/config.vsh.yaml +++ b/src/methods/ss_opm/ss_opm/config.vsh.yaml @@ -10,8 +10,9 @@ description: | 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. + task's files. The output is z-scored per cell (the competition scored per-cell correlations + only), so each test cell gets its own target level and spread back, predicted from its input + profile by a ridge regression fitted on the training cells. 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 822b74c6..614063df 100644 --- a/src/methods/ss_opm/ss_opm_common.py +++ b/src/methods/ss_opm/ss_opm_common.py @@ -466,13 +466,6 @@ def informative_cells(targets): 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) 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 112c6fd4..c93b4f05 100644 --- a/src/methods/ss_opm/ss_opm_predict/config.vsh.yaml +++ b/src/methods/ss_opm/ss_opm_predict/config.vsh.yaml @@ -9,6 +9,7 @@ resources: - type: python_script path: script.py - path: ../ss_opm_common.py + - path: /src/utils/cell_scale.py engines: - type: docker image: openproblems/base_pytorch_nvidia:1 diff --git a/src/methods/ss_opm/ss_opm_predict/script.py b/src/methods/ss_opm/ss_opm_predict/script.py index 5c42909e..bccadc2e 100644 --- a/src/methods/ss_opm/ss_opm_predict/script.py +++ b/src/methods/ss_opm/ss_opm_predict/script.py @@ -18,6 +18,7 @@ ## VIASH END sys.path.append(meta["resources_dir"]) +from cell_scale import CellScale # noqa: E402 from ss_opm_common import ( # noqa: E402 apply_runtime_patches, apply_standardization, @@ -77,9 +78,16 @@ 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"]) +# The network outputs per-cell z-scores, so a non-finite value is set to 0, the cell's mean. +predictions = np.nan_to_num(predictions, nan=0.0, posinf=0.0, neginf=0.0) +cell_scale_path = os.path.join(par["input_model"], "cell_scale.npz") +if os.path.exists(cell_scale_path): + test_counts = to_sparse_csr(input_test_mod1.layers["counts"]) if "counts" in input_test_mod1.layers else None + predictions = CellScale.load(cell_scale_path).apply(predictions, test_inputs, test_counts) +else: + # models trained before the per-cell scale was introduced store one global affine map + rescaling = task_info["prediction_rescaling"] + predictions = predictions * rescaling["slope"] + rescaling["intercept"] assert predictions.shape == (input_test_mod1.n_obs, mod2_var.shape[0]) # ---- Write ---- 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 c7b15269..2efde543 100644 --- a/src/methods/ss_opm/ss_opm_train/config.vsh.yaml +++ b/src/methods/ss_opm/ss_opm_train/config.vsh.yaml @@ -46,12 +46,13 @@ arguments: 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. + Number of training cells whose predictions are used to fit the shrinkage of the restored + per-cell spread (see src/utils/cell_scale.py). resources: - type: python_script path: script.py - path: ../ss_opm_common.py + - path: /src/utils/cell_scale.py engines: - type: docker image: openproblems/base_pytorch_nvidia:1 diff --git a/src/methods/ss_opm/ss_opm_train/script.py b/src/methods/ss_opm/ss_opm_train/script.py index 7aa8e3f8..bb8431aa 100644 --- a/src/methods/ss_opm/ss_opm_train/script.py +++ b/src/methods/ss_opm/ss_opm_train/script.py @@ -27,6 +27,7 @@ ## VIASH END sys.path.append(meta["resources_dir"]) +from cell_scale import CellScale # noqa: E402 from ss_opm_common import ( # noqa: E402 ROW_BLOCK, BatchSingularVectors, @@ -36,7 +37,6 @@ compute_batch_input_medians, compute_cell_statistics, download_reference_files, - fit_prediction_rescaling, fit_standardization, gene_symbols_from_var_names, informative_cells, @@ -98,6 +98,8 @@ def _float32_train_step_forward(self, batch, training_length_ratio): print(f"Task type: {task_type}", flush=True) train_inputs = to_sparse_csr(input_train_mod1.layers["normalized"]).astype(np.float32) +# raw input counts, only for the sequencing-depth features of the per-cell scale model +train_counts = to_sparse_csr(input_train_mod1.layers["counts"]) if "counts" in input_train_mod1.layers else None 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 @@ -114,6 +116,7 @@ def _float32_train_step_forward(self, batch, training_length_ratio): 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] + train_counts = train_counts[keep] if train_counts is not None else None # ---- Metadata: the derived files of the original pipeline, rebuilt from the h5ad files ---- print("Computing cell statistics...", flush=True) @@ -233,8 +236,9 @@ def _float32_train_step_forward(self, batch, training_length_ratio): gc.collect() # ---- 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. +# The network outputs per-cell z-scores (the competition only scored per-cell correlations). Every cell gets its own +# target level and spread back, predicted from its input profile (src/utils/cell_scale.py); the shrinkage of the spread +# is fitted on predictions for a subsample of training cells. Per-cell correlations are 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)) @@ -243,8 +247,14 @@ def _float32_train_step_forward(self, batch, training_length_ratio): 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) +cell_scale = CellScale().fit(train_inputs, train_targets, train_counts) +cell_scale.fit_spread( + train_predictions, + train_inputs[rescaling_cells], + train_targets[rescaling_cells], + train_counts[rescaling_cells] if train_counts is not None else None, +) +print(f"Per-cell scale: spread shrinkage {cell_scale.spread_:.4f}", flush=True) # ---- Save ---- print("Saving model...", flush=True) @@ -257,6 +267,7 @@ def _float32_train_step_forward(self, batch, training_length_ratio): 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")) +cell_scale.save(os.path.join(par["output"], "cell_scale.npz")) save_json( os.path.join(par["output"], "task_info.json"), { @@ -267,7 +278,6 @@ def _float32_train_step_forward(self, batch, training_length_ratio): "day_pattern": par["day_pattern"], "donor_pattern": par["donor_pattern"], "cell_statistics_standardization": cell_statistics_standardization, - "prediction_rescaling": rescaling, }, ) print("Done!", flush=True) diff --git a/src/utils/cell_scale.py b/src/utils/cell_scale.py new file mode 100644 index 00000000..b3abaf84 --- /dev/null +++ b/src/utils/cell_scale.py @@ -0,0 +1,123 @@ +"""Restore each cell's level and spread for predictions that only get the per-cell profile right. + +Methods written for the NeurIPS 2022 Kaggle competition, which scored nothing but the per-cell Pearson correlation, +predict a z-scored profile for each cell. Mapped back with one global slope and intercept, every predicted cell ends up +with the same mean and standard deviation, although the true per-cell mean varies several-fold between cells; per-gene +correlations, overall correlations and RMSE/MAE then score that missing information, not the method. + +`CellScale` learns each training cell's target mean and log standard deviation from its input profile (a ridge +regression on a truncated SVD of the normalized input plus two sequencing-depth features) and gives every standardized +prediction its own mean back, and its own standard deviation times one shrinkage factor. The factor is fitted by least +squares on predictions for training cells (out-of-fold where the method has them): where the predicted profile +correlates only weakly with the truth, the error-minimising spread is smaller than the cell's standard deviation. +Per-cell correlations are unchanged by construction. + +The fitted model is a handful of arrays, saved with `numpy.savez`, so that it can be loaded in another image without +unpickling library objects. +""" + +import numpy as np +import scipy.sparse as sp +from sklearn.decomposition import TruncatedSVD +from sklearn.linear_model import RidgeCV + +RIDGE_ALPHAS = np.logspace(-2, 4, 13) +EPSILON = 1e-3 # keeps log(sd) finite for constant target vectors + + +def row_mean_std(matrix): + """Mean and standard deviation of every row of a dense or sparse matrix.""" + if sp.issparse(matrix): + matrix = sp.csr_matrix(matrix, dtype=np.float64) + mean = np.asarray(matrix.mean(axis=1)).ravel() + mean_of_squares = np.asarray(matrix.multiply(matrix).mean(axis=1)).ravel() + else: + matrix = np.asarray(matrix, dtype=np.float64) + mean = matrix.mean(axis=1) + mean_of_squares = (matrix**2).mean(axis=1) + return mean, np.sqrt(np.maximum(mean_of_squares - mean**2, 0)) + + +def standardize_rows(matrix): + """Z-score every row; a constant row becomes all zeros.""" + matrix = np.asarray(matrix, dtype=np.float64) + mean, std = row_mean_std(matrix) + return (matrix - mean[:, None]) / np.where(std > 0, std, 1)[:, None] + + +def depth_features(counts): + """log1p of the total and of the number of detected features per cell.""" + counts = sp.csr_matrix(counts) + total = np.asarray(counts.sum(axis=1)).ravel() + detected = np.diff(counts.indptr) + return np.column_stack([np.log1p(np.maximum(total, 0)), np.log1p(detected)]) + + +class CellScale: + """Predict each cell's target mean and standard deviation from its input, and apply them to standardized predictions. + + `inputs` is the normalized input modality (cells x input features), `counts` its raw counts (or the normalized + values when there are none; only used for the two depth features), `targets` the normalized target modality. + """ + + def __init__(self, n_components=64, random_state=0): + self.n_components = n_components + self.random_state = random_state + self.spread_ = 1.0 + + def _features(self, inputs, counts): + inputs = sp.csr_matrix(inputs, dtype=np.float32) + projected = np.asarray(inputs @ self.components_.T) + features = np.hstack([projected, depth_features(inputs if counts is None else counts)]) + return (features - self.feature_mean_) / self.feature_std_ + + def fit(self, inputs, targets, counts=None): + inputs = sp.csr_matrix(inputs, dtype=np.float32) + n_components = min(self.n_components, min(inputs.shape) - 1) + self.components_ = TruncatedSVD(n_components=n_components, random_state=self.random_state).fit(inputs).components_ + raw_features = np.hstack([np.asarray(inputs @ self.components_.T), depth_features(inputs if counts is None else counts)]) + self.feature_mean_ = raw_features.mean(axis=0) + self.feature_std_ = np.where(raw_features.std(axis=0) > 0, raw_features.std(axis=0), 1) + target_mean, target_std = row_mean_std(targets) + ridge = RidgeCV(alphas=RIDGE_ALPHAS).fit( + (raw_features - self.feature_mean_) / self.feature_std_, + np.column_stack([target_mean, np.log(target_std + EPSILON)]), + ) + self.coef_, self.intercept_ = ridge.coef_, ridge.intercept_ + return self + + def predict(self, inputs, counts=None): + """Predicted target mean and standard deviation of every cell.""" + fitted = self._features(inputs, counts) @ self.coef_.T + self.intercept_ + return fitted[:, 0], np.maximum(np.exp(fitted[:, 1]) - EPSILON, 0) + + def fit_spread(self, predictions, inputs, targets, counts=None, block_size=1000): + """Least-squares shrinkage of the predicted deviations from the cell mean, fitted on training cells.""" + mean, std = self.predict(inputs, counts) + targets = sp.csr_matrix(targets) if sp.issparse(targets) else np.asarray(targets) + covariance, variance = 0.0, 0.0 + for start in range(0, predictions.shape[0], block_size): + rows = slice(start, start + block_size) + deviations = standardize_rows(predictions[rows]) * std[rows, None] + block_targets = targets[rows].toarray() if sp.issparse(targets) else targets[rows] + covariance += float(((block_targets - mean[rows, None]) * deviations).sum()) + variance += float((deviations**2).sum()) + self.spread_ = covariance / variance if variance > 0 else 0.0 + return self + + def apply(self, predictions, inputs, counts=None): + """Give each cell of `predictions` (any per-cell scale) its predicted mean and shrunken standard deviation.""" + mean, std = self.predict(inputs, counts) + return standardize_rows(predictions) * (self.spread_ * std)[:, None] + mean[:, None] + + def save(self, path): + np.savez(path, components=self.components_, feature_mean=self.feature_mean_, feature_std=self.feature_std_, + coef=self.coef_, intercept=self.intercept_, spread=self.spread_) + + @classmethod + def load(cls, path): + arrays = np.load(path) + model = cls(n_components=arrays["components"].shape[0]) + model.components_, model.feature_mean_, model.feature_std_ = arrays["components"], arrays["feature_mean"], arrays["feature_std"] + model.coef_, model.intercept_, model.spread_ = arrays["coef"], arrays["intercept"], float(arrays["spread"]) + return model From 5ba37b562c130f82951e85c5da2b6d446bf02c5c Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Mon, 5 Oct 2026 07:26:45 +0200 Subject: [PATCH 2/8] scbutterfly: fix 2022 peak names, model every target feature, predict on the scored scale - chrom_utils.parse_chrom: split peak names on the first ':' or '-'. NeurIPS 2022 names peaks chr1:start-end, so every peak became its own chromosome with its own encoder block: pbmc_multiome was OOM-killed (ATAC -> GEX, 174,978 "chromosomes") or hit the 8 h limit (GEX -> ATAC, 7,002 "chromosomes") on rc7. - Keep scButterfly's HVG selection and peak filter for the input modality only, so every target gene/peak is modelled instead of zero-filled. - target_scale.py: convert GEX predictions exactly from scButterfly's median-library-size log1p to log CP10k; map ATAC predictions onto the dataset's TF-IDF with one least-squares line per peak (slope >= 0), fitted on the held-out validation cells. - Patch Model.__init__ to serve cells as dense rows of the sparse matrix instead of densifying both modalities (92 GB for the 2022 ATAC input). - Model bundles without metadata["model_all_target_features"] keep the old preprocessing. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 6 + src/methods/scbutterfly/butterfly_common.py | 127 +++++++++++++----- src/methods/scbutterfly/chrom_utils.py | 24 ++-- .../scbutterfly/scbutterfly/config.vsh.yaml | 8 +- .../scbutterfly_predict/config.vsh.yaml | 1 + .../scbutterfly/scbutterfly_predict/script.py | 33 +++-- .../scbutterfly_train/config.vsh.yaml | 2 +- .../scbutterfly/scbutterfly_train/script.py | 1 + src/methods/scbutterfly/target_scale.py | 62 +++++++++ 9 files changed, 202 insertions(+), 62 deletions(-) create mode 100644 src/methods/scbutterfly/target_scale.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 831dc115..9dc1b3fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -58,6 +58,12 @@ ## BUG FIXES +* `scbutterfly`: Parse the chromosome of `chr1:1000-2000` peak names, not only `chr1-1000-2000`. NeurIPS 2022 names its peaks the former way, so every peak became its own chromosome, each with its own block in the ATAC encoder: `pbmc_multiome` was OOM-killed when GEX was the target (174,978 "chromosomes") and hit the 8 h limit when ATAC was (7,002 "chromosomes", 5.7 h for one pretraining epoch). + +* `scbutterfly`: Model every target feature and predict it on the scale of the target's `normalized` layer. scButterfly kept only the 3,000 most variable genes and the peaks open in at least 0.5% of the cells, also when they were the target, so the other genes and peaks were predicted as zero; and it predicts its own preprocessing of the target. GEX predictions are now converted exactly from scButterfly's median-library-size scaling to log CP10k, and ATAC predictions (a sigmoid of the max-scaled TF-IDF) are mapped onto the dataset's TF-IDF by one least-squares line per peak, fitted on the training cells held out for validation. The HVG selection and peak filter still apply to the input modality. `scbutterfly_predict` reproduces the old preprocessing for model bundles trained before this change. + +* `scbutterfly`: Serve the model's training and test cells one sparse row at a time instead of densifying both whole modalities, 92 GB for the ATAC input of `pbmc_multiome` alone. + * `cellmapper_linear`: Write the unmasked variants as `mask_var: "none"` rather than `mask_var: null`. Viash drops null values from the config, so these variants fell back to the default `mask_var: "hvg"` and duplicated their `-hvg` counterparts. * `guanlab_dengkw_pm`: Restore the consensus scheme of the original submission -- five reshuffles of the batches into two halves, ten kernel ridge models averaged. The port had replaced it with a single fixed two-way split for ADT pairs and leave-one-batch-out otherwise, so the result depended on the order the batches happened to come in. `--n_repeats` and `--seed` are now arguments; the unused `--distance_method` and `--n_pcs` are gone (PR #32). diff --git a/src/methods/scbutterfly/butterfly_common.py b/src/methods/scbutterfly/butterfly_common.py index a21cccee..8bbbf619 100644 --- a/src/methods/scbutterfly/butterfly_common.py +++ b/src/methods/scbutterfly/butterfly_common.py @@ -30,6 +30,7 @@ def apply_runtime_patches(): - BCE on a float32 sigmoid can drift >1.0 -> clamp input to [0,1] - scButterfly hardcodes ``.cuda()``; make it a no-op when no GPU is present - ``TFIDF`` builds dense (n_peaks, n_cells) tiles -> sparse equivalent + - ``Model`` densifies both whole modalities -> densify one cell at a time """ import torch import torch.nn as nn @@ -78,6 +79,10 @@ def _safe_bce(input, target, *args, **kwargs): data_processing.TFIDF = _sparse_tfidf logger.info("Patched scButterfly TFIDF with the sparse implementation.") + from scButterfly import train_model + + train_model.Model.__init__ = _rows_on_demand(train_model.Model.__init__) + def _sparse_tfidf(count_mat): """Sparse, O(nnz) drop-in for ``scButterfly.data_processing.TFIDF``. @@ -120,15 +125,51 @@ def _sparse_tfidf(count_mat): return out, cell_totals, idf +class _DenseRows: + """A CSR matrix that hands out one cell at a time as a dense float32 row. + + ``Model.__init__`` stores ``RNA_data.X.toarray()`` and ``ATAC_data.X.toarray()``: + every cell of both modalities, training and test, as dense float32. At NeurIPS 2022 + Multiome scale (131k cells x 175k peaks after peak filtering) the ATAC matrix alone + is 92 GB. The model only ever reads ``.shape`` and single rows (``data[cell, :]`` in + its datasets), so this serves those rows from the sparse matrix instead. + """ + + def __init__(self, matrix): + self.matrix = csr_matrix(matrix, dtype=np.float32) + self.shape = self.matrix.shape + + def __getitem__(self, index): + cell = index[0] if isinstance(index, tuple) else index + start, end = self.matrix.indptr[cell], self.matrix.indptr[cell + 1] + row = np.zeros(self.shape[1], dtype=np.float32) + row[self.matrix.indices[start:end]] = self.matrix.data[start:end] + return row + + +def _rows_on_demand(original_init): + """Wrap ``Model.__init__`` to keep the data sparse (see :class:`_DenseRows`).""" + + def init(self, RNA_data, ATAC_data, *args, **kwargs): + # The networks are sized from the dim lists, not the data, so build them on + # zero cells and attach the real matrices afterwards. + original_init(self, RNA_data[:0], ATAC_data[:0], *args, **kwargs) + self.RNA_data_obs, self.ATAC_data_obs = RNA_data.obs, ATAC_data.obs + # Upstream stores the RNA var as ATAC_data_var too; nothing reads either. + self.RNA_data_var, self.ATAC_data_var = RNA_data.var, RNA_data.var + self.RNA_data, self.ATAC_data = _DenseRows(RNA_data.X), _DenseRows(ATAC_data.X) + + return init + + @contextlib.contextmanager def suppress_unused_postprocessing(): """Make ``sc.pp.pca``/``sc.pp.neighbors`` no-ops for the duration of the block. ``Model.test`` runs both, on *both* predicted matrices, whenever it is not asked to draw figures — and there is no flag to turn it off. Nothing reads the results - back: :func:`extract_predictions` only touches ``.X`` and ``.var_names``. On the - ATAC side it is a PCA over a near-dense ``n_test x n_peaks`` matrix, which - dominates the predict step. + back: :func:`predict_cells` only touches ``.X``. On the ATAC side it is a PCA over + a near-dense ``n_test x n_peaks`` matrix, which dominates the predict step. """ orig_pca, orig_neighbors = sc.pp.pca, sc.pp.neighbors sc.pp.pca = lambda *args, **kwargs: None @@ -210,13 +251,22 @@ def _concat_blocks(train_block, test_block): return out -def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly): +def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly, + model_all_target_features=True): """Build + preprocess + construct a Butterfly for the given data. + scButterfly's preprocessing keeps the ``n_top_genes`` most variable genes and drops + peaks open in fewer than 0.5% of the cells. With ``model_all_target_features`` that + only applies to the input modality: the benchmark scores every target feature, and + one the model drops could only be predicted as a constant zero. ``False`` reproduces + models trained before this option existed. + Returns a dict with the constructed ``butterfly``, the resolved ``direction``, - ``test_id`` (indices of the test cells), ``chrom_list``, and the preprocessed - target-modality var names. Deterministic given the same inputs, so train and - predict produce identical architecture/preprocessing. + ``test_id`` and ``validation_id`` (row indices of the test and held-out training + cells), ``chrom_list``, the preprocessed and the original target var names, and + ``rna_target_sum``, the total scButterfly scales every cell's GEX counts to. + Deterministic given the same inputs, so train and predict produce identical + architecture/preprocessing. """ direction, mod1, mod2 = detect_direction(train_mod1, train_mod2) logger.info("Direction: %s (mod1=%s, mod2=%s)", direction, mod1, mod2) @@ -268,9 +318,18 @@ def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly): sort_index, chrom_list = chrom_utils.sorted_chrom_order(ATAC_data) ATAC_data = chrom_utils.apply_sort(ATAC_data, sort_index) + # What scanpy's normalize_total scales every cell to: the median count total of the + # cells that have any counts. + rna_totals = np.asarray(RNA_data.X.sum(axis=1)).ravel() + rna_target_sum = float(np.median(rna_totals[rna_totals > 0])) + butterfly = Butterfly() butterfly.load_data(RNA_data, ATAC_data, train_id_final, test_id, validation_id) - butterfly.data_preprocessing(n_top_genes=n_top_genes) + butterfly.data_preprocessing( + n_top_genes=n_top_genes, + use_hvg=not (model_all_target_features and direction == "ATAC2GEX"), + filter_features=not (model_all_target_features and direction == "GEX2ATAC"), + ) butterfly.augmentation(aug_type=None) # scButterfly casts to float32 only on the CUDA path; force float32 for CPU. @@ -298,31 +357,39 @@ def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly): "n_train": n_train, "n_test": n_test, "test_id": test_id, + "validation_id": validation_id, "chrom_list": chrom_list, "target_p_var_names": list(target_p.var_names), + "target_var_names": list(train_mod2.var_names), + "rna_target_sum": rna_target_sum, } -def extract_predictions(built, A2R_predict, R2A_predict, target_var_names): - """Select the target-modality prediction, name its vars, and scatter to target order.""" - direction = built["direction"] - n_train, n_test = built["n_train"], built["n_test"] - pred = A2R_predict if direction == "ATAC2GEX" else R2A_predict - - # tensor2adata gives integer var names; assign the preprocessed target var names. - p_names = built["target_p_var_names"] - if pred.n_vars == len(p_names): - pred.var_names = p_names - else: - logger.warning("Prediction vars (%d) != preprocessed target vars (%d)", - pred.n_vars, len(p_names)) - - if pred.n_obs == (n_train + n_test): - pred = pred[built["test_id"]].copy() - elif pred.n_obs != n_test: - logger.warning("Unexpected prediction rows: %d (n_test=%d)", pred.n_obs, n_test) +def predict_cells(built, cell_ids, batch_size, model_path=None): + """scButterfly's prediction of the target modality for rows ``cell_ids``. - out = chrom_utils.scatter_to_target(pred, target_var_names) - overlap = len(set(pred.var_names) & set(target_var_names)) - logger.info("Var overlap with target: %d / %d", overlap, len(target_var_names)) - return out + Returned as a dense float32 array in the target's original var order, still on + scButterfly's scale (see ``target_scale``). ``model_path`` loads the trained + weights first; leave it out once they are loaded. + """ + # Model.test also runs PCA + a neighbour graph on both predicted matrices with no + # way to opt out; none of it is read back. See suppress_unused_postprocessing. + with suppress_unused_postprocessing(): + A2R_predict, R2A_predict = built["butterfly"].model.test( + test_id_r=cell_ids, + test_id_a=cell_ids, + batch_size=batch_size, + model_path=model_path, + load_model=model_path is not None, + test_cluster=False, + test_figure=False, + output_data=False, + return_predict=True, + ) + prediction = A2R_predict if built["direction"] == "ATAC2GEX" else R2A_predict + # tensor2adata gives integer var names; name them, then scatter by name, which + # undoes the chromosome sort and zero-fills target features the model dropped. + prediction.var_names = built["target_p_var_names"] + n_modelled = len(set(prediction.var_names) & set(built["target_var_names"])) + logger.info("Modelled target features: %d / %d", n_modelled, len(built["target_var_names"])) + return chrom_utils.scatter_to_target(prediction, built["target_var_names"]) diff --git a/src/methods/scbutterfly/chrom_utils.py b/src/methods/scbutterfly/chrom_utils.py index fe5a612a..aa2d0205 100644 --- a/src/methods/scbutterfly/chrom_utils.py +++ b/src/methods/scbutterfly/chrom_utils.py @@ -4,7 +4,7 @@ chromosome) and reads ``ATAC_data.var.chrom`` during model construction. It also assumes peaks are contiguous per chromosome. The predict-modality ATAC h5ads have no ``chrom`` column and peaks are in arbitrary chromosome order, but the peak names -encode the chromosome (e.g. ``chr17-6651156-6652045``). +encode the chromosome (``chr17-6651156-6652045`` or ``chr17:6651156-6652045``). This module parses the chromosome from peak names, produces a peak ordering that groups peaks contiguously per chromosome (with the matching ``chrom_list``), and @@ -16,23 +16,15 @@ import numpy as np from scipy.sparse import issparse -# Matches a leading chromosome token like "chr17", "chrX", "17", "X" followed by -# a ':' or '-' delimiter. Used as a fallback when the simple split does not yield -# a recognisable chromosome. -_CHROM_RE = re.compile(r"^chr?([0-9XYMxym]+)[:\-]") - def parse_chrom(name): - """Return the chromosome token for a peak name (e.g. 'chr17').""" - name = str(name) - token = name.split("-", 1)[0] - if token.lower().startswith("chr"): - return token - m = _CHROM_RE.match(name) - if m: - return "chr" + m.group(1) - # Unparseable — bucket everything unknown together so it still forms one group. - return token + """Return the contig of a peak name. + + NeurIPS 2021 names peaks ``chr17-6651156-6652045`` and NeurIPS 2022 + ``chr17:6651156-6652045``; both also have peaks on unplaced scaffolds + (``GL000195.1:32231-33125``), which keep their scaffold as their own group. + """ + return re.split(r"[:\-]", str(name), maxsplit=1)[0] def sorted_chrom_order(atac_adata): diff --git a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml index be01fbe7..795f463b 100644 --- a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml @@ -7,8 +7,12 @@ description: | adversarial translator that learns to convert between paired single-cell modalities. For the predict-modality task it is trained on paired Multiome cells and used to translate the held-out test modality. Chromosome grouping for the ATAC branch is parsed from peak - coordinates in the feature names. Only Multiome GEX<->ATAC is supported; CITE-seq (ADT) - datasets are not handled. + coordinates in the feature names. scButterfly's feature selection (highly variable genes, + peak filtering) applies to the input modality only, so every target feature is modelled. + Predictions are mapped from scButterfly's own preprocessing onto the target layer: exactly + to log CP10k for GEX, and by a per-peak least-squares line fitted on held-out training + cells for ATAC. Only Multiome GEX<->ATAC is supported; CITE-seq (ADT) datasets are not + handled. references: doi: - 10.1038/s41467-024-47418-x diff --git a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml index 7fba55f5..865dc606 100644 --- a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml @@ -12,6 +12,7 @@ resources: type: python_script - path: ../butterfly_common.py - path: ../chrom_utils.py + - path: ../target_scale.py - path: /src/utils/exit_codes.py test_resources: - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap diff --git a/src/methods/scbutterfly/scbutterfly_predict/script.py b/src/methods/scbutterfly/scbutterfly_predict/script.py index f1fe8b2f..a32b0f40 100644 --- a/src/methods/scbutterfly/scbutterfly_predict/script.py +++ b/src/methods/scbutterfly/scbutterfly_predict/script.py @@ -22,6 +22,7 @@ sys.path.append(meta["resources_dir"]) import butterfly_common +import target_scale butterfly_common.apply_runtime_patches() from scButterfly.butterfly import Butterfly @@ -45,24 +46,30 @@ built = butterfly_common.build_butterfly( train_mod1, train_mod2, test_mod1, n_top_genes=metadata["n_top_genes"], Butterfly=Butterfly, + # Models trained before every target feature was modelled lack the key. + model_all_target_features=metadata.get("model_all_target_features", False), ) -butterfly = built["butterfly"] +batch_size = metadata["batch_size"] logger.info("Loading trained weights and predicting...") -# test_model also runs PCA + a neighbour graph on both predicted matrices with no -# way to opt out; none of it is read back. See suppress_unused_postprocessing. -with butterfly_common.suppress_unused_postprocessing(): - A2R_predict, R2A_predict = butterfly.test_model( - batch_size=metadata["batch_size"], - model_path=par["input_model"], - load_model=True, - ) - -target_var_names = list(train_mod2.var_names) -test_predictions = butterfly_common.extract_predictions( - built, A2R_predict, R2A_predict, target_var_names, +test_predictions = butterfly_common.predict_cells( + built, built["test_id"], batch_size, model_path=par["input_model"], ) +# scButterfly predicts its own preprocessing of the target; map it onto the scale of +# the target's normalized layer (see target_scale). +if built["direction"] == "ATAC2GEX": + logger.info("Converting predictions to log CP10k (scButterfly target sum %.1f)...", built["rna_target_sum"]) + test_predictions = target_scale.to_log_cp10k(test_predictions, built["rna_target_sum"]) +else: + logger.info("Calibrating peak predictions on %d held-out training cells...", len(built["validation_id"])) + calibration = target_scale.PeakCalibration().fit( + butterfly_common.predict_cells(built, built["validation_id"], batch_size), + train_mod2.layers["normalized"][built["validation_id"]], + ) + logger.info("Peaks with zero slope: %d / %d", (calibration.slope_ == 0).sum(), calibration.slope_.size) + test_predictions = calibration.apply(test_predictions) + # --------------------------------------------------------------------------- # Write predictions. # --------------------------------------------------------------------------- diff --git a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml index 0308521b..abf20596 100644 --- a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml @@ -44,7 +44,7 @@ arguments: - name: "--n_top_genes" type: integer default: 3000 - description: Number of highly variable genes kept by scButterfly RNA preprocessing. + description: Number of highly variable genes scButterfly keeps when GEX is the input (GEX -> ATAC). When GEX is the target, every gene is modelled. engines: - type: docker image: python:3.9 diff --git a/src/methods/scbutterfly/scbutterfly_train/script.py b/src/methods/scbutterfly/scbutterfly_train/script.py index ed8dbedf..20c42310 100644 --- a/src/methods/scbutterfly/scbutterfly_train/script.py +++ b/src/methods/scbutterfly/scbutterfly_train/script.py @@ -68,6 +68,7 @@ metadata = { "direction": built["direction"], "n_top_genes": par["n_top_genes"], + "model_all_target_features": True, "batch_size": par["batch_size"], "target_var_names": list(train_mod2.var_names), "dataset_id": train_mod1.uns.get("dataset_id", ""), diff --git a/src/methods/scbutterfly/target_scale.py b/src/methods/scbutterfly/target_scale.py new file mode 100644 index 00000000..c36420fd --- /dev/null +++ b/src/methods/scbutterfly/target_scale.py @@ -0,0 +1,62 @@ +"""Map scButterfly's predictions onto the scale of the target's ``normalized`` layer. + +scButterfly predicts its own preprocessing of the target modality, which is not the +layer the benchmark scores: + +- GEX: ``log1p`` of each cell's counts scaled to a total of ``target_sum``, the median + count total of the cells (scanpy's ``normalize_total`` default), where the dataset + has log CP10k. With every gene modelled, the two differ by a constant factor inside + the log, so :func:`to_log_cp10k` converts exactly. +- ATAC: binarized peaks, TF-IDF-transformed and divided by the global maximum, + predicted through a sigmoid, where the dataset has its own TF-IDF (0 to ~20). There is + no closed form, so :class:`PeakCalibration` fits one least-squares line per peak on + cells held out from training. + +Both maps are increasing, so per-gene and per-peak rankings stay as the model made them +(a peak whose fitted slope is zero becomes constant, see below); what they fix is the +level and spread that the per-cell profiles, overall correlations and RMSE/MAE score. +""" + +import numpy as np +import scipy.sparse as sp + +LOG_CP10K_TOTAL = 1e4 + + +def to_log_cp10k(predictions, target_sum): + """Convert ``log1p(counts / total * target_sum)`` to ``log1p(counts / total * 1e4)``. + + The RNA decoder ends in a LeakyReLU, so small negative outputs (no expression) are + clipped to zero first. + """ + counts_per_target_sum = np.expm1(np.clip(predictions, 0, None)) + return np.log1p(counts_per_target_sum * (LOG_CP10K_TOTAL / target_sum)).astype(np.float32) + + +class PeakCalibration: + """Per-peak least-squares line from scButterfly's ATAC output to the target layer. + + A peak whose held-out predictions do not correlate positively with the truth gets + slope zero, i.e. its mean: a negative slope would turn the model's ranking of the + cells around, which it is not trained to do. + """ + + def fit(self, predictions, targets): + """Fit on ``predictions`` (cells x peaks, dense) and ``targets`` (dense or sparse).""" + predictions = np.asarray(predictions, dtype=np.float64) + prediction_mean = predictions.mean(axis=0) + centered = predictions - prediction_mean + if sp.issparse(targets): + target_mean = np.asarray(targets.mean(axis=0)).ravel() + covariance = np.asarray(sp.csr_matrix(targets).multiply(centered).sum(axis=0)).ravel() + else: + targets = np.asarray(targets, dtype=np.float64) + target_mean = targets.mean(axis=0) + covariance = (targets * centered).sum(axis=0) + variance = (centered**2).sum(axis=0) + self.slope_ = np.where(variance > 0, np.maximum(covariance, 0) / np.where(variance > 0, variance, 1), 0) + self.intercept_ = target_mean - self.slope_ * prediction_mean + return self + + def apply(self, predictions): + return (np.asarray(predictions) * self.slope_ + self.intercept_).astype(np.float32) From 339fa09617265950744262a2d7310e9767e1bc38 Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Mon, 5 Oct 2026 09:23:57 +0200 Subject: [PATCH 3/8] scbutterfly: densify training batches on the device instead of per cell in DataLoader workers Co-Authored-By: Claude Opus 5.5 --- src/methods/scbutterfly/butterfly_common.py | 90 +++++++++++++++------ 1 file changed, 64 insertions(+), 26 deletions(-) diff --git a/src/methods/scbutterfly/butterfly_common.py b/src/methods/scbutterfly/butterfly_common.py index 8bbbf619..f0eb750b 100644 --- a/src/methods/scbutterfly/butterfly_common.py +++ b/src/methods/scbutterfly/butterfly_common.py @@ -16,6 +16,7 @@ import numpy as np import pandas as pd import scanpy as sc +import torch from scipy.sparse import csr_matrix, issparse import chrom_utils @@ -30,9 +31,9 @@ def apply_runtime_patches(): - BCE on a float32 sigmoid can drift >1.0 -> clamp input to [0,1] - scButterfly hardcodes ``.cuda()``; make it a no-op when no GPU is present - ``TFIDF`` builds dense (n_peaks, n_cells) tiles -> sparse equivalent - - ``Model`` densifies both whole modalities -> densify one cell at a time + - ``Model`` densifies both whole modalities -> keep them sparse + - per-cell ``DataLoader`` batches -> batches densified on the device """ - import torch import torch.nn as nn import torch.nn.functional as F @@ -81,7 +82,8 @@ def _safe_bce(input, target, *args, **kwargs): from scButterfly import train_model - train_model.Model.__init__ = _rows_on_demand(train_model.Model.__init__) + train_model.Model.__init__ = _keep_data_sparse(train_model.Model.__init__) + train_model.DataLoader = _SparseBatchLoader def _sparse_tfidf(count_mat): @@ -125,31 +127,16 @@ def _sparse_tfidf(count_mat): return out, cell_totals, idf -class _DenseRows: - """A CSR matrix that hands out one cell at a time as a dense float32 row. +def _keep_data_sparse(original_init): + """Wrap ``Model.__init__`` so that it keeps the data as float32 CSR matrices. - ``Model.__init__`` stores ``RNA_data.X.toarray()`` and ``ATAC_data.X.toarray()``: - every cell of both modalities, training and test, as dense float32. At NeurIPS 2022 - Multiome scale (131k cells x 175k peaks after peak filtering) the ATAC matrix alone - is 92 GB. The model only ever reads ``.shape`` and single rows (``data[cell, :]`` in - its datasets), so this serves those rows from the sparse matrix instead. + Upstream stores ``RNA_data.X.toarray()`` and ``ATAC_data.X.toarray()``: every cell + of both modalities, training and test, as dense float32. At NeurIPS 2022 Multiome + scale (131k cells x 175k peaks after peak filtering) the ATAC matrix alone is 92 GB. + The model reads only their ``.shape`` and, through its datasets, batches of rows, + which :class:`_SparseBatchLoader` serves from the sparse matrices. """ - def __init__(self, matrix): - self.matrix = csr_matrix(matrix, dtype=np.float32) - self.shape = self.matrix.shape - - def __getitem__(self, index): - cell = index[0] if isinstance(index, tuple) else index - start, end = self.matrix.indptr[cell], self.matrix.indptr[cell + 1] - row = np.zeros(self.shape[1], dtype=np.float32) - row[self.matrix.indices[start:end]] = self.matrix.data[start:end] - return row - - -def _rows_on_demand(original_init): - """Wrap ``Model.__init__`` to keep the data sparse (see :class:`_DenseRows`).""" - def init(self, RNA_data, ATAC_data, *args, **kwargs): # The networks are sized from the dim lists, not the data, so build them on # zero cells and attach the real matrices afterwards. @@ -157,11 +144,62 @@ def init(self, RNA_data, ATAC_data, *args, **kwargs): self.RNA_data_obs, self.ATAC_data_obs = RNA_data.obs, ATAC_data.obs # Upstream stores the RNA var as ATAC_data_var too; nothing reads either. self.RNA_data_var, self.ATAC_data_var = RNA_data.var, RNA_data.var - self.RNA_data, self.ATAC_data = _DenseRows(RNA_data.X), _DenseRows(ATAC_data.X) + self.RNA_data = csr_matrix(RNA_data.X, dtype=np.float32) + self.ATAC_data = csr_matrix(ATAC_data.X, dtype=np.float32) return init +class _SparseBatchLoader: + """Stand-in for the ``DataLoader`` scButterfly builds over its datasets. + + ``DataLoader`` workers assembled every batch from dense per-cell rows on the CPU, + and the training loop copied it to the GPU. With every target feature modelled, a + NeurIPS 2022 Multiome cell is 198k floats (23k genes, 175k peaks), so a batch of 64 + is 51 MB and data loading dominated each step: the GPU sat at ~28% and one ATAC + pretraining epoch took 128 s on a V100. This loader gathers a batch's rows from the + CSR matrices, ships only their stored entries to the device and scatters them into + a dense tensor there. + + Batches hold the same values in the same layout (RNA columns, then ATAC), and come + back on the device as float32, so the training loop's ``.cuda().to(torch.float32)`` + is a no-op. Shuffling draws a ``torch.randperm``, like ``DataLoader``, so it follows + the seed scButterfly sets. + """ + + def __init__(self, dataset, batch_size, shuffle=False, num_workers=0, drop_last=False): + if hasattr(dataset, "id_list"): # Single_omics_dataset + self.blocks = [(dataset.dataset, np.asarray(dataset.id_list))] + else: # RNA_ATAC_dataset + self.blocks = [ + (dataset.RNA_dataset, np.asarray(dataset.id_list_r)), + (dataset.ATAC_dataset, np.asarray(dataset.id_list_a)), + ] + self.n_cells = len(self.blocks[0][1]) + self.batch_size, self.shuffle, self.drop_last = batch_size, shuffle, drop_last + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + def __len__(self): + if self.drop_last: + return self.n_cells // self.batch_size + return -(-self.n_cells // self.batch_size) + + def __iter__(self): + order = torch.randperm(self.n_cells).numpy() if self.shuffle else np.arange(self.n_cells) + for start in range(0, len(self) * self.batch_size, self.batch_size): + positions = order[start:start + self.batch_size] + yield torch.cat([self._dense_rows(matrix, cells[positions]) for matrix, cells in self.blocks], dim=1) + + def _dense_rows(self, matrix, cells): + rows = matrix[cells] + row_of_entry = np.repeat(np.arange(len(cells)), np.diff(rows.indptr)) + dense = torch.zeros((len(cells), matrix.shape[1]), dtype=torch.float32, device=self.device) + dense[torch.from_numpy(row_of_entry).to(self.device), torch.from_numpy(rows.indices.astype(np.int64)).to(self.device)] = ( + torch.from_numpy(rows.data).to(self.device) + ) + return dense + + @contextlib.contextmanager def suppress_unused_postprocessing(): """Make ``sc.pp.pca``/``sc.pp.neighbors`` no-ops for the duration of the block. From 01f51add1040d2f6b4e1960df6ea6a190f27231d Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Mon, 5 Oct 2026 19:30:38 +0200 Subject: [PATCH 4/8] scbutterfly: one global scale line, lower peak memory, larger resource labels - target_scale.ScaleCalibration: map both GEX and ATAC predictions with one least-squares line fitted on the held-out validation cells. On the 2021 and 2022 Multiome datasets it matched or beat per-feature lines and, for GEX, the exact log CP10k conversion (which inflates the decoder's conditional means for sparse genes: RMSE 0.553 vs 0.518 on bmmc_multiome); it leaves every correlation as is. - predict_cells: give the direction Model.test does not need a single cell. - read_modality drops unused layers/obsm; _to_counts_X takes the counts layer out of its input; build_butterfly frees intermediate copies. - train: veryhighmem + veryhightime; predict: veryhighmem. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 4 +- src/methods/scbutterfly/butterfly_common.py | 58 +++++++++---- .../scbutterfly/scbutterfly/config.vsh.yaml | 7 +- .../scbutterfly_predict/config.vsh.yaml | 2 +- .../scbutterfly/scbutterfly_predict/script.py | 28 +++---- .../scbutterfly_train/config.vsh.yaml | 2 +- .../scbutterfly/scbutterfly_train/script.py | 6 +- src/methods/scbutterfly/target_scale.py | 82 ++++++++----------- 8 files changed, 100 insertions(+), 89 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9dc1b3fd..73986e10 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,9 +60,9 @@ * `scbutterfly`: Parse the chromosome of `chr1:1000-2000` peak names, not only `chr1-1000-2000`. NeurIPS 2022 names its peaks the former way, so every peak became its own chromosome, each with its own block in the ATAC encoder: `pbmc_multiome` was OOM-killed when GEX was the target (174,978 "chromosomes") and hit the 8 h limit when ATAC was (7,002 "chromosomes", 5.7 h for one pretraining epoch). -* `scbutterfly`: Model every target feature and predict it on the scale of the target's `normalized` layer. scButterfly kept only the 3,000 most variable genes and the peaks open in at least 0.5% of the cells, also when they were the target, so the other genes and peaks were predicted as zero; and it predicts its own preprocessing of the target. GEX predictions are now converted exactly from scButterfly's median-library-size scaling to log CP10k, and ATAC predictions (a sigmoid of the max-scaled TF-IDF) are mapped onto the dataset's TF-IDF by one least-squares line per peak, fitted on the training cells held out for validation. The HVG selection and peak filter still apply to the input modality. `scbutterfly_predict` reproduces the old preprocessing for model bundles trained before this change. +* `scbutterfly`: Model every target feature and predict it on the scale of the target's `normalized` layer. scButterfly kept only the 3,000 most variable genes and the peaks open in at least 0.5% of the cells, also when they were the target, so the other genes and peaks were predicted as zero; and it predicts its own preprocessing of the target. Predictions are now mapped onto the target layer by one least-squares line, fitted on the training cells scButterfly holds out for validation. It scored as well as or better than one line per feature, and, for GEX, than the exact conversion from scButterfly's median-library-size scaling to log CP10k: the decoder predicts a conditional mean in its own log space, which the exact conversion inflates for sparsely expressed genes. The HVG selection and peak filter still apply to the input modality. `scbutterfly_predict` reproduces the old preprocessing for model bundles trained before this change. -* `scbutterfly`: Serve the model's training and test cells one sparse row at a time instead of densifying both whole modalities, 92 GB for the ATAC input of `pbmc_multiome` alone. +* `scbutterfly`: Keep the data sparse and densify each batch on the GPU. The model densified both whole modalities (92 GB for the ATAC input of `pbmc_multiome` alone), and its `DataLoader` workers built every batch from dense per-cell rows, so data loading dominated each step. Training on `pbmc_multiome` (ATAC -> GEX) now takes 5.9 h on a V100 instead of 10.0 h. Unused layers and intermediate copies are also dropped early. `scbutterfly_train` asks for `veryhighmem` and `veryhightime`, `scbutterfly_predict` for `veryhighmem`: the model rebuilds scButterfly's preprocessing in both, which peaks at ~100 GB on `pbmc_multiome`, and the T4 the benchmark runs on is slower than the V100. * `cellmapper_linear`: Write the unmasked variants as `mask_var: "none"` rather than `mask_var: null`. Viash drops null values from the config, so these variants fell back to the default `mask_var: "hvg"` and duplicated their `-hvg` counterparts. diff --git a/src/methods/scbutterfly/butterfly_common.py b/src/methods/scbutterfly/butterfly_common.py index f0eb750b..358d3be4 100644 --- a/src/methods/scbutterfly/butterfly_common.py +++ b/src/methods/scbutterfly/butterfly_common.py @@ -144,12 +144,17 @@ def init(self, RNA_data, ATAC_data, *args, **kwargs): self.RNA_data_obs, self.ATAC_data_obs = RNA_data.obs, ATAC_data.obs # Upstream stores the RNA var as ATAC_data_var too; nothing reads either. self.RNA_data_var, self.ATAC_data_var = RNA_data.var, RNA_data.var - self.RNA_data = csr_matrix(RNA_data.X, dtype=np.float32) - self.ATAC_data = csr_matrix(ATAC_data.X, dtype=np.float32) + self.RNA_data, self.ATAC_data = _float32_csr(RNA_data.X), _float32_csr(ATAC_data.X) return init +def _float32_csr(matrix): + """``matrix`` as float32 CSR, without a copy when it already is one.""" + matrix = matrix.tocsr() if issparse(matrix) else csr_matrix(matrix) + return matrix.astype(np.float32, copy=False) + + class _SparseBatchLoader: """Stand-in for the ``DataLoader`` scButterfly builds over its datasets. @@ -230,9 +235,26 @@ def detect_direction(train_mod1, train_mod2): return direction, mod1, mod2 +def read_modality(path, layers=("counts",)): + """Read an h5ad, keeping only ``layers`` besides ``obs``, ``var`` and ``uns``. + + scButterfly reads only the counts (predict also the target's ``normalized`` layer, + for the scale map). The NeurIPS 2022 Multiome ATAC files hold two 9 GB layers. + """ + adata = ad.read_h5ad(path) + for key in [key for key in adata.layers.keys() if key not in layers]: + del adata.layers[key] + for key in list(adata.obsm.keys()): + del adata.obsm[key] + return adata + + def _to_counts_X(adata): """AnnData whose .X is the raw counts layer as float32, carrying obs and var. + Takes the counts layer out of ``adata``: nothing reads the float64 original again, + and for the NeurIPS 2022 ATAC input it is 9 GB. + Counts are stored as float64; scButterfly's model is float32 and its data loader does not cast, so feed float32 to avoid a dtype mismatch. @@ -242,7 +264,7 @@ def _to_counts_X(adata): duplicates the largest matrices in the file to throw them away a moment later. """ return ad.AnnData( - X=adata.layers["counts"].astype(np.float32), + X=adata.layers.pop("counts").astype(np.float32), obs=adata.obs.copy(), var=adata.var.copy(), ) @@ -299,10 +321,11 @@ def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly, one the model drops could only be predicted as a constant zero. ``False`` reproduces models trained before this option existed. + Consumes the ``counts`` layers of the three inputs (see :func:`_to_counts_X`). + Returns a dict with the constructed ``butterfly``, the resolved ``direction``, ``test_id`` and ``validation_id`` (row indices of the test and held-out training - cells), ``chrom_list``, the preprocessed and the original target var names, and - ``rna_target_sum``, the total scButterfly scales every cell's GEX counts to. + cells), ``chrom_list``, and the preprocessed and the original target var names. Deterministic given the same inputs, so train and predict produce identical architecture/preprocessing. """ @@ -340,6 +363,7 @@ def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly, RNA_data = _concat_blocks(rna_train_c, rna_test_c) ATAC_data = _concat_blocks(atac_train_c, atac_test_c) + del rna_train_c, rna_test_c, atac_train_c, atac_test_c train_id = list(range(n_train)) test_id = list(range(n_train, n_train + n_test)) @@ -356,25 +380,23 @@ def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly, sort_index, chrom_list = chrom_utils.sorted_chrom_order(ATAC_data) ATAC_data = chrom_utils.apply_sort(ATAC_data, sort_index) - # What scanpy's normalize_total scales every cell to: the median count total of the - # cells that have any counts. - rna_totals = np.asarray(RNA_data.X.sum(axis=1)).ravel() - rna_target_sum = float(np.median(rna_totals[rna_totals > 0])) - butterfly = Butterfly() butterfly.load_data(RNA_data, ATAC_data, train_id_final, test_id, validation_id) + del RNA_data, ATAC_data # load_data keeps its own copies butterfly.data_preprocessing( n_top_genes=n_top_genes, use_hvg=not (model_all_target_features and direction == "ATAC2GEX"), filter_features=not (model_all_target_features and direction == "GEX2ATAC"), ) butterfly.augmentation(aug_type=None) + # Only the preprocessed copies are read from here on. + butterfly.RNA_data = butterfly.ATAC_data = None # scButterfly casts to float32 only on the CUDA path; force float32 for CPU. for _attr in ("RNA_data_p", "ATAC_data_p"): _adata = getattr(butterfly, _attr, None) if _adata is not None and _adata.X is not None: - _adata.X = _adata.X.astype(np.float32) + _adata.X = _adata.X.astype(np.float32, copy=False) # Rebuild chrom_list from the preprocessed (peak-filtered) ATAC so model dims match. atac_p = getattr(butterfly, "ATAC_data_p", None) @@ -399,7 +421,6 @@ def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly, "chrom_list": chrom_list, "target_p_var_names": list(target_p.var_names), "target_var_names": list(train_mod2.var_names), - "rna_target_sum": rna_target_sum, } @@ -410,12 +431,15 @@ def predict_cells(built, cell_ids, batch_size, model_path=None): scButterfly's scale (see ``target_scale``). ``model_path`` loads the trained weights first; leave it out once they are loaded. """ - # Model.test also runs PCA + a neighbour graph on both predicted matrices with no - # way to opt out; none of it is read back. See suppress_unused_postprocessing. + # Model.test predicts both directions; give the unused one a single cell, which + # for ATAC -> GEX on NeurIPS 2022 saves a cells x 175k-peak matrix. It also runs + # PCA + a neighbour graph on both predicted matrices with no way to opt out; none + # of it is read back. See suppress_unused_postprocessing. + gex_to_atac = built["direction"] == "GEX2ATAC" with suppress_unused_postprocessing(): A2R_predict, R2A_predict = built["butterfly"].model.test( - test_id_r=cell_ids, - test_id_a=cell_ids, + test_id_r=cell_ids if gex_to_atac else cell_ids[:1], + test_id_a=cell_ids[:1] if gex_to_atac else cell_ids, batch_size=batch_size, model_path=model_path, load_model=model_path is not None, @@ -424,7 +448,7 @@ def predict_cells(built, cell_ids, batch_size, model_path=None): output_data=False, return_predict=True, ) - prediction = A2R_predict if built["direction"] == "ATAC2GEX" else R2A_predict + prediction = R2A_predict if gex_to_atac else A2R_predict # tensor2adata gives integer var names; name them, then scatter by name, which # undoes the chromosome sort and zero-fills target features the model dropped. prediction.var_names = built["target_p_var_names"] diff --git a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml index 795f463b..a2714fc2 100644 --- a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml @@ -9,10 +9,9 @@ description: | the held-out test modality. Chromosome grouping for the ATAC branch is parsed from peak coordinates in the feature names. scButterfly's feature selection (highly variable genes, peak filtering) applies to the input modality only, so every target feature is modelled. - Predictions are mapped from scButterfly's own preprocessing onto the target layer: exactly - to log CP10k for GEX, and by a per-peak least-squares line fitted on held-out training - cells for ATAC. Only Multiome GEX<->ATAC is supported; CITE-seq (ADT) datasets are not - handled. + Predictions are mapped from scButterfly's own preprocessing onto the target layer by one + least-squares line fitted on held-out training cells. Only Multiome GEX<->ATAC is + supported; CITE-seq (ADT) datasets are not handled. references: doi: - 10.1038/s41467-024-47418-x diff --git a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml index 865dc606..9bba2e48 100644 --- a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml @@ -33,4 +33,4 @@ runners: - type: executable - type: nextflow directives: - label: [highmem, hightime, midcpu, highsharedmem, gpu] + label: [veryhighmem, hightime, midcpu, highsharedmem, gpu] diff --git a/src/methods/scbutterfly/scbutterfly_predict/script.py b/src/methods/scbutterfly/scbutterfly_predict/script.py index a32b0f40..43d1c4d3 100644 --- a/src/methods/scbutterfly/scbutterfly_predict/script.py +++ b/src/methods/scbutterfly/scbutterfly_predict/script.py @@ -31,9 +31,10 @@ # Load data + model metadata # --------------------------------------------------------------------------- logger.info("Reading input files...") -train_mod1 = ad.read_h5ad(par["input_train_mod1"]) -train_mod2 = ad.read_h5ad(par["input_train_mod2"]) -test_mod1 = ad.read_h5ad(par["input_test_mod1"]) +train_mod1 = butterfly_common.read_modality(par["input_train_mod1"]) +# The target's normalized layer is what the scale map is fitted to. +train_mod2 = butterfly_common.read_modality(par["input_train_mod2"], layers=("counts", "normalized")) +test_mod1 = butterfly_common.read_modality(par["input_test_mod1"]) with open(os.path.join(par["input_model"], "metadata.pkl"), "rb") as f: metadata = pickle.load(f) @@ -57,18 +58,15 @@ ) # scButterfly predicts its own preprocessing of the target; map it onto the scale of -# the target's normalized layer (see target_scale). -if built["direction"] == "ATAC2GEX": - logger.info("Converting predictions to log CP10k (scButterfly target sum %.1f)...", built["rna_target_sum"]) - test_predictions = target_scale.to_log_cp10k(test_predictions, built["rna_target_sum"]) -else: - logger.info("Calibrating peak predictions on %d held-out training cells...", len(built["validation_id"])) - calibration = target_scale.PeakCalibration().fit( - butterfly_common.predict_cells(built, built["validation_id"], batch_size), - train_mod2.layers["normalized"][built["validation_id"]], - ) - logger.info("Peaks with zero slope: %d / %d", (calibration.slope_ == 0).sum(), calibration.slope_.size) - test_predictions = calibration.apply(test_predictions) +# the target's normalized layer with a line fitted on the held-out training cells +# (see target_scale). +logger.info("Fitting the scale map on %d held-out training cells...", len(built["validation_id"])) +calibration = target_scale.ScaleCalibration().fit( + butterfly_common.predict_cells(built, built["validation_id"], batch_size), + train_mod2.layers["normalized"][built["validation_id"]], +) +logger.info("Target = %.4f + %.4f x scButterfly output", calibration.intercept_, calibration.slope_) +test_predictions = calibration.apply(test_predictions) # --------------------------------------------------------------------------- # Write predictions. diff --git a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml index abf20596..4e7257a0 100644 --- a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml @@ -61,4 +61,4 @@ runners: - type: executable - type: nextflow directives: - label: [highmem, hightime, midcpu, highsharedmem, gpu] + label: [veryhighmem, veryhightime, midcpu, highsharedmem, gpu] diff --git a/src/methods/scbutterfly/scbutterfly_train/script.py b/src/methods/scbutterfly/scbutterfly_train/script.py index 20c42310..0c8ce853 100644 --- a/src/methods/scbutterfly/scbutterfly_train/script.py +++ b/src/methods/scbutterfly/scbutterfly_train/script.py @@ -34,9 +34,9 @@ # Load data # --------------------------------------------------------------------------- logger.info("Reading input files...") -train_mod1 = ad.read_h5ad(par["input_train_mod1"]) -train_mod2 = ad.read_h5ad(par["input_train_mod2"]) -test_mod1 = ad.read_h5ad(par["input_test_mod1"]) +train_mod1 = butterfly_common.read_modality(par["input_train_mod1"]) +train_mod2 = butterfly_common.read_modality(par["input_train_mod2"]) +test_mod1 = butterfly_common.read_modality(par["input_test_mod1"]) # --------------------------------------------------------------------------- # Build + construct the model, then train (weights written to /model). diff --git a/src/methods/scbutterfly/target_scale.py b/src/methods/scbutterfly/target_scale.py index c36420fd..14fffd2d 100644 --- a/src/methods/scbutterfly/target_scale.py +++ b/src/methods/scbutterfly/target_scale.py @@ -3,58 +3,48 @@ scButterfly predicts its own preprocessing of the target modality, which is not the layer the benchmark scores: -- GEX: ``log1p`` of each cell's counts scaled to a total of ``target_sum``, the median - count total of the cells (scanpy's ``normalize_total`` default), where the dataset - has log CP10k. With every gene modelled, the two differ by a constant factor inside - the log, so :func:`to_log_cp10k` converts exactly. +- GEX: ``log1p`` of each cell's counts scaled to the median count total of the cells + (scanpy's ``normalize_total`` default) instead of 1e4. - ATAC: binarized peaks, TF-IDF-transformed and divided by the global maximum, - predicted through a sigmoid, where the dataset has its own TF-IDF (0 to ~20). There is - no closed form, so :class:`PeakCalibration` fits one least-squares line per peak on - cells held out from training. - -Both maps are increasing, so per-gene and per-peak rankings stay as the model made them -(a peak whose fitted slope is zero becomes constant, see below); what they fix is the -level and spread that the per-cell profiles, overall correlations and RMSE/MAE score. + predicted through a sigmoid (0 to 1), where the dataset has its own TF-IDF (0 to ~20). + +:class:`ScaleCalibration` maps them with one least-squares line, fitted on training +cells scButterfly held out for validation. On the NeurIPS 2021 and 2022 Multiome +datasets it scored as well as or better than the alternatives on every metric: one line +per gene or peak (noisier, and it reorders each cell's profile), and for GEX the exact +conversion ``log1p(expm1(p) * 1e4 / target_sum)``. The exact conversion is exact for a +value, not for the model's prediction of one: an MSE-trained decoder predicts a +conditional mean in its own log space, and the concave conversion inflates it for +sparsely expressed genes (RMSE 0.553 against 0.518 for the line on ``bmmc_multiome``). + +The line is increasing, so every correlation is exactly the model's; it fixes the level +and spread that RMSE scores. """ import numpy as np import scipy.sparse as sp -LOG_CP10K_TOTAL = 1e4 - - -def to_log_cp10k(predictions, target_sum): - """Convert ``log1p(counts / total * target_sum)`` to ``log1p(counts / total * 1e4)``. - - The RNA decoder ends in a LeakyReLU, so small negative outputs (no expression) are - clipped to zero first. - """ - counts_per_target_sum = np.expm1(np.clip(predictions, 0, None)) - return np.log1p(counts_per_target_sum * (LOG_CP10K_TOTAL / target_sum)).astype(np.float32) - - -class PeakCalibration: - """Per-peak least-squares line from scButterfly's ATAC output to the target layer. - - A peak whose held-out predictions do not correlate positively with the truth gets - slope zero, i.e. its mean: a negative slope would turn the model's ranking of the - cells around, which it is not trained to do. - """ - - def fit(self, predictions, targets): - """Fit on ``predictions`` (cells x peaks, dense) and ``targets`` (dense or sparse).""" - predictions = np.asarray(predictions, dtype=np.float64) - prediction_mean = predictions.mean(axis=0) - centered = predictions - prediction_mean - if sp.issparse(targets): - target_mean = np.asarray(targets.mean(axis=0)).ravel() - covariance = np.asarray(sp.csr_matrix(targets).multiply(centered).sum(axis=0)).ravel() - else: - targets = np.asarray(targets, dtype=np.float64) - target_mean = targets.mean(axis=0) - covariance = (targets * centered).sum(axis=0) - variance = (centered**2).sum(axis=0) - self.slope_ = np.where(variance > 0, np.maximum(covariance, 0) / np.where(variance > 0, variance, 1), 0) + +class ScaleCalibration: + """One least-squares line from scButterfly's output to the target layer.""" + + def fit(self, predictions, targets, block_size=1000): + """Fit on ``predictions`` (cells x features, dense) and ``targets`` (dense or sparse). + + A negative slope, a model anticorrelated with its target, is set to 0. + """ + n_entries = predictions.size + prediction_sum, prediction_square_sum = 0.0, 0.0 + for start in range(0, predictions.shape[0], block_size): + block = predictions[start:start + block_size].astype(np.float64) + prediction_sum += block.sum() + prediction_square_sum += (block**2).sum() + targets = sp.csr_matrix(targets, dtype=np.float64) + prediction_mean = prediction_sum / n_entries + target_mean = targets.sum() / n_entries + covariance = targets.multiply(predictions).sum() / n_entries - prediction_mean * target_mean + variance = prediction_square_sum / n_entries - prediction_mean**2 + self.slope_ = max(covariance, 0.0) / variance if variance > 0 else 0.0 self.intercept_ = target_mean - self.slope_ * prediction_mean return self From 1eeb08e32272abe4c97f2d417f2d3238708ade9a Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Mon, 5 Oct 2026 19:33:45 +0200 Subject: [PATCH 5/8] scbutterfly: rescale rather than flatten predictions when held-out cells show no positive association Co-Authored-By: Claude Opus 5.5 --- src/methods/scbutterfly/target_scale.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/methods/scbutterfly/target_scale.py b/src/methods/scbutterfly/target_scale.py index 14fffd2d..4fcf0389 100644 --- a/src/methods/scbutterfly/target_scale.py +++ b/src/methods/scbutterfly/target_scale.py @@ -31,7 +31,10 @@ class ScaleCalibration: def fit(self, predictions, targets, block_size=1000): """Fit on ``predictions`` (cells x features, dense) and ``targets`` (dense or sparse). - A negative slope, a model anticorrelated with its target, is set to 0. + If the held-out cells show no positive association, which only an essentially + untrained model produces (e.g. two epochs on the test resources), the line + matches the target's mean and standard deviation instead, so the predictions + are rescaled rather than replaced by a constant. """ n_entries = predictions.size prediction_sum, prediction_square_sum = 0.0, 0.0 @@ -44,7 +47,13 @@ def fit(self, predictions, targets, block_size=1000): target_mean = targets.sum() / n_entries covariance = targets.multiply(predictions).sum() / n_entries - prediction_mean * target_mean variance = prediction_square_sum / n_entries - prediction_mean**2 - self.slope_ = max(covariance, 0.0) / variance if variance > 0 else 0.0 + target_variance = (targets.data**2).sum() / n_entries - target_mean**2 + if variance <= 0: + self.slope_ = 0.0 + elif covariance > 0: + self.slope_ = covariance / variance + else: + self.slope_ = np.sqrt(max(target_variance, 0.0) / variance) self.intercept_ = target_mean - self.slope_ * prediction_mean return self From 2198e4cfd6a654a6b0e0cf50ef4231c663482e05 Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Mon, 5 Oct 2026 22:58:52 +0200 Subject: [PATCH 6/8] scbutterfly: state the scale-map comparison precisely Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- src/methods/scbutterfly/target_scale.py | 14 ++++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 73986e10..f8e0dcc4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,7 +60,7 @@ * `scbutterfly`: Parse the chromosome of `chr1:1000-2000` peak names, not only `chr1-1000-2000`. NeurIPS 2022 names its peaks the former way, so every peak became its own chromosome, each with its own block in the ATAC encoder: `pbmc_multiome` was OOM-killed when GEX was the target (174,978 "chromosomes") and hit the 8 h limit when ATAC was (7,002 "chromosomes", 5.7 h for one pretraining epoch). -* `scbutterfly`: Model every target feature and predict it on the scale of the target's `normalized` layer. scButterfly kept only the 3,000 most variable genes and the peaks open in at least 0.5% of the cells, also when they were the target, so the other genes and peaks were predicted as zero; and it predicts its own preprocessing of the target. Predictions are now mapped onto the target layer by one least-squares line, fitted on the training cells scButterfly holds out for validation. It scored as well as or better than one line per feature, and, for GEX, than the exact conversion from scButterfly's median-library-size scaling to log CP10k: the decoder predicts a conditional mean in its own log space, which the exact conversion inflates for sparsely expressed genes. The HVG selection and peak filter still apply to the input modality. `scbutterfly_predict` reproduces the old preprocessing for model bundles trained before this change. +* `scbutterfly`: Model every target feature and predict it on the scale of the target's `normalized` layer. scButterfly kept only the 3,000 most variable genes and the peaks open in at least 0.5% of the cells, also when they were the target, so the other genes and peaks were predicted as zero; and it predicts its own preprocessing of the target. Predictions are now mapped onto the target layer by one least-squares line, fitted on the training cells scButterfly holds out for validation. Apart from MAE, it came within 0.005 of the best alternative on every metric, while one line per feature lowered per-cell Spearman and, for GEX, the exact conversion from scButterfly's median-library-size scaling to log CP10k raised RMSE: the decoder predicts a conditional mean in its own log space, which the exact conversion inflates for sparsely expressed genes. The HVG selection and peak filter still apply to the input modality. `scbutterfly_predict` reproduces the old preprocessing for model bundles trained before this change. * `scbutterfly`: Keep the data sparse and densify each batch on the GPU. The model densified both whole modalities (92 GB for the ATAC input of `pbmc_multiome` alone), and its `DataLoader` workers built every batch from dense per-cell rows, so data loading dominated each step. Training on `pbmc_multiome` (ATAC -> GEX) now takes 5.9 h on a V100 instead of 10.0 h. Unused layers and intermediate copies are also dropped early. `scbutterfly_train` asks for `veryhighmem` and `veryhightime`, `scbutterfly_predict` for `veryhighmem`: the model rebuilds scButterfly's preprocessing in both, which peaks at ~100 GB on `pbmc_multiome`, and the T4 the benchmark runs on is slower than the V100. diff --git a/src/methods/scbutterfly/target_scale.py b/src/methods/scbutterfly/target_scale.py index 4fcf0389..93ccfa69 100644 --- a/src/methods/scbutterfly/target_scale.py +++ b/src/methods/scbutterfly/target_scale.py @@ -10,12 +10,14 @@ :class:`ScaleCalibration` maps them with one least-squares line, fitted on training cells scButterfly held out for validation. On the NeurIPS 2021 and 2022 Multiome -datasets it scored as well as or better than the alternatives on every metric: one line -per gene or peak (noisier, and it reorders each cell's profile), and for GEX the exact -conversion ``log1p(expm1(p) * 1e4 / target_sum)``. The exact conversion is exact for a -value, not for the model's prediction of one: an MSE-trained decoder predicts a -conditional mean in its own log space, and the concave conversion inflates it for -sparsely expressed genes (RMSE 0.553 against 0.518 for the line on ``bmmc_multiome``). +datasets it came within 0.005 of the best alternative on every metric but MAE, which +favours scButterfly's raw output for staying near the zeros. Each alternative lost +clearly somewhere: one line per gene or peak reorders each cell's profile (per-cell +Spearman -0.023 on ``bmmc_multiome`` GEX -> ATAC), and the exact GEX conversion +``log1p(expm1(p) * 1e4 / target_sum)`` is exact for a value, not for the model's +prediction of one: an MSE-trained decoder predicts a conditional mean in its own log +space, which the concave conversion inflates for sparsely expressed genes (RMSE 0.553 +against 0.518 for the line on ``bmmc_multiome`` ATAC -> GEX). The line is increasing, so every correlation is exactly the model's; it fixes the level and spread that RMSE scores. From f07253331f61bcb4a4f935646a93ec789133775b Mon Sep 17 00:00:00 2001 From: Vladimir Shitov Date: Mon, 5 Oct 2026 23:14:20 +0200 Subject: [PATCH 7/8] scbutterfly: reuse the per-cell scale model of src/utils/cell_scale.py Replace scbutterfly's own scale map (target_scale.py) with CellScale, which ss_opm and senkin_tmp already use: scbutterfly_train fits its ridge model on the training cells and its spread shrinkage on the held-out validation cells (predicted with the saved weights) and saves cell_scale.npz; scbutterfly_predict applies it. Bundles without the file keep scButterfly's own output. Also reuse library functions in chrom_utils: np.argsort/np.unique for the chromosome order and counts, itertools.groupby for run lengths, pd.Index.get_indexer to scatter predictions by name; and share the held-out split through validation_cells(). Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- src/methods/scbutterfly/butterfly_common.py | 24 +++--- src/methods/scbutterfly/chrom_utils.py | 77 +++++-------------- .../scbutterfly/scbutterfly/config.vsh.yaml | 7 +- .../scbutterfly_predict/config.vsh.yaml | 2 +- .../scbutterfly/scbutterfly_predict/script.py | 29 ++++--- .../scbutterfly_train/config.vsh.yaml | 1 + .../scbutterfly/scbutterfly_train/script.py | 38 ++++++++- src/methods/scbutterfly/target_scale.py | 63 --------------- 9 files changed, 89 insertions(+), 154 deletions(-) delete mode 100644 src/methods/scbutterfly/target_scale.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f8e0dcc4..ff46e3ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -60,7 +60,7 @@ * `scbutterfly`: Parse the chromosome of `chr1:1000-2000` peak names, not only `chr1-1000-2000`. NeurIPS 2022 names its peaks the former way, so every peak became its own chromosome, each with its own block in the ATAC encoder: `pbmc_multiome` was OOM-killed when GEX was the target (174,978 "chromosomes") and hit the 8 h limit when ATAC was (7,002 "chromosomes", 5.7 h for one pretraining epoch). -* `scbutterfly`: Model every target feature and predict it on the scale of the target's `normalized` layer. scButterfly kept only the 3,000 most variable genes and the peaks open in at least 0.5% of the cells, also when they were the target, so the other genes and peaks were predicted as zero; and it predicts its own preprocessing of the target. Predictions are now mapped onto the target layer by one least-squares line, fitted on the training cells scButterfly holds out for validation. Apart from MAE, it came within 0.005 of the best alternative on every metric, while one line per feature lowered per-cell Spearman and, for GEX, the exact conversion from scButterfly's median-library-size scaling to log CP10k raised RMSE: the decoder predicts a conditional mean in its own log space, which the exact conversion inflates for sparsely expressed genes. The HVG selection and peak filter still apply to the input modality. `scbutterfly_predict` reproduces the old preprocessing for model bundles trained before this change. +* `scbutterfly`: Model every target feature and predict it on the scale of the target's `normalized` layer. scButterfly kept only the 3,000 most variable genes and the peaks open in at least 0.5% of the cells, also when they were the target, so the other genes and peaks were predicted as zero; and it predicts its own preprocessing of the target. Every predicted cell now gets the level and spread of the target layer from the per-cell scale model of `src/utils/cell_scale.py`, as in `ss_opm` and `senkin_tmp`. `scbutterfly_train` fits its ridge model on the training cells and its spread shrinkage on the cells scButterfly holds out for early stopping, and saves `cell_scale.npz`; `scbutterfly_predict` applies it. The HVG selection and peak filter still apply to the input modality. `scbutterfly_predict` reproduces the old preprocessing for model bundles trained before this change, and writes scButterfly's own output for bundles without `cell_scale.npz`. * `scbutterfly`: Keep the data sparse and densify each batch on the GPU. The model densified both whole modalities (92 GB for the ATAC input of `pbmc_multiome` alone), and its `DataLoader` workers built every batch from dense per-cell rows, so data loading dominated each step. Training on `pbmc_multiome` (ATAC -> GEX) now takes 5.9 h on a V100 instead of 10.0 h. Unused layers and intermediate copies are also dropped early. `scbutterfly_train` asks for `veryhighmem` and `veryhightime`, `scbutterfly_predict` for `veryhighmem`: the model rebuilds scButterfly's preprocessing in both, which peaks at ~100 GB on `pbmc_multiome`, and the T4 the benchmark runs on is slower than the V100. diff --git a/src/methods/scbutterfly/butterfly_common.py b/src/methods/scbutterfly/butterfly_common.py index 358d3be4..1f0b6e62 100644 --- a/src/methods/scbutterfly/butterfly_common.py +++ b/src/methods/scbutterfly/butterfly_common.py @@ -311,6 +311,17 @@ def _concat_blocks(train_block, test_block): return out +def validation_cells(n_train): + """The ~10% of the training cells held out for early stopping (row indices). + + Deterministic, so train and predict hold out the same cells; ``scbutterfly_train`` + also fits the spread of the per-cell scale on them. + """ + shuffled = list(range(n_train)) + np.random.RandomState(0).shuffle(shuffled) + return sorted(shuffled[:max(1, int(0.1 * n_train))]) + + def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly, model_all_target_features=True): """Build + preprocess + construct a Butterfly for the given data. @@ -368,13 +379,8 @@ def build_butterfly(train_mod1, train_mod2, test_mod1, n_top_genes, Butterfly, train_id = list(range(n_train)) test_id = list(range(n_train, n_train + n_test)) - # Deterministic ~10% validation carve (not five_fold_split_dataset). - rng = np.random.RandomState(0) - shuffled = list(train_id) - rng.shuffle(shuffled) - n_val = max(1, int(0.1 * n_train)) - validation_id = sorted(shuffled[:n_val]) - train_id_final = sorted(shuffled[n_val:]) + validation_id = validation_cells(n_train) + train_id_final = sorted(set(train_id) - set(validation_id)) # Chromosome ordering (peaks contiguous per chrom + chrom_list). sort_index, chrom_list = chrom_utils.sorted_chrom_order(ATAC_data) @@ -428,8 +434,8 @@ def predict_cells(built, cell_ids, batch_size, model_path=None): """scButterfly's prediction of the target modality for rows ``cell_ids``. Returned as a dense float32 array in the target's original var order, still on - scButterfly's scale (see ``target_scale``). ``model_path`` loads the trained - weights first; leave it out once they are loaded. + scButterfly's scale (``cell_scale.CellScale`` maps it onto the target's). ``model_path`` + loads the trained weights first; leave it out once they are loaded. """ # Model.test predicts both directions; give the unused one a single cell, which # for ATAC -> GEX on NeurIPS 2022 saves a cells x 175k-peak matrix. It also runs diff --git a/src/methods/scbutterfly/chrom_utils.py b/src/methods/scbutterfly/chrom_utils.py index aa2d0205..0ca55d30 100644 --- a/src/methods/scbutterfly/chrom_utils.py +++ b/src/methods/scbutterfly/chrom_utils.py @@ -11,9 +11,11 @@ scatters a predicted matrix back into a target var order by feature name. """ +import itertools import re import numpy as np +import pandas as pd from scipy.sparse import issparse @@ -33,29 +35,16 @@ def sorted_chrom_order(atac_adata): Returns ------- sort_index : np.ndarray - Indices that reorder ``atac_adata`` so peaks are grouped per chromosome. + Indices that reorder ``atac_adata`` so peaks are grouped per chromosome, in + sorted chromosome order; the stable sort keeps the original order within one. chrom_list : list[int] Number of peaks per chromosome, in the sorted order. ``sum == n_peaks``. """ chroms = np.array([parse_chrom(v) for v in atac_adata.var_names]) - # Stable ordering of chromosomes; stable argsort keeps peaks deterministic - # within a chromosome (preserving original relative order). - chrom_order = sorted(set(chroms)) - rank = {c: i for i, c in enumerate(chrom_order)} - keys = np.array([rank[c] for c in chroms]) - sort_index = np.argsort(keys, kind="stable") - - sorted_chroms = chroms[sort_index] - chrom_list = [] - last = None - for c in sorted_chroms: - if c != last: - chrom_list.append(1) - last = c - else: - chrom_list[-1] += 1 - assert sum(chrom_list) == atac_adata.n_vars - return sort_index, chrom_list + sort_index = np.argsort(chroms, kind="stable") + # np.unique returns the chromosomes sorted, i.e. in the order of sort_index. + _, chrom_list = np.unique(chroms, return_counts=True) + return sort_index, chrom_list.tolist() def chrom_counts(atac_adata): @@ -67,18 +56,10 @@ def chrom_counts(atac_adata): keeps the grouping contiguous, so recounting here yields a valid ``chrom_list``. """ if "chrom" in atac_adata.var.columns: - chroms = list(atac_adata.var["chrom"]) + chroms = atac_adata.var["chrom"] else: - chroms = [parse_chrom(v) for v in atac_adata.var_names] - counts = [] - last = None - for c in chroms: - if c != last: - counts.append(1) - last = c - else: - counts[-1] += 1 - return counts + chroms = map(parse_chrom, atac_adata.var_names) + return [len(list(peaks)) for _, peaks in itertools.groupby(chroms)] def apply_sort(atac_adata, sort_index): @@ -96,35 +77,15 @@ def scatter_to_target(pred_adata, target_var_names): """Scatter a prediction into ``target_var_names`` order, by feature name. Predicted features may be a subset of and/or in a different order from the - target modality's vars (scButterfly can subset RNA to HVGs or filter ATAC - peaks). Scattering by name simultaneously (a) restores the original peak - order and (b) fills any missing target features with zeros. + target modality's vars (scButterfly sorts peaks by chromosome, and models trained + before every target feature was modelled subset genes and peaks). Scattering by + name restores the original order and fills missing target features with zeros. Returns a dense float32 ndarray of shape ``(n_cells, len(target_var_names))``. """ - X = pred_adata.X - if issparse(X): - X = X.toarray() - X = np.asarray(X, dtype=np.float32) - - n_cells = X.shape[0] - out = np.zeros((n_cells, len(target_var_names)), dtype=np.float32) - - target_pos = {name: i for i, name in enumerate(target_var_names)} - src_cols, dst_cols = [], [] - for src_col, name in enumerate(pred_adata.var_names): - dst = target_pos.get(name) - if dst is not None: - src_cols.append(src_col) - dst_cols.append(dst) - - # Copy in column blocks: one fancy-index assignment per block rather than one - # per feature (up to ~116k of them for ATAC), while keeping the temporary that - # fancy indexing allocates bounded. - src_cols = np.asarray(src_cols, dtype=np.intp) - dst_cols = np.asarray(dst_cols, dtype=np.intp) - block = 4096 - for start in range(0, src_cols.size, block): - sl = slice(start, start + block) - out[:, dst_cols[sl]] = X[:, src_cols[sl]] + X = pred_adata.X.toarray() if issparse(pred_adata.X) else np.asarray(pred_adata.X) + source_columns = pd.Index(pred_adata.var_names).get_indexer(target_var_names) + modelled = source_columns >= 0 + out = np.zeros((X.shape[0], len(target_var_names)), dtype=np.float32) + out[:, modelled] = X[:, source_columns[modelled]] return out diff --git a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml index a2714fc2..7a513a67 100644 --- a/src/methods/scbutterfly/scbutterfly/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly/config.vsh.yaml @@ -9,9 +9,10 @@ description: | the held-out test modality. Chromosome grouping for the ATAC branch is parsed from peak coordinates in the feature names. scButterfly's feature selection (highly variable genes, peak filtering) applies to the input modality only, so every target feature is modelled. - Predictions are mapped from scButterfly's own preprocessing onto the target layer by one - least-squares line fitted on held-out training cells. Only Multiome GEX<->ATAC is - supported; CITE-seq (ADT) datasets are not handled. + scButterfly predicts its own preprocessing of the target, so every predicted cell is given + the level and spread of the target layer with the per-cell scale model shared with ss_opm + and senkin_tmp (src/utils/cell_scale.py). Only Multiome GEX<->ATAC is supported; CITE-seq + (ADT) datasets are not handled. references: doi: - 10.1038/s41467-024-47418-x diff --git a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml index 9bba2e48..9d668a7e 100644 --- a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml @@ -12,8 +12,8 @@ resources: type: python_script - path: ../butterfly_common.py - path: ../chrom_utils.py - - path: ../target_scale.py - path: /src/utils/exit_codes.py + - path: /src/utils/cell_scale.py test_resources: - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap diff --git a/src/methods/scbutterfly/scbutterfly_predict/script.py b/src/methods/scbutterfly/scbutterfly_predict/script.py index 43d1c4d3..c2983dc0 100644 --- a/src/methods/scbutterfly/scbutterfly_predict/script.py +++ b/src/methods/scbutterfly/scbutterfly_predict/script.py @@ -4,6 +4,7 @@ import sys import anndata as ad +import numpy as np from scipy.sparse import csc_matrix logging.basicConfig(level=logging.INFO) @@ -22,7 +23,7 @@ sys.path.append(meta["resources_dir"]) import butterfly_common -import target_scale +from cell_scale import CellScale butterfly_common.apply_runtime_patches() from scButterfly.butterfly import Butterfly @@ -32,9 +33,10 @@ # --------------------------------------------------------------------------- logger.info("Reading input files...") train_mod1 = butterfly_common.read_modality(par["input_train_mod1"]) -# The target's normalized layer is what the scale map is fitted to. -train_mod2 = butterfly_common.read_modality(par["input_train_mod2"], layers=("counts", "normalized")) -test_mod1 = butterfly_common.read_modality(par["input_test_mod1"]) +train_mod2 = butterfly_common.read_modality(par["input_train_mod2"]) +test_mod1 = butterfly_common.read_modality(par["input_test_mod1"], layers=("counts", "normalized")) +# The per-cell scale reads the test inputs; build_butterfly takes the counts out of test_mod1. +test_inputs, test_counts = test_mod1.layers.pop("normalized"), test_mod1.layers["counts"] with open(os.path.join(par["input_model"], "metadata.pkl"), "rb") as f: metadata = pickle.load(f) @@ -57,16 +59,13 @@ built, built["test_id"], batch_size, model_path=par["input_model"], ) -# scButterfly predicts its own preprocessing of the target; map it onto the scale of -# the target's normalized layer with a line fitted on the held-out training cells -# (see target_scale). -logger.info("Fitting the scale map on %d held-out training cells...", len(built["validation_id"])) -calibration = target_scale.ScaleCalibration().fit( - butterfly_common.predict_cells(built, built["validation_id"], batch_size), - train_mod2.layers["normalized"][built["validation_id"]], -) -logger.info("Target = %.4f + %.4f x scButterfly output", calibration.intercept_, calibration.slope_) -test_predictions = calibration.apply(test_predictions) +# scButterfly predicts its own preprocessing of the target; give every cell the level +# and spread of the target's normalized layer (fitted in scbutterfly_train). +cell_scale_path = os.path.join(par["input_model"], "cell_scale.npz") +if os.path.exists(cell_scale_path): + test_predictions = CellScale.load(cell_scale_path).apply(test_predictions, test_inputs, test_counts) +else: + logger.warning("The model has no cell_scale.npz (trained before the per-cell scale): writing scButterfly's own output.") # --------------------------------------------------------------------------- # Write predictions. @@ -81,7 +80,7 @@ logger.info("Writing predictions...") adata_out = ad.AnnData( - layers={"normalized": csc_matrix(test_predictions)}, + layers={"normalized": csc_matrix(test_predictions.astype(np.float32))}, obs=test_mod1.obs, var=train_mod2.var, uns={ diff --git a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml index 4e7257a0..f3542921 100644 --- a/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly_train/config.vsh.yaml @@ -17,6 +17,7 @@ resources: - path: ../butterfly_common.py - path: ../chrom_utils.py - path: /src/utils/exit_codes.py + - path: /src/utils/cell_scale.py test_resources: - path: /resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap dest: resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap diff --git a/src/methods/scbutterfly/scbutterfly_train/script.py b/src/methods/scbutterfly/scbutterfly_train/script.py index 0c8ce853..b9425829 100644 --- a/src/methods/scbutterfly/scbutterfly_train/script.py +++ b/src/methods/scbutterfly/scbutterfly_train/script.py @@ -3,7 +3,7 @@ import pickle import sys -import anndata as ad +from scipy.sparse import csr_matrix logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -26,6 +26,7 @@ sys.path.append(meta["resources_dir"]) import butterfly_common +from cell_scale import CellScale butterfly_common.apply_runtime_patches() from scButterfly.butterfly import Butterfly @@ -34,10 +35,29 @@ # Load data # --------------------------------------------------------------------------- logger.info("Reading input files...") -train_mod1 = butterfly_common.read_modality(par["input_train_mod1"]) -train_mod2 = butterfly_common.read_modality(par["input_train_mod2"]) +train_mod1 = butterfly_common.read_modality(par["input_train_mod1"], layers=("counts", "normalized")) +train_mod2 = butterfly_common.read_modality(par["input_train_mod2"], layers=("counts", "normalized")) test_mod1 = butterfly_common.read_modality(par["input_test_mod1"]) +# --------------------------------------------------------------------------- +# Per-cell scale of the predictions (src/utils/cell_scale.py, as in ss_opm and +# senkin_tmp). scButterfly predicts its own preprocessing of the target, not the +# target's normalized layer; CellScale gives every predicted cell the level and spread +# of the target. Its ridge model is fitted on all training cells here, its spread on +# the cells held out for early stopping once the model is trained. build_butterfly +# takes the counts out of the inputs, so slice what the spread needs first. +# --------------------------------------------------------------------------- +logger.info("Fitting the per-cell scale model...") +train_inputs = csr_matrix(train_mod1.layers.pop("normalized")) +train_targets = csr_matrix(train_mod2.layers.pop("normalized")) +train_counts = csr_matrix(train_mod1.layers["counts"]) +cell_scale = CellScale().fit(train_inputs, train_targets, train_counts) +validation_id = butterfly_common.validation_cells(train_mod1.n_obs) +validation_inputs, validation_targets, validation_counts = ( + matrix[validation_id] for matrix in (train_inputs, train_targets, train_counts) +) +del train_inputs, train_targets, train_counts + # --------------------------------------------------------------------------- # Build + construct the model, then train (weights written to /model). # --------------------------------------------------------------------------- @@ -60,9 +80,19 @@ output_path=par["output"], ) +# Spread of the per-cell scale, from the held-out cells predicted with the weights +# train_model saved, which are the ones scbutterfly_predict loads. +validation_predictions = butterfly_common.predict_cells( + built, validation_id, par["batch_size"], model_path=par["output"], +) +cell_scale.fit_spread(validation_predictions, validation_inputs, validation_targets, validation_counts) +logger.info("Per-cell scale: spread shrinkage %.4f", cell_scale.spread_) +cell_scale.save(os.path.join(par["output"], "cell_scale.npz")) + # --------------------------------------------------------------------------- # Persist the metadata predict needs to reconstruct the model deterministically. -# The trained weights live in /model/*.pt (written by train_model). +# The trained weights live in /model/*.pt (written by train_model), the +# per-cell scale in /cell_scale.npz. # --------------------------------------------------------------------------- logger.info("Saving model metadata...") metadata = { diff --git a/src/methods/scbutterfly/target_scale.py b/src/methods/scbutterfly/target_scale.py deleted file mode 100644 index 93ccfa69..00000000 --- a/src/methods/scbutterfly/target_scale.py +++ /dev/null @@ -1,63 +0,0 @@ -"""Map scButterfly's predictions onto the scale of the target's ``normalized`` layer. - -scButterfly predicts its own preprocessing of the target modality, which is not the -layer the benchmark scores: - -- GEX: ``log1p`` of each cell's counts scaled to the median count total of the cells - (scanpy's ``normalize_total`` default) instead of 1e4. -- ATAC: binarized peaks, TF-IDF-transformed and divided by the global maximum, - predicted through a sigmoid (0 to 1), where the dataset has its own TF-IDF (0 to ~20). - -:class:`ScaleCalibration` maps them with one least-squares line, fitted on training -cells scButterfly held out for validation. On the NeurIPS 2021 and 2022 Multiome -datasets it came within 0.005 of the best alternative on every metric but MAE, which -favours scButterfly's raw output for staying near the zeros. Each alternative lost -clearly somewhere: one line per gene or peak reorders each cell's profile (per-cell -Spearman -0.023 on ``bmmc_multiome`` GEX -> ATAC), and the exact GEX conversion -``log1p(expm1(p) * 1e4 / target_sum)`` is exact for a value, not for the model's -prediction of one: an MSE-trained decoder predicts a conditional mean in its own log -space, which the concave conversion inflates for sparsely expressed genes (RMSE 0.553 -against 0.518 for the line on ``bmmc_multiome`` ATAC -> GEX). - -The line is increasing, so every correlation is exactly the model's; it fixes the level -and spread that RMSE scores. -""" - -import numpy as np -import scipy.sparse as sp - - -class ScaleCalibration: - """One least-squares line from scButterfly's output to the target layer.""" - - def fit(self, predictions, targets, block_size=1000): - """Fit on ``predictions`` (cells x features, dense) and ``targets`` (dense or sparse). - - If the held-out cells show no positive association, which only an essentially - untrained model produces (e.g. two epochs on the test resources), the line - matches the target's mean and standard deviation instead, so the predictions - are rescaled rather than replaced by a constant. - """ - n_entries = predictions.size - prediction_sum, prediction_square_sum = 0.0, 0.0 - for start in range(0, predictions.shape[0], block_size): - block = predictions[start:start + block_size].astype(np.float64) - prediction_sum += block.sum() - prediction_square_sum += (block**2).sum() - targets = sp.csr_matrix(targets, dtype=np.float64) - prediction_mean = prediction_sum / n_entries - target_mean = targets.sum() / n_entries - covariance = targets.multiply(predictions).sum() / n_entries - prediction_mean * target_mean - variance = prediction_square_sum / n_entries - prediction_mean**2 - target_variance = (targets.data**2).sum() / n_entries - target_mean**2 - if variance <= 0: - self.slope_ = 0.0 - elif covariance > 0: - self.slope_ = covariance / variance - else: - self.slope_ = np.sqrt(max(target_variance, 0.0) / variance) - self.intercept_ = target_mean - self.slope_ * prediction_mean - return self - - def apply(self, predictions): - return (np.asarray(predictions) * self.slope_ + self.intercept_).astype(np.float32) From 78eb776ae98afd9efff658f7725245500ab813eb Mon Sep 17 00:00:00 2001 From: Vladimir Shitov <35199218+VladimirShitov@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:31:53 +0200 Subject: [PATCH 8/8] Use a bigger gpu node --- src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml index 9d668a7e..c4289f5f 100644 --- a/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml +++ b/src/methods/scbutterfly/scbutterfly_predict/config.vsh.yaml @@ -33,4 +33,4 @@ runners: - type: executable - type: nextflow directives: - label: [veryhighmem, hightime, midcpu, highsharedmem, gpu] + label: [veryhighmem, hightime, midcpu, highsharedmem, biggpu]