Restore the per-cell scale of senkin_tmp and ss_opm predictions (#80)

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: Vladimir Shitov <vladimir.shitov@helmholtz-munich.de>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Author: 3 people

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)