diff --git a/README.md b/README.md index d092e80..282e4c9 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,7 @@ Its `"table"` AnnData must contain: | obs column | content | | --- | --- | | `cell_type` (configurable) | cell-type label for each cell | -| `spatial_group` (configurable) | binary spatial region label (e.g. `0` = outside tumour, `1` = inside tumour) | +| `spatial_group` (configurable) | spatial region label with two values to compare (e.g. `0`/`1`, or `"out_of_tumor"`/`"in_tumor"`). Which value is the target is chosen in Step 3 with `--spatial-group-target` / `--spatial-group-reference`, defaulting to `1` / `0` | | `center_x`, `center_y` | cell centroid in microns | The zarr must also expose the following SpatialData elements, used to render the per-cell annotation panels (Step 1): @@ -90,6 +90,12 @@ python scripts/export.py \ --layer counts ``` +`--annotation-mode` selects the actions offered in Step 2, and defaults to +`accept_correct_reject`. Use `--annotation-mode accept_reject` to drop the relabelling +option. The value is saved to `config.json`; because the cell-type vocabulary is written +there too (always, whatever the mode), you can switch modes afterwards by editing +`config.json`, without re-exporting the panels. + `--target-proportion` controls the fraction of cells of interest in the subsample. Cells of interest are upweighted accordingly (importance sampling); the unnormalized weight for each sampled cell is stored in `metadata.csv` for downstream use. `--layer` selects which expression matrix to read: the named `.layers` entry holding the raw counts (e.g. `counts`), or `.X` when omitted. The value is saved to `config.json` and reused throughout the workflow — the same layer feeds the top-gene panels here in Step 1 and the CSDE model in Step 3, so set it once at export time. **It must point at raw counts**, since the noise model (Poisson / negative binomial) assumes integer counts; pointing it at normalised or log-transformed values will produce invalid results. @@ -101,9 +107,9 @@ The script writes: ├── images/ │ ├── cell_.png # one panel per cell │ └── ... -├── config.json # all export arguments (read by annotate.py) +├── config.json # export arguments + cell_type_vocabulary (read by annotate.py) ├── metadata.csv # cell_id, cell_type, image_path, sampling_weight, center_x, center_y -└── annotations.json # {cell_id: true/false} — written by annotate.py +└── annotations.json # {cell_id: {action, label}} — written by annotate.py ``` Each panel contains: @@ -127,12 +133,27 @@ A simple JSON mapping gene names to colours: ## Step 2 — Manual validation (`scripts/annotate.py`) -For each exported image, an annotator decides whether the cell is **correct** — meaning it is both properly **segmented** and properly **labelled**. A cell should be rejected (marked incorrect) when either check fails: +For each exported image, the annotator runs two checks in order: + +1. **Segmentation** — is the cell boundary (left panel) consistent with the nuclei / membrane staining, or does it merge two cells or clip part of one? +2. **Cell-type label** — are the top expressed genes (right panel) consistent with the assigned label? + +which lead to one of three actions: + +| action | when | effect | +| --- | --- | --- | +| **accept** | segmentation fine, label fine | the cell keeps its automated label | +| **correct** | segmentation fine, label wrong | the annotator picks the right cell type | +| **reject** | segmentation inadequate | the cell is excluded from both compared groups | -- **Segmentation** — the cell boundary (left panel) is not consistent with the nuclei / membrane staining, e.g. it merges two cells or clips part of one. -- **Cell-type label** — the top expressed genes (right panel) include genes unlikely to be expressed by the cell type of interest, suggesting the automated label is wrong. +Correcting a cell revises only its **cell type**; its spatial region is treated as reliable +and is always taken from the automated pipeline. So correcting a cell *into* the cell type +of interest is what places it in the target or reference group, according to the region it +already sits in — this is the case an accept/reject workflow cannot express. -The result is a boolean column `is_correct` added to `metadata.csv`, which becomes `adata_gt` in Step 3. +Segmentation is never edited: an accepted or corrected cell keeps the automated expression +counts. Rejection therefore doubles as a quality-control filter for cells whose +quantification cannot be trusted at all. ```bash streamlit run scripts/annotate.py -- --dir /path/to/annotation_dir @@ -142,10 +163,19 @@ The `--` is required: it tells Streamlit to pass everything after it to the scri VS Code Remote forwards the Streamlit port automatically. Open the URL printed in the terminal, then use: -- **`1`** — label as correct -- **`2`** — label as incorrect +| key | `accept_correct_reject` (default) | `accept_reject` | +| --- | --- | --- | +| **`1`** | accept | accept | +| **`2`** | correct | reject | +| **`3`** | reject | — | -Progress is saved after every keypress to `annotations.json`. Re-running the command resumes from where you left off. You can also start annotating while `export.py` is still running — the UI picks up newly exported cells automatically. +Pressing **`2`** in `accept_correct_reject` mode opens a cell-type selector below the panel +— type a few characters to filter, then pick the label. Nothing is written until you +choose one, so pressing `2` by mistake is harmless: hit Cancel and the cell stays +unannotated. + +Progress is saved after every keypress to `annotations.json`, as +`{cell_id: {"action": ..., "label": ...}}` (`label` is set only for corrections). Re-running the command resumes from where you left off. You can also start annotating while `export.py` is still running — the UI picks up newly exported cells automatically. --- @@ -157,11 +187,27 @@ python scripts/differential_expression.py --dir /path/to/annotation_dir Reads all export settings from `config.json` and writes gene-level results to `/results.csv`. +The three-way comparison is built here: cells of interest in spatial group `0` (reference) +and group `1` (target) form the two compared populations, and everything else — including +rejected cells — is collapsed into a third group. Both the automated labels and the manual +ones are built the same way; only the cell type differs between them. The script prints a +summary of the annotations first (counts per action, plus how many cells the curation moved +into and out of the compared groups), which is the quickest check that the annotations were +read as intended. + +If your region column is not encoded as `1` / `0`, set `--spatial-group-target` and +`--spatial-group-reference` to the two values you want to compare; the script reports the +values it found if they don't match. The target region is the one positive log-fold changes +refer to, so swapping the two flips the sign of every result — this is deliberately not +inferred for you, even when the column has exactly two values. + | option | default | description | |---|---|---| | `--dir` | *(required)* | annotation directory (output of steps 1 & 2) | | `--out` | `/results.csv` | output CSV path | | `--spatial-group-key` | `spatial_group` | obs column encoding the two spatial populations | +| `--spatial-group-target` | `1` | value of that column identifying the target region (group 1) | +| `--spatial-group-reference` | `0` | value of that column identifying the reference region (group 0) | | `--n-cells-expressed-threshold` | `10` | min annotated cells expressing a gene for it to be tested | | `--noise-model` | `poisson` | `poisson` or `nb` (negative binomial) | diff --git a/pyproject.toml b/pyproject.toml index 529d2f6..b17a9be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dev = [ "isort", "flake8", ] -annotate = ["streamlit"] +annotate = ["streamlit>=1.30"] cuda12 = ["jax[cuda12]"] cuda13 = ["jax[cuda13]"] diff --git a/scripts/annotate.py b/scripts/annotate.py index 8419020..bcb6c97 100644 --- a/scripts/annotate.py +++ b/scripts/annotate.py @@ -4,6 +4,18 @@ Usage ----- streamlit run scripts/annotate.py -- --dir /path/to/annotations/R2_macrophages + +The available actions are set by ``annotation_mode`` in config.json (written by +export.py): + +- ``accept_reject`` 1 = accept, 2 = reject +- ``accept_correct_reject`` 1 = accept, 2 = correct, 3 = reject + +Writes ``annotations.json`` as +``{cell_id: {"action": "accept" | "correct" | "reject", "label": str | None}}``. +``label`` is set only for corrections; the manual cell type of an accepted cell is +resolved from the automated label downstream. This schema is validated by +``csde.read_annotations`` — keep the two in step. """ import argparse @@ -14,6 +26,20 @@ import streamlit as st import streamlit.components.v1 as components +ACCEPT, CORRECT, REJECT = "accept", "correct", "reject" + +# Action -> (button glyph, button label). The glyph is what the keyboard handler +# matches on, so each must be unique. +ACTION_STYLE = { + ACCEPT: ("✓", "Accept"), + CORRECT: ("✎", "Correct"), + REJECT: ("✗", "Reject"), +} +MODE_ACTIONS = { + "accept_reject": [ACCEPT, REJECT], + "accept_correct_reject": [ACCEPT, CORRECT, REJECT], +} + def parse_args(): p = argparse.ArgumentParser() @@ -34,20 +60,43 @@ def save_annotations(annotations: dict, annotation_dir: Path) -> None: json.dump(annotations, f, indent=2) +def format_status(record: dict | None) -> str: + if record is None: + return "not annotated" + glyph, label = ACTION_STYLE[record["action"]] + if record["action"] == CORRECT: + return f"{glyph} corrected → **{record['label']}**" + return f"{glyph} {label.lower()}ed" + + def main(): args = parse_args() annotation_dir = Path(args.dir) cell_type_of_interest = "cell of interest" + annotation_mode = "accept_correct_reject" + vocabulary = [] config_path = annotation_dir / "config.json" if config_path.exists(): with open(config_path) as f: config = json.load(f) cell_type_of_interest = config.get("cell_type_of_interest", cell_type_of_interest) + annotation_mode = config.get("annotation_mode", annotation_mode) + vocabulary = config.get("cell_type_vocabulary", []) + + actions = MODE_ACTIONS[annotation_mode] st.set_page_config(layout="wide", page_title="Cell Annotator") st.title(f"Cell Annotation — {cell_type_of_interest}") + if CORRECT in actions and not vocabulary: + st.error( + "`cell_type_vocabulary` is missing from config.json, so corrections " + "cannot offer any cell type. Re-run scripts/export.py, or set " + '`"annotation_mode": "accept_reject"` in config.json.' + ) + st.stop() + metadata_path = annotation_dir / "metadata.csv" if not metadata_path.exists(): st.warning("metadata.csv not found — waiting for export.py to write the first cell.") @@ -103,13 +152,14 @@ def main(): row = metadata.iloc[idx] cell_id = row["cell_id"] existing = annotations.get(cell_id) - status = "✓ correct" if existing is True else ("✗ incorrect" if existing is False else "not annotated") - st.subheader(f"Cell `{cell_id}` — predicted: **{row['cell_type']}** — {status}") + st.subheader( + f"Cell `{cell_id}` — predicted: **{row['cell_type']}** — {format_status(existing)}" + ) st.image(str(row["image_path"]), use_container_width=True) - def annotate(is_correct: bool) -> None: - annotations[cell_id] = is_correct + def annotate(action: str, label: str | None = None) -> None: + annotations[cell_id] = {"action": action, "label": label} save_annotations(annotations, annotation_dir) # Advance to next unannotated after annotating remaining = metadata.index[~metadata["cell_id"].isin(annotations)] @@ -119,32 +169,74 @@ def annotate(is_correct: bool) -> None: elif remaining.any(): st.session_state.current_idx = int(remaining[0]) - col1, col2, _ = st.columns([1, 1, 4]) - with col1: - if st.button("✓ Correct [1]", type="primary", use_container_width=True): - annotate(True) - st.rerun() - with col2: - if st.button("✗ Incorrect [2]", use_container_width=True): - annotate(False) + select_key = f"correction_choice_{cell_id}" + pending = st.session_state.get("pending_correction") == cell_id + + if pending: + # The action buttons are removed while a correction is pending, so the + # digit shortcuts below match nothing and cannot fire. The selector is + # rendered under the pan + # el to keep the image visible while choosing. + st.markdown(f"**Correcting cell `{cell_id}`** — pick the right cell type:") + sel_col, cancel_col, _ = st.columns([2, 1, 3]) + with sel_col: + choice = st.selectbox( + "Corrected cell type", + [t for t in vocabulary if t != str(row["cell_type"])], + index=None, + placeholder="Type to filter…", + label_visibility="collapsed", + key=select_key, + ) + with cancel_col: + if st.button("Cancel", use_container_width=True): + st.session_state.pending_correction = None + st.rerun() + if choice is not None: + annotate(CORRECT, label=choice) + st.session_state.pending_correction = None st.rerun() + else: + cols = st.columns([1] * len(actions) + [6 - len(actions)]) + for col, action in zip(cols, actions): + glyph, label = ACTION_STYLE[action] + key_hint = actions.index(action) + 1 + with col: + clicked = st.button( + f"{glyph} {label} [{key_hint}]", + type="primary" if action == ACCEPT else "secondary", + use_container_width=True, + ) + if clicked: + if action == CORRECT: + # Nothing is written yet: the cell stays unannotated until a + # label is picked. + st.session_state.pending_correction = cell_id + st.session_state.pop(select_key, None) + else: + annotate(action) + st.rerun() - # Keyboard shortcuts: 1/2 annotate, ←/→ navigate - components.html(""" + # Keyboard shortcuts: digits annotate, ←/→ navigate + key_to_glyph = { + str(i + 1): ACTION_STYLE[action][0] for i, action in enumerate(actions) + } + components.html(f""" """, height=0) diff --git a/scripts/differential_expression.py b/scripts/differential_expression.py index d3092a7..f7daf8a 100644 --- a/scripts/differential_expression.py +++ b/scripts/differential_expression.py @@ -24,6 +24,13 @@ def parse_args(): help="Output CSV path (default: /results.csv).") p.add_argument("--spatial-group-key", default="spatial_group", help="obs column encoding the two spatial populations.") + p.add_argument("--spatial-group-target", default="1", + help="Value of --spatial-group-key identifying the target region " + "(group 1). Positive log-fold changes mean higher expression " + "here.") + p.add_argument("--spatial-group-reference", default="0", + help="Value of --spatial-group-key identifying the reference " + "region (group 0).") p.add_argument("--n-cells-expressed-threshold", type=int, default=10, help="Min cells expressing a gene for it to be tested.") p.add_argument("--noise-model", default="poisson", choices=["poisson", "nb"]) @@ -42,21 +49,28 @@ def main(): inputs = prepare_csde_inputs( annotation_dir=annotation_dir, spatial_group_key=args.spatial_group_key, + spatial_group_target=args.spatial_group_target, + spatial_group_reference=args.spatial_group_reference, layer=layer, n_cells_expressed_threshold=args.n_cells_expressed_threshold, ) adata_gt = inputs["adata_gt"] adata_other = inputs["adata_other"] + summary = inputs["summary"] + + print("Manual annotations:") + for key, value in summary.items(): + print(f" {key}: {value}") results = run_csde( adata_pred=adata_other, adata_gt=adata_gt, pred_cell_pop_key="prediction", + gt_cell_pop_key="annotation", cell_pop_a=0, cell_pop_b=1, - gt_key="is_correct", layer_name=layer, - importance_weights=adata_gt.obs["sampling_weight"].values, + importance_weights=1.0 / adata_gt.obs["sampling_weight"].values, noise_model=args.noise_model, ) diff --git a/scripts/export.py b/scripts/export.py index 7030c03..073cfdf 100644 --- a/scripts/export.py +++ b/scripts/export.py @@ -35,6 +35,11 @@ def parse_args(): help="JSON file mapping gene name → colour.") p.add_argument("--image-channel", default="DAPI") p.add_argument("--n-cells", type=int, default=600) + p.add_argument("--annotation-mode", + choices=["accept_correct_reject", "accept_reject"], + default="accept_correct_reject", + help="Actions offered by annotate.py. accept_correct_reject " + "additionally lets the annotator relabel a cell.") p.add_argument("--delta", type=float, default=50.0, help="Half-width of the spatial crop around each cell (microns).") p.add_argument("--n-top-genes", type=int, default=15) @@ -58,8 +63,10 @@ def main(): annotation_dir = Path(args.out) annotation_dir.mkdir(parents=True, exist_ok=True) + cell_types = sdata["table"].obs[args.cell_type_key].dropna().unique() + config = vars(args) | {"cell_type_vocabulary": sorted(map(str, cell_types))} with open(annotation_dir / "config.json", "w") as f: - json.dump(vars(args), f, indent=2) + json.dump(config, f, indent=2) metadata = export_cell_panels( sdata=sdata, diff --git a/src/csde/__init__.py b/src/csde/__init__.py index 13d6ad5..8adedfe 100644 --- a/src/csde/__init__.py +++ b/src/csde/__init__.py @@ -1,4 +1,10 @@ -from .annotation import export_cell_panels, load_annotations, prepare_csde_inputs +from .annotation import ( + ANNOTATION_ACTIONS, + export_cell_panels, + load_annotations, + prepare_csde_inputs, + read_annotations, +) from .api import run_csde from .model_nb import NBIntercept, NBInterceptModule from .model_poisson import PoissonIntercept, PoissonInterceptModule @@ -22,4 +28,6 @@ "export_cell_panels", "load_annotations", "prepare_csde_inputs", + "read_annotations", + "ANNOTATION_ACTIONS", ] diff --git a/src/csde/annotation.py b/src/csde/annotation.py index 252757f..f05ebf4 100644 --- a/src/csde/annotation.py +++ b/src/csde/annotation.py @@ -5,10 +5,80 @@ from pathlib import Path import matplotlib.pyplot as plt +import numpy as np import pandas as pd from .spatial_utils import plot_region, plot_top_genes, subsample_cells +# --- annotations.json schema ------------------------------------------------- +# {cell_id: {"action": "accept" | "correct" | "reject", "label": str | None}} +# +# ``label`` is set only for "correct"; the manual cell type of an accepted cell +# is resolved from the automated label at read time, so a cell type is never +# recorded in two places. scripts/annotate.py writes this file and mirrors these +# constants; the validation below is what keeps the two in step. +ACCEPT = "accept" +CORRECT = "correct" +REJECT = "reject" +ANNOTATION_ACTIONS = (ACCEPT, CORRECT, REJECT) + + +def _resolve_group_value(value, observed: set): + """ + Match a requested spatial-group value against the values in the obs column. + + Exact equality first, then string equality, so that a value coming from the + command line (always a string) matches an integer-encoded column. Returns the + value as it appears in the column, or the input unchanged when there is no + match — the caller reports that. + """ + if value in observed: + return value + by_str = {str(seen): seen for seen in observed} + return by_str.get(str(value), value) + + +def read_annotations(annotation_dir: str | Path) -> dict: + """ + Load ``annotations.json`` and validate it against the schema above. + + Returns + ------- + dict + ``{cell_id: {"action": ..., "label": ...}}``, with ``cell_id`` as ``str``. + """ + annotation_dir = Path(annotation_dir) + ann_path = annotation_dir / "annotations.json" + if not ann_path.exists(): + raise FileNotFoundError( + f"No annotations found at {ann_path}. Run scripts/annotate.py first." + ) + with open(ann_path) as f: + records = json.load(f) + + for cell_id, record in records.items(): + if not isinstance(record, dict) or "action" not in record: + raise ValueError( + f"Malformed annotation for cell '{cell_id}': {record!r}. Expected " + '{"action": "accept" | "correct" | "reject", "label": str | None}.' + ) + action, label = record["action"], record.get("label") + if action not in ANNOTATION_ACTIONS: + raise ValueError( + f"Unknown action '{action}' for cell '{cell_id}'. " + f"Expected one of {ANNOTATION_ACTIONS}." + ) + if action == CORRECT and label is None: + raise ValueError( + f"Cell '{cell_id}' is annotated as '{CORRECT}' but carries no label." + ) + if action != CORRECT and label is not None: + raise ValueError( + f"Cell '{cell_id}' is annotated as '{action}' but carries label " + f"'{label}'. Only '{CORRECT}' annotations may set a label." + ) + return {str(cell_id): record for cell_id, record in records.items()} + def prepare_csde_inputs( annotation_dir: str | Path, @@ -24,15 +94,23 @@ def prepare_csde_inputs( Reads config.json (cell_type_key, cell_type_of_interest, sdata path), metadata.csv (sampling_weight per annotated cell), and annotations.json - (is_correct per cell) from annotation_dir. + (accept / correct / reject per cell) from annotation_dir. + + This function owns the translation from annotation scheme to model labels: + :func:`~csde.run_csde` consumes ``.obs["annotation"]`` as-is. Label encoding in .obs["prediction"] / .obs["annotation"]: 1 — cell_type_of_interest with spatial_group == 1 (target) 0 — cell_type_of_interest with spatial_group == 0 (reference) 2 — all other cells - For GT annotation labels, cells predicted as cell_type_of_interest but - marked incorrect (is_correct=False) are reassigned to class 2. + ``prediction`` uses the automated cell type; ``annotation`` uses the manual + cell type, which is the automated one for an *accepted* cell, the annotator's + choice for a *corrected* one, and undefined (hence class 2) for a *rejected* + one. A corrected cell keeps its automated spatial group: only the cell-type + component of the label is curated. Correcting a cell *into* + cell_type_of_interest therefore moves it from class 2 into class 0 or 1, + which is the case an accept/reject scheme cannot express. Parameters ---------- @@ -59,14 +137,16 @@ def prepare_csde_inputs( adata_gt : AnnData Annotated cells. obs columns added: ``prediction`` (int 0/1/2), - ``annotation`` (int 0/1/2), ``is_correct`` (bool), - ``sampling_weight`` (float). Genes are filtered. + ``annotation`` (int 0/1/2), ``manual_action`` (str), + ``manual_cell_type`` (str or None), ``sampling_weight`` (float). + Genes are filtered. adata_other : AnnData All unannotated cells. obs column added: ``prediction`` (int 0/1/2). Same gene set as adata_gt. + summary : dict + Per-action counts plus ``n_promoted`` / ``n_demoted``, the number of + cells the manual curation moved into / out of the compared groups. """ - import numpy as np - annotation_dir = Path(annotation_dir) with open(annotation_dir / "config.json") as f: @@ -74,13 +154,7 @@ def prepare_csde_inputs( cell_type_key = config["cell_type_key"] cell_type_of_interest = config["cell_type_of_interest"] - ann_path = annotation_dir / "annotations.json" - if not ann_path.exists(): - raise FileNotFoundError( - f"No annotations found at {ann_path}. Run scripts/annotate.py first." - ) - with open(ann_path) as f: - annotations = json.load(f) # {cell_id: True/False} + annotations = read_annotations(annotation_dir) metadata = pd.read_csv(annotation_dir / "metadata.csv") metadata["cell_id"] = metadata["cell_id"].astype(str) @@ -105,6 +179,38 @@ def prepare_csde_inputs( adata.obs_names = adata.obs_names.astype(str) adata = adata[adata.obs[cell_type_key].notna()].copy() + if cell_type_of_interest not in set(adata.obs[cell_type_key].unique()): + raise ValueError( + f"cell_type_of_interest={cell_type_of_interest!r} was not found in obs " + f"column '{cell_type_key}', which contains " + f"{sorted(map(str, adata.obs[cell_type_key].unique()))}." + ) + + # Fail before any label math: a region value that does not occur would send + # every cell to class 2, and the resulting emptiness would only surface much + # later as a confusing "population not found" error. Not auto-detected from a + # two-level column on purpose — which value becomes the target sets the sign + # of every log-fold change, so it has to be chosen explicitly. + observed_groups = set(adata.obs[spatial_group_key].unique()) + spatial_group_target = _resolve_group_value(spatial_group_target, observed_groups) + spatial_group_reference = _resolve_group_value( + spatial_group_reference, observed_groups + ) + missing = [ + value + for value in (spatial_group_target, spatial_group_reference) + if value not in observed_groups + ] + if missing: + raise ValueError( + f"spatial_group_target={spatial_group_target!r} and " + f"spatial_group_reference={spatial_group_reference!r}: " + f"{missing!r} not found in obs column '{spatial_group_key}', which " + f"contains {sorted(map(str, observed_groups))}. Pass " + "--spatial-group-target / --spatial-group-reference to choose the two " + "regions to compare." + ) + # --- Prediction labels (automated, all cells) --- is_coi = (adata.obs[cell_type_key] == cell_type_of_interest).values spatial_group = adata.obs[spatial_group_key].values @@ -120,22 +226,96 @@ def prepare_csde_inputs( adata_gt = adata[annotated_mask].copy() - is_correct_arr = np.array( - [annotations[cid] for cid in adata_gt.obs_names], dtype=bool + # --- Manual cell type: accepted keeps the automated label, corrected takes + # the annotator's, rejected has none --- + actions = np.array( + [annotations[cid]["action"] for cid in adata_gt.obs_names], dtype=object + ) + manual_cell_type = [] + for cell_id, predicted_type in zip( + adata_gt.obs_names, adata_gt.obs[cell_type_key] + ): + record = annotations[cell_id] + if record["action"] == REJECT: + manual_cell_type.append(None) + elif record["action"] == CORRECT: + manual_cell_type.append(record["label"]) + else: + manual_cell_type.append(predicted_type) + + # The vocabulary is closed: a corrected label absent from the data would + # otherwise fall through to class 2 and be indistinguishable from a rejection. + vocabulary = set(adata.obs[cell_type_key].unique()) + unknown = sorted( + { + label + for label, action in zip(manual_cell_type, actions) + if action == CORRECT and label not in vocabulary + } + ) + if unknown: + raise ValueError( + f"Corrected labels {unknown} are not present in " + f"'{cell_type_key}' of the SpatialData table. Known cell types: " + f"{sorted(vocabulary)}." + ) + + adata_gt.obs["manual_action"] = actions + adata_gt.obs["manual_cell_type"] = pd.Series( + manual_cell_type, index=adata_gt.obs_names, dtype=object ) - adata_gt.obs["is_correct"] = is_correct_arr # --- Annotation (GT) labels --- - is_coi_gt = (adata_gt.obs[cell_type_key] == cell_type_of_interest).values + # The spatial group is always the automated one: manual curation revises the + # cell type, never the region. + is_coi_gt = np.array( + [label == cell_type_of_interest for label in manual_cell_type] + ) spatial_group_gt = adata_gt.obs[spatial_group_key].values annotation = np.full(len(adata_gt), 2, dtype=int) - annotation[is_coi_gt & is_correct_arr & (spatial_group_gt == spatial_group_target)] = 1 - annotation[is_coi_gt & is_correct_arr & (spatial_group_gt == spatial_group_reference)] = 0 + annotation[is_coi_gt & (spatial_group_gt == spatial_group_target)] = 1 + annotation[is_coi_gt & (spatial_group_gt == spatial_group_reference)] = 0 adata_gt.obs["annotation"] = annotation adata_gt.obs["sampling_weight"] = adata_gt.obs_names.map(sampling_weights).values + prediction_gt = adata_gt.obs["prediction"].values + summary = { + "n_annotated": int(len(adata_gt)), + **{ + f"n_{action}": int((actions == action).sum()) + for action in ANNOTATION_ACTIONS + }, + # Cells the curation moved into / out of the two compared groups. + "n_promoted": int(((prediction_gt == 2) & (annotation != 2)).sum()), + "n_demoted": int(((prediction_gt != 2) & (annotation == 2)).sum()), + } + + # Both label sets must populate both compared groups. Checked here, rather + # than left to fail at fit time, so the message can name the region and the + # counts: this is usually an annotation-budget problem, not a wiring one. + for labels, source in ( + (prediction_gt, "the automated pipeline"), + (annotation, "manual annotation"), + ): + for group, region in ( + (0, spatial_group_reference), + (1, spatial_group_target), + ): + if (labels == group).sum() > 0: + continue + n_in_groups = int((labels != 2).sum()) + raise ValueError( + f"No annotated cell falls in group {group} " + f"({cell_type_of_interest!r} in region {region!r}) according to " + f"{source}, so its expression cannot be estimated. Of " + f"{len(adata_gt)} annotated cells, {n_in_groups} are " + f"{cell_type_of_interest!r} across both regions. Annotate more " + "cells, or raise --target-proportion at export time to sample " + f"more {cell_type_of_interest!r}." + ) + # --- Gene filter: expressed in >= threshold pred-target/ref cells in adata_gt --- # pred_mask = adata_gt.obs["prediction"].isin([0, 1]) pred_mask = adata_gt.obs["annotation"].isin([0, 1]) @@ -150,30 +330,40 @@ def prepare_csde_inputs( adata_gt = adata_gt[:, gene_mask].copy() adata_other = adata[~annotated_mask][:, gene_mask].copy() - return {"adata_gt": adata_gt, "adata_other": adata_other} + return {"adata_gt": adata_gt, "adata_other": adata_other, "summary": summary} def load_annotations(annotation_dir: str | Path) -> pd.DataFrame: """ Merge ``metadata.csv`` and ``annotations.json`` into a single DataFrame. - Returns only annotated cells, with an added boolean ``is_correct`` column. - Pass the result as ``adata_gt`` to :func:`~csde.run_csde`. + Returns only annotated cells, with added ``action`` (accept / correct / + reject) and ``manual_cell_type`` columns. The latter is the automated + ``cell_type`` for accepted cells, the annotator's choice for corrected ones, + and None for rejected ones. + + This is a convenience view for inspecting annotations; the model labels are + built by :func:`prepare_csde_inputs`. """ annotation_dir = Path(annotation_dir) metadata = pd.read_csv(annotation_dir / "metadata.csv") metadata["cell_id"] = metadata["cell_id"].astype(str) - ann_path = annotation_dir / "annotations.json" - if not ann_path.exists(): - raise FileNotFoundError( - f"No annotations found at {ann_path}. Run scripts/annotate.py first." - ) - with open(ann_path) as f: - annotations = json.load(f) + annotations = read_annotations(annotation_dir) - metadata["is_correct"] = metadata["cell_id"].map(annotations) - return metadata[metadata["is_correct"].notna()].copy() + metadata["action"] = metadata["cell_id"].map( + {cell_id: record["action"] for cell_id, record in annotations.items()} + ) + metadata = metadata[metadata["action"].notna()].copy() + metadata["manual_cell_type"] = [ + None + if action == REJECT + else (annotations[cell_id]["label"] if action == CORRECT else predicted) + for cell_id, action, predicted in zip( + metadata["cell_id"], metadata["action"], metadata["cell_type"] + ) + ] + return metadata def export_cell_panels( diff --git a/src/csde/api.py b/src/csde/api.py index 056d806..676cba2 100644 --- a/src/csde/api.py +++ b/src/csde/api.py @@ -13,24 +13,26 @@ def _map_cell_types( cell_type_col: str, cell_pop_a: str, cell_pop_b: str, + context: str = "dataset", + hint: str = "", ) -> np.ndarray: """ Map cell types to a simplified 3-class representation. 0: cell_pop_a (Reference) 1: cell_pop_b (Target) 2: Other + + ``context`` and ``hint`` only shape the error raised when a population is + empty; the mapping itself is identical for the automated and manual labels. """ labels = np.full(len(obs), 2, dtype=int) - # Check if cell types exist - if cell_pop_a not in obs[cell_type_col].values: - raise ValueError( - f"Cell population '{cell_pop_a}' not found in column '{cell_type_col}'" - ) - if cell_pop_b not in obs[cell_type_col].values: - raise ValueError( - f"Cell population '{cell_pop_b}' not found in column '{cell_type_col}'" - ) + for cell_pop in (cell_pop_a, cell_pop_b): + if cell_pop not in obs[cell_type_col].values: + raise ValueError( + f"No cells assigned to population '{cell_pop}' in column " + f"'{cell_type_col}' of the {context}.{hint}" + ) labels[obs[cell_type_col] == cell_pop_a] = 0 labels[obs[cell_type_col] == cell_pop_b] = 1 @@ -42,9 +44,9 @@ def run_csde( adata_pred: anndata.AnnData, adata_gt: anndata.AnnData, pred_cell_pop_key: str, + gt_cell_pop_key: str, cell_pop_a: str, cell_pop_b: str, - gt_key: str, layer_name: Optional[str] = None, importance_weights: Optional[np.ndarray] = None, noise_model: str = "poisson", @@ -60,9 +62,14 @@ def run_csde( adata_pred: AnnData object containing cells with prediction-based assignments only. adata_gt: AnnData object containing cells with ground-truth assignments. pred_cell_pop_key: Column in .obs containing the prediction-based cell population labels. + Read from both ``adata_pred`` and ``adata_gt``. + gt_cell_pop_key: Column in adata_gt.obs containing the manually curated cell + population labels. Built upstream by + :func:`~csde.prepare_csde_inputs`, which resolves the annotation scheme + (accept/reject or accept/correct/reject) into labels; this function only + consumes them. cell_pop_a: Name of the first cell population (reference group). cell_pop_b: Name of the second cell population (target group). - gt_key: Boolean column in adata_gt.obs indicating if the prediction is correct. layer_name: Layer in adata.layers to use for expression counts. If None, uses .X. importance_weights: Optional 1-D array of importance weights for the ground-truth observations. Will be normalized to sum to n_obs internally. @@ -77,24 +84,35 @@ def run_csde( - p_value_adj: Benjamini-Hochberg multiplicity-adjusted p-value. """ - # create simplified 3-class representation for predictions (pop_a, pop_b, other) + # create simplified 3-class representation (pop_a, pop_b, other). + # The automated and manual labels are mapped identically; the annotation + # scheme that produced the manual labels is resolved upstream. y_pred_unl = _map_cell_types( - adata_pred.obs, pred_cell_pop_key, cell_pop_a, cell_pop_b + adata_pred.obs, + pred_cell_pop_key, + cell_pop_a, + cell_pop_b, + context="unlabeled set", ) y_pred_gt_set = _map_cell_types( - adata_gt.obs, pred_cell_pop_key, cell_pop_a, cell_pop_b + adata_gt.obs, + pred_cell_pop_key, + cell_pop_a, + cell_pop_b, + context="automated labels of the manually annotated set", + ) + y_gt = _map_cell_types( + adata_gt.obs, + gt_cell_pop_key, + cell_pop_a, + cell_pop_b, + context="manually annotated set", + hint=( + " No annotated cell was curated into this group, so its expression " + "cannot be estimated. Annotate more cells, or increase the " + "importance-sampling weight of the cell type of interest." + ), ) - - # logic to construct gt labels based on boolean column gt_key - # - if predicted as a and correct (gt_key=true) -> gt is a (0) - # - if predicted as b and correct (gt_key=true) -> gt is b (1) - # - else -> gt is other (2) - y_gt = np.full(len(adata_gt), 2, dtype=int) - is_correct = adata_gt.obs[gt_key].values.astype(bool) - is_pred_a = (adata_gt.obs[pred_cell_pop_key] == cell_pop_a).values - is_pred_b = (adata_gt.obs[pred_cell_pop_key] == cell_pop_b).values - y_gt[is_pred_a & is_correct] = 0 - y_gt[is_pred_b & is_correct] = 1 def get_X(adata): if layer_name: diff --git a/tests/test_csde.py b/tests/test_csde.py index 4a63e41..636c7f0 100644 --- a/tests/test_csde.py +++ b/tests/test_csde.py @@ -1,10 +1,13 @@ +import json +import tempfile import unittest +from pathlib import Path import anndata import numpy as np import pandas as pd -from csde import run_csde +from csde import prepare_csde_inputs, run_csde class TestCSDE(unittest.TestCase): @@ -21,10 +24,11 @@ def setUp(self): n_gt = 50 X_gt = np.random.poisson(lam=2.0, size=(n_gt, n_genes)).astype(float) + is_correct = np.random.choice([True, False], size=n_gt) obs_gt = pd.DataFrame( { "cell_type": np.random.choice(["TypeA", "TypeB", "TypeC"], size=n_gt), - "is_correct": np.random.choice([True, False], size=n_gt), + "is_correct": is_correct, } ) self.adata_gt = anndata.AnnData(X=X_gt, obs=obs_gt) @@ -34,15 +38,24 @@ def setUp(self): self.adata_pred.obs.iloc[1, 0] = "TypeB" self.adata_gt.obs.iloc[0, 0] = "TypeA" self.adata_gt.obs.iloc[1, 0] = "TypeB" + self.adata_gt.obs.iloc[0, 1] = True + self.adata_gt.obs.iloc[1, 1] = True + + # Manual labels: run_csde consumes a label column, not a boolean. + self.adata_gt.obs["manual_cell_type"] = np.where( + self.adata_gt.obs["is_correct"].values, + self.adata_gt.obs["cell_type"].values, + "Rejected", + ) def test_run_csde(self): res = run_csde( adata_pred=self.adata_pred, adata_gt=self.adata_gt, pred_cell_pop_key="cell_type", + gt_cell_pop_key="manual_cell_type", cell_pop_a="TypeA", cell_pop_b="TypeB", - gt_key="is_correct", optimizer="gd", optimizer_kwargs={"n_iter": 10}, # Fast run ) @@ -60,45 +73,47 @@ def test_run_csde_with_importance_weights(self): rng = np.random.default_rng(0) importance_weights = rng.uniform(0.5, 2.0, size=n_gt) - res = run_csde( - adata_pred=self.adata_pred, - adata_gt=self.adata_gt, - pred_cell_pop_key="cell_type", - cell_pop_a="TypeA", - cell_pop_b="TypeB", - gt_key="is_correct", - optimizer="gd", - optimizer_kwargs={"n_iter": 10}, - importance_weights=importance_weights, - noise_model="poisson", - ) - - self.assertIsInstance(res, pd.DataFrame) - self.assertEqual(len(res), 10) - self.assertListEqual( - list(res.columns), ["log_fold_change", "p_value", "p_value_adj"] - ) - self.assertTrue(not res.isnull().values.any()) + for noise_model in ("poisson", "nb"): + res = run_csde( + adata_pred=self.adata_pred, + adata_gt=self.adata_gt, + pred_cell_pop_key="cell_type", + gt_cell_pop_key="manual_cell_type", + cell_pop_a="TypeA", + cell_pop_b="TypeB", + optimizer="gd", + optimizer_kwargs={"n_iter": 10}, + importance_weights=importance_weights, + noise_model=noise_model, + ) - res = run_csde( - adata_pred=self.adata_pred, - adata_gt=self.adata_gt, - pred_cell_pop_key="cell_type", - cell_pop_a="TypeA", - cell_pop_b="TypeB", - gt_key="is_correct", - optimizer="gd", - optimizer_kwargs={"n_iter": 10}, - importance_weights=importance_weights, - noise_model="nb", - ) + self.assertIsInstance(res, pd.DataFrame) + self.assertEqual(len(res), 10) + self.assertListEqual( + list(res.columns), ["log_fold_change", "p_value", "p_value_adj"] + ) + self.assertTrue(not res.isnull().values.any()) - self.assertIsInstance(res, pd.DataFrame) - self.assertEqual(len(res), 10) - self.assertListEqual( - list(res.columns), ["log_fold_change", "p_value", "p_value_adj"] + def test_missing_gt_population_raises(self): + # No annotated cell curated into TypeB -> the group cannot be estimated. + adata_gt = self.adata_gt.copy() + adata_gt.obs["manual_cell_type"] = np.where( + adata_gt.obs["manual_cell_type"] == "TypeB", + "Rejected", + adata_gt.obs["manual_cell_type"], ) - self.assertTrue(not res.isnull().values.any()) + with self.assertRaises(ValueError) as ctx: + run_csde( + adata_pred=self.adata_pred, + adata_gt=adata_gt, + pred_cell_pop_key="cell_type", + gt_cell_pop_key="manual_cell_type", + cell_pop_a="TypeA", + cell_pop_b="TypeB", + optimizer="gd", + optimizer_kwargs={"n_iter": 10}, + ) + self.assertIn("manually annotated set", str(ctx.exception)) def test_importance_weights_wrong_shape(self): from csde.model_poisson import PoissonIntercept as InterceptRegression @@ -121,5 +136,152 @@ def test_importance_weights_wrong_shape(self): ) +COI = "macrophage" + +# cell_id -> (automated cell_type, spatial_group) +CELLS = { + "c0": (COI, 0), # accept -> 0 + "c1": (COI, 1), # accept -> 1 + "c2": (COI, 1), # reject -> 2 + "c3": (COI, 1), # correct away from COI -> 2 + "c4": ("fibroblast", 1), # correct into COI -> 1 (promotion) + "c5": ("fibroblast", 0), # correct into COI -> 0 (promotion) + "c6": ("fibroblast", 0), # accept -> 2 + "c7": (COI, 1), # unannotated + "c8": ("fibroblast", 0), # unannotated +} +ANNOTATIONS = { + "c0": {"action": "accept", "label": None}, + "c1": {"action": "accept", "label": None}, + "c2": {"action": "reject", "label": None}, + "c3": {"action": "correct", "label": "fibroblast"}, + "c4": {"action": "correct", "label": COI}, + "c5": {"action": "correct", "label": COI}, + "c6": {"action": "accept", "label": None}, +} + + +def _build_sdata(): + cell_ids = list(CELLS) + obs = pd.DataFrame( + { + "cell_type": [CELLS[c][0] for c in cell_ids], + "spatial_group": [CELLS[c][1] for c in cell_ids], + }, + index=cell_ids, + ) + adata = anndata.AnnData(X=np.ones((len(cell_ids), 4)), obs=obs) + adata.var_names = [f"Gene_{i}" for i in range(4)] + # prepare_csde_inputs only ever does sdata["table"]. + return {"table": adata} + + +def _write_annotation_dir(tmpdir: Path, annotations: dict) -> Path: + with open(tmpdir / "config.json", "w") as f: + json.dump( + {"cell_type_key": "cell_type", "cell_type_of_interest": COI}, + f, + ) + pd.DataFrame( + { + "cell_id": list(annotations), + "cell_type": [CELLS[c][0] for c in annotations], + "sampling_weight": [1.0] * len(annotations), + } + ).to_csv(tmpdir / "metadata.csv", index=False) + with open(tmpdir / "annotations.json", "w") as f: + json.dump(annotations, f) + return tmpdir + + +class TestPrepareCsdeInputs(unittest.TestCase): + """Label construction: the path where a bug yields wrong DE rather than a crash.""" + + def _run(self, annotations=None): + with tempfile.TemporaryDirectory() as tmp: + annotation_dir = _write_annotation_dir( + Path(tmp), ANNOTATIONS if annotations is None else annotations + ) + return prepare_csde_inputs( + annotation_dir=annotation_dir, + sdata=_build_sdata(), + n_cells_expressed_threshold=1, + ) + + def test_label_construction(self): + obs = self._run()["adata_gt"].obs + expected = {"c0": 0, "c1": 1, "c2": 2, "c3": 2, "c4": 1, "c5": 0, "c6": 2} + self.assertEqual(obs["annotation"].to_dict(), expected) + + def test_promotion_changes_the_label(self): + # The case an accept/reject scheme cannot express: automated says 2, + # manual curation moves the cell into a compared group. + obs = self._run()["adata_gt"].obs + for cell_id, expected_group in (("c4", 1), ("c5", 0)): + self.assertEqual(obs.loc[cell_id, "prediction"], 2) + self.assertEqual(obs.loc[cell_id, "annotation"], expected_group) + self.assertEqual(obs.loc[cell_id, "manual_cell_type"], COI) + + def test_manual_cell_type_resolution(self): + obs = self._run()["adata_gt"].obs + self.assertEqual(obs.loc["c0", "manual_cell_type"], COI) # accepted + self.assertEqual(obs.loc["c3", "manual_cell_type"], "fibroblast") # corrected + self.assertTrue(pd.isna(obs.loc["c2", "manual_cell_type"])) # rejected + + def test_summary_counts(self): + summary = self._run()["summary"] + self.assertEqual(summary["n_annotated"], 7) + self.assertEqual(summary["n_accept"], 3) + self.assertEqual(summary["n_correct"], 3) + self.assertEqual(summary["n_reject"], 1) + self.assertEqual(summary["n_promoted"], 2) # c4, c5 + self.assertEqual(summary["n_demoted"], 2) # c2, c3 + + def test_unannotated_cells_go_to_adata_other(self): + inputs = self._run() + self.assertEqual(set(inputs["adata_other"].obs_names), {"c7", "c8"}) + self.assertEqual(inputs["adata_other"].obs["prediction"].to_dict(), {"c7": 1, "c8": 2}) + + def test_unknown_corrected_label_raises(self): + annotations = dict(ANNOTATIONS) + annotations["c4"] = {"action": "correct", "label": "not_a_cell_type"} + with self.assertRaises(ValueError) as ctx: + self._run(annotations) + self.assertIn("not_a_cell_type", str(ctx.exception)) + + def test_malformed_annotations_raise(self): + cases = [ + {"c0": True}, # the old boolean format is no longer accepted + {"c0": {"action": "maybe", "label": None}}, + {"c0": {"action": "correct", "label": None}}, + {"c0": {"action": "accept", "label": "fibroblast"}}, + ] + for annotations in cases: + with self.subTest(annotations=annotations): + with self.assertRaises(ValueError): + self._run(annotations) + + def test_accept_reject_only_matches_previous_behaviour(self): + # Strict extension: with no corrections, labels equal the old + # (is_coi & is_correct & spatial_group) formula. + annotations = { + cell_id: { + "action": "accept" if cell_id in ("c0", "c1", "c6") else "reject", + "label": None, + } + for cell_id in ANNOTATIONS + } + obs = self._run(annotations)["adata_gt"].obs + expected = {} + for cell_id in annotations: + cell_type, spatial_group = CELLS[cell_id] + is_correct = annotations[cell_id]["action"] == "accept" + expected[cell_id] = ( + spatial_group if (cell_type == COI and is_correct) else 2 + ) + self.assertEqual(obs["annotation"].to_dict(), expected) + self.assertEqual(self._run(annotations)["summary"]["n_promoted"], 0) + + if __name__ == "__main__": unittest.main()