bioptimus.data.cohort - Bioptimus Documentation

Cohort manifest for pairing WSIs with bulk RNA and clinical metadata. Provides Cohort, a typed registry that tracks the mapping between patients, wsis, bulk RNA samples, timepoints, labels, and arbitrary clinical metadata. WSIRecord is the single source of truth for each WSI — including per-model outputs and processing status. Construction paths:

From a user-provided CSV via Cohort.from_csv.
Auto-matching from directories via Cohort.from_directories.
From a YAML manifest via Cohort.load (for resume).

The cohort can be serialised to YAML via Cohort.save and loaded back for full reproducibility and resume support. Example:

from bioptimus.data.cohort import Cohort

# Auto-match from directories:
cohort = Cohort.from_directories(
    wsi_dir="/data/wsis",
    bulk_rna_dir="/data/rna",
)

# Or load a saved manifest (resume):
cohort = Cohort.load(".output/brca/run_1/manifest.yaml")

# Add labels post-hoc:
cohort.add_labels({"patient_001": {"mutation": "BRAF"}})

# Iterate:
for record in cohort:
    print(record.patient_id, record.wsi_path)

ModelOutput

@dataclass
class ModelOutput()

Per-model output record for a single WSI. Attributes:

model_name - Name of the model that produced this output.
embedding_path - Path to the embedding output file.
prediction_path - Path to the prediction output file.
tiles_csv_path - Path to the tile coordinates CSV.
tiles_dir - Directory containing exported tile images.
modalities - Input modalities used for this output (e.g. ["image"] or ["image", "bulk_rna"]).
timestamp - ISO 8601 string of when the output was produced.
status - Processing status ("pending", "done", "error").
error - Error message if status is "error".
base_dir - Manifest directory that relative output paths are anchored to for existence checks. Set by Cohort.load and never serialised. None resolves relative paths against the current working directory.

is_done

@property
def is_done() -> bool

Returns True if status is done.

has_embedding

@property
def has_embedding() -> bool

Returns True if an embedding path is set and exists.

has_prediction

@property
def has_prediction() -> bool

Returns True if a prediction path is set and exists.

set_stage_output

def set_stage_output(stage: InferenceMode | str,
                     modalities: list[str],
                     path: Path,
                     timestamp: str | None = None) -> None

Records an output for a specific stage and modality combo. Updates both the per-modality tracking dict and the top-level convenience fields (embedding_path / prediction_path). The combined InferenceMode.PREDICTION_WITH_EMBEDDING stage records the single output file under both fields, since it carries the prediction and the embedding together.

The output stage as an InferenceMode (or a mode string; legacy short spellings are accepted).

Modalities used for this output.

Output file path.

ISO 8601 timestamp.

get_stage_path

def get_stage_path(stage: InferenceMode | str,
                   modalities: list[str] | None = None) -> Path | None

Returns the output path for a stage and modality combo. For the combined stage both path fields point at the same file, so the embedding field is returned.

The output stage as an InferenceMode (or a mode string; legacy short spellings are accepted).

Modality combination to look up. Falls back to the top-level path when None.

Returns: Path resolved against base_dir, or None if not recorded.

to_dict

def to_dict() -> dict[str, Any]

Serialises to a plain dict for YAML output. Returns: A dictionary representation of the model output.

from_dict

@classmethod
def from_dict(cls,
              data: dict[str, Any],
              base_dir: Path | None = None) -> ModelOutput

Constructs from a plain dict (YAML round-trip).

Dictionary with serialised model output fields.

Manifest directory to anchor relative paths to for existence checks.

Returns: A new ModelOutput instance.

WSIRecord

@dataclass
class WSIRecord()

Single WSI entry — the single source of truth. Tracks inputs (WSI, RNA, mask), per-model outputs, and arbitrary annotations for one WSI within a cohort. Attributes:

patient_id - Unique patient identifier.
wsi_id - Unique WSI identifier (typically the filename stem).
wsi_path - Path to the WSI file.
bulk_rna_id - Identifier for the paired bulk RNA sample.
bulk_rna_path - Path to the bulk RNA CSV/TSV file.
mask_path - Path to a tissue mask (pre-computed or cached).
mask_timestamp - When the mask was computed/discovered.
timepoint - Temporal ordering label (e.g. "t0", "t1").
outputs - Per-model output records keyed by model name.
labels - Arbitrary categorical annotations.
metadata - Arbitrary clinical/treatment metadata.
base_dir - Manifest directory that relative paths are anchored to for existence checks. Set by Cohort.load and never serialised. None resolves relative paths against the current working directory.

get_output

def get_output(model_name: str) -> ModelOutput

Returns the output for model_name, creating if absent.

Model identifier string.

Returns: The ModelOutput for the given model.

is_stage_done

def is_stage_done(model_name: str, stage: InferenceMode | str) -> bool

Checks whether a stage output exists on disk. When modality_outputs are tracked, the check is scoped to the record’s current available_modalities so that linking new data (e.g. bulk RNA) automatically surfaces pending work without requiring force=True. The combined stage requires both its path fields, so it is done only once both outputs exist.

Model identifier.

The output stage as an InferenceMode (or a mode string; legacy short spellings are accepted).

Returns: True if the stage’s output path(s) exist on disk.

has_mask

@property
def has_mask() -> bool

Returns True if a mask path is set and exists on disk.

merge_saved_outputs

def merge_saved_outputs(saved: WSIRecord) -> int

Merges output/mask metadata from a saved (manifest) record. Used to resume after a restart: outputs recorded in a previous manifest are folded into this live record, which may already have some outputs from the current process. WSI/RNA paths are left untouched so the user’s current directory layout wins. Adopts saved.base_dir unconditionally so that relative paths merged from saved keep resolving against the manifest directory rather than the process CWD, regardless of whether a mask is present (existence checks read the record’s base_dir, not the output’s).

The corresponding record loaded from a manifest.

Returns: Number of output entries merged in (new models plus new modality combinations added to an existing model output).

available_modalities

@property
def available_modalities() -> list[str]

Returns the input modalities available for this record. Always includes "image". Includes "bulk_rna" when bulk_rna_path is set.

to_dict

def to_dict() -> dict[str, Any]

Serialises to a plain dict for YAML output. Returns: A dictionary representation of the WSI record.

from_dict

@classmethod
def from_dict(cls,
              data: dict[str, Any],
              base_dir: Path | None = None) -> WSIRecord

Constructs from a plain dict (YAML round-trip).

Dictionary with serialised WSI record fields.

Manifest directory to anchor relative paths to for existence checks.

Returns: A new WSIRecord instance.

PatientRecord

@dataclass
class PatientRecord()

Groups all WSIs belonging to a single patient. Attributes:

patient_id - Unique patient identifier.
wsis - Ordered list of WSI records (by timepoint).
labels - Patient-level categorical annotations.
metadata - Patient-level clinical metadata.

Cohort

class Cohort(records: list[WSIRecord] | None = None)

Typed cohort manifest — single source of truth for an experiment. A cohort is an ordered collection of WSIRecord entries grouped by patient. It stores all information needed to reproduce and resume an inference run: file locations, pairing logic, timepoints, per-model outputs, labels, and clinical metadata. Construction: Use from_csv, from_directories, or load rather than calling the constructor directly.

Initial list of WSIRecord entries.

from_csv

@classmethod
@log_failures
def from_csv(cls,
             path: str | Path,
             *,
             wsi_dir: str | Path | None = None,
             bulk_rna_dir: str | Path | None = None,
             columns: dict[str, str] | None = None) -> Cohort

Creates a cohort from a user-provided CSV manifest. Required columns:

wsi_id: WSI identifier (filename stem or full name).

Optional columns:

patient_id: if absent, derived from wsi_id.
bulk_rna_id: paired RNA sample identifier.
timepoint: temporal label.
wsi_path: explicit path override.
bulk_rna_path: explicit path override.
Any other columns are treated as labels if prefixed with label_ or as metadata otherwise.

When wsi_dir or bulk_rna_dir are provided, paths are resolved by joining the directory with the corresponding ID (plus matching extension found on disk). The columns dict maps canonical field names to the actual CSV column names. Only the keys that differ from the defaults need to be specified. For example: Recognised canonical keys: wsi_id, patient_id, bulk_rna_id, timepoint, wsi_path, bulk_rna_path.

Cohort.from_csv(
    "cohort.csv",
    columns={
        "wsi_id": "imageid",
        "patient_id": "subject_id",
    },
)

Path to the CSV manifest.

Directory to resolve WSI paths from.

Directory to resolve bulk RNA paths from.

Mapping of canonical field names to actual CSV column names. Unmapped fields fall back to their canonical name.

Returns: A populated Cohort.

from_directories

@classmethod
@log_failures
def from_directories(
        cls,
        wsi_dir: str | Path,
        bulk_rna_dir: str | Path | None = None,
        *,
        mask_dir: str | Path | None = None,
        patient_id_fn: Callable[[str], str] | None = None) -> Cohort

Creates a cohort by auto-matching files from directories. WSI and bulk RNA files are paired by exact filename stem equality (case-sensitive). Patient IDs default to the shared stem. When multiple wsis share the same patient ID, they are assigned sequential timepoints (t0, t1, …) in alphabetical order.

Directory containing WSI files.

Directory containing bulk RNA files. When None, wsis are registered without RNA.

Directory containing pre-computed masks (PNG files whose stem matches the WSI stem).

Optional callable that maps a filename stem to a patient ID.

Returns: A populated Cohort.

save

@log_failures
def save(path: str | Path) -> Path

Persists the cohort manifest to a YAML file. The YAML includes all inputs, per-model outputs, status, labels, and metadata — everything needed to resume.

Destination file path.

Returns: The resolved output path.

load

@classmethod
@log_failures
def load(cls, path: str | Path) -> Cohort

Loads a cohort from a previously saved YAML manifest.

Path to the YAML manifest file.

Returns: A fully-populated Cohort.

upsert

def upsert(wsi_id: str,
           *,
           patient_id: str | None = None,
           wsi_path: Path | None = None,
           bulk_rna_path: Path | None = None,
           mask_path: Path | None = None,
           timepoint: str | None = None) -> WSIRecord

Inserts or updates a WSI record. If a record with wsi_id exists, updates the non-None fields. Otherwise creates a new record.

WSI identifier.

Patient identifier (defaults to wsi_id).

Path to WSI file.

Path to bulk RNA file.

Path to tissue mask.

Timepoint label.

Returns: The inserted or updated WSIRecord.

get_wsi

def get_wsi(wsi_id: str) -> WSIRecord | None

Returns the record for wsi_id, or None.

Unique WSI identifier (typically the file stem).

Returns: The matching WSIRecord, or None if not found.

link_bulk_rna

def link_bulk_rna(bulk_rna_dir: str | Path,
                  *,
                  extensions: set[str] | None = None,
                  separator: str | None = None,
                  gene_column: str | None = None,
                  value_column: str | None = None,
                  strip_version: bool = False) -> int

Links bulk RNA files to existing WSI records by stem. Scans bulk_rna_dir for recognised RNA files and pairs them with existing records by exact filename-stem equality. Records that already have a bulk_rna_path are skipped. This enables late-binding of bulk RNA data: build a cohort from WSIs first, then call this method to attach RNA when it becomes available.

Directory containing bulk RNA files.

File extensions to match. Defaults to {".csv", ".tsv"}.

Column delimiter for parsing. When None the separator is inferred from the file extension.

Column name containing gene identifiers. Required (with value_column) for long-format files (e.g. GDC/TCGA gene quantification TSVs).

Column name containing expression values to read (e.g. "tpm_unstranded").

When True, strips version suffixes from gene identifiers (e.g. ENSG…00003.15 → ENSG…00003).

Returns: Number of records that were linked.

add_labels

def add_labels(labels: dict[str, dict[str, Any]],
               *,
               by: str = "patient_id") -> None

Adds categorical labels to records.

Mapping of {identifier: {label_name: value}}.

Key to match on — "patient_id" or "wsi_id".

add_metadata

def add_metadata(metadata: dict[str, dict[str, Any]],
                 *,
                 by: str = "patient_id") -> None

Adds clinical/treatment metadata to records.

Mapping of {identifier: {field: value}}.

Key to match on — "patient_id" or "wsi_id".

add_labels_from_csv

@log_failures
def add_labels_from_csv(path: str | Path, *, by: str = "patient_id") -> None

Adds labels and metadata from a CSV file. The CSV must have a column matching by (default patient_id). Columns prefixed with label_ are treated as labels; remaining columns as metadata.

Path to the labels CSV.

Join key column name.

num_patients

@property
def num_patients() -> int

Number of unique patients in the cohort.

patients

@property
def patients() -> dict[str, PatientRecord]

Patient-indexed view of the cohort.

get_patient

def get_patient(patient_id: str) -> PatientRecord

Returns the PatientRecord for a given patient.

Patient identifier string.

Returns: The matching PatientRecord. Raises:

KeyError - If patient_id is not found.

wsi_ids

@property
def wsi_ids() -> list[str]

Ordered list of all WSI IDs.

patient_ids

@property
def patient_ids() -> list[str]

Ordered list of unique patient IDs.

wsi_paths

@property
def wsi_paths() -> list[Path | None]

Ordered list of WSI paths (may contain None).

pending

def pending(model_name: str, stage: InferenceMode | str) -> list[WSIRecord]

Returns wsis that have not completed a stage for a model. Checks whether the output file exists on disk. Records with missing WSI paths are excluded.

Model identifier.

The output stage as an InferenceMode (or a mode string; legacy short spellings are accepted).

Returns: List of WSIRecord entries still needing processing.

to_csv

@log_failures
def to_csv(path: str | Path) -> Path

Writes input columns to a CSV file (labels + metadata). For full persistence including outputs, use save.

Destination file path.

Returns: The resolved output path.

summary

def summary() -> str

Returns a human-readable summary of the cohort. Returns: Multi-line summary string.

​ModelOutput

​is_done

​has_embedding

​has_prediction

​set_stage_output

​get_stage_path

​to_dict

​from_dict

​WSIRecord

​get_output

​is_stage_done

​has_mask

​merge_saved_outputs

​available_modalities

​to_dict

​from_dict

​PatientRecord

​Cohort

​from_csv

​from_directories

​save

​load

​upsert

​get_wsi

​link_bulk_rna

​add_labels

​add_metadata

​add_labels_from_csv

​num_patients

​patients

​get_patient

​wsi_ids

​patient_ids

​wsi_paths

​pending

​to_csv

​summary

ModelOutput

is_done

has_embedding

has_prediction

set_stage_output

get_stage_path

to_dict

from_dict

WSIRecord

get_output

is_stage_done

has_mask

merge_saved_outputs

available_modalities

to_dict

from_dict

PatientRecord

Cohort

from_csv

from_directories

save

load

upsert

get_wsi

link_bulk_rna

add_labels

add_metadata

add_labels_from_csv

num_patients

patients

get_patient

wsi_ids

patient_ids

wsi_paths

pending

to_csv

summary