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