run_guide_assignment.py

"""GUIDE ASSIGNMENT: from raw gRNA counts -> a per-cell perturbation label, then validate it.

Each cell's guide identity isn't given directly — it's called from the raw guide-count matrix (gdo).
We assign each cell to its dominant guide (above a count threshold), map the guide to its target gene,
and check concordance against the paper's own labels. This is step 0 of every Perturb-seq analysis.

  python run_guide_assignment.py --threshold 5
"""
from __future__ import annotations
import argparse
import re
import numpy as np
import pandas as pd
import scipy.sparse as sp
from src import data


def guide_to_target(g: str) -> str:
    if g == "Negative":
        return "Negative"
    base = re.sub(r"g\d+$", "", g)              # 'STAT1g2' -> 'STAT1', 'NTg3' -> 'NT'
    return "NT" if base in ("NT", "eGFP") else base


def run(threshold: float = 5.0):
    mdata = data.load()
    gdo = data.get_gdo(mdata)
    print(f"guide matrix: {gdo.shape[0]} cells x {gdo.shape[1]} guides; raw counts max={float(gdo.X.max()):.0f}")

    # Assign each cell to its dominant guide if that guide's count >= threshold, else 'Negative'.
    # (This is what pertpy's pt.pp.GuideAssignment.assign_to_max_guide does; we do it explicitly —
    # it's clearer, and pertpy 1.0.3's AnnData dispatch has a bug.)
    X = gdo.X.toarray() if sp.issparse(gdo.X) else np.asarray(gdo.X)
    top_idx = X.argmax(axis=1)
    top_val = X.max(axis=1)
    guides = np.asarray(gdo.var_names)
    gdo.obs["assigned_guide"] = np.where(top_val >= threshold, guides[top_idx], "Negative")
    gdo.obs["assigned_target"] = [guide_to_target(g) for g in gdo.obs["assigned_guide"]]

    pred = gdo.obs["assigned_target"].to_numpy()
    truth = mdata.obs["gene_target"].to_numpy()
    assigned = pred != "Negative"
    conc = float((pred[assigned] == truth[assigned]).mean())

    print(f"\nthreshold = {threshold}")
    print(f"  cells assigned a guide : {assigned.mean():.1%}  ({assigned.sum()}/{len(pred)})")
    print(f"  unassigned ('Negative'): {(~assigned).mean():.1%}")
    print(f"  target concordance vs paper labels (assigned cells): {conc:.1%}")
    top = pd.Series(pred[assigned]).value_counts().head(8)
    print("\n  most common assigned targets:")
    for t, n in top.items():
        print(f"    {t:10s} {n}")
    print("\nTakeaway: guide assignment is a real upstream step — threshold + 'one dominant guide'\n"
          "decisions (and the unassigned/multiplet rate) directly shape every downstream result.")
    return conc


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--threshold", type=float, default=5.0, help="min count for the top guide to assign a cell")
    run(ap.parse_args().threshold)