In-cohort Cell Phenotyping

The pipeline run_cell_probe_tasks evaluates the quality of cell-level representations with in-cohort linear probes for cell phenotyping. The model computes a token per cell; a logistic regression trained on the standardized tokens of the training patients then predicts the cell labels of the test patients.

Running the Benchmark

python -m spora_bench.benchmarks.run_cell_probe_tasks \
  model_config=<path-to-model-config> \
  datasets_config=<path-to-datasets-config> \
  benchmarks_config=<path-to-benchmarks-config>

For example, VirTues on hoch2022multiplexed:

python -m spora_bench.benchmarks.run_cell_probe_tasks \
  model_config=configs/models/virtues.yaml \
  datasets_config=configs/datasets/virtues/hoch2022multiplexed.yaml \
  benchmarks_config=configs/benchmarks/cell_probe/hoch2022multiplexed.yaml

Some models, for instance KRONOS, require distinct dataset configs to account for differences in standardization; for KRONOS, use configs/datasets/kronos/.

Outputs

Cell tokens are cached in output_dir/<model_name>/cell_tokens/<dataset>.h5ad and reused by later runs. For each task, output_dir/<model_name>/results/ receives a classification report, a bootstrapped classification report (1,000 resamples), and a confusion matrix as Parquet files.

Benchmark Config

datasets:
  <dataset_name>:
    benchmarks:
      cell_probe:
        <task_name>:
          label_col: <column_name>
          excluded_classes: <list of classes to exclude>

label_col is a column of metadata/cells.parquet. Cells with a missing label or a class listed in excluded_classes are removed before training and evaluation.

The datasets and configs used in the paper are listed in the Task Zoo.