Restore the per-cell scale of senkin_tmp and ss_opm predictions

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

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

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

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

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

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)

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

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 ----

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

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)

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