src/ml/gene_programs.py

"""Gene programs for the "cell-dividing neuron" extension.

The base virtual-cell tutorial predicts a perturbation's effect on the full
highly-variable gene panel. This module layers *biological meaning* on top of
that panel by grouping genes into named **programs** and providing a way to
score how strongly a given expression profile expresses each program.

Three programs are defined, chosen around the "cell-dividing neuron" idea. The
base dataset is the Virtual Cell Challenge H1 human embryonic stem cell (hESC)
line: a fast-*dividing*, pluripotent cell that is the canonical starting point
for *neuronal* differentiation. So the interesting biology lives on two axes:

  - CELL_DIVISION  : cell-cycle / mitosis genes. High score == actively dividing.
  - NEURON_IDENTITY: pan-neuronal markers. High score == neuron-like.
  - PLURIPOTENCY   : stem-cell identity markers (the hESC ground state).

A "cell-dividing neuron" is the (biologically unusual) corner where the
CELL_DIVISION score and the NEURON_IDENTITY score are *both* high. The model
extension lets us ask which gene knockdowns push a cell toward that corner.

Gene symbols are human HGNC symbols, matching `adata.var_names` in the VCC data.
The cell-cycle lists are the widely used Tirosh et al. (2016, Science) S-phase
and G2/M marker sets that Scanpy ships for `score_genes_cell_cycle`, with a few
legacy symbols updated to current HGNC names.
"""

from __future__ import annotations

import numpy as np


# --- S phase (Tirosh et al. 2016) -------------------------------------------
S_PHASE_GENES = [
    "MCM5", "PCNA", "TYMS", "FEN1", "MCM2", "MCM4", "RRM1", "UNG", "GINS2",
    "MCM6", "CDCA7", "DTL", "PRIM1", "UHRF1", "CENPU", "HELLS", "RFC2",
    "RPA2", "NASP", "RAD51AP1", "GMNN", "WDR76", "SLBP", "CCNE2", "UBR7",
    "POLD3", "MSH2", "ATAD2", "RAD51", "RRM2", "CDC45", "CDC6", "EXO1",
    "TIPIN", "DSCC1", "BLM", "CASP8AP2", "USP1", "CLSPN", "POLA1", "CHAF1B",
    "BRIP1", "E2F8",
]

# --- G2 / M phase (Tirosh et al. 2016) --------------------------------------
G2M_PHASE_GENES = [
    "HMGB2", "CDK1", "NUSAP1", "UBE2C", "BIRC5", "TPX2", "TOP2A", "NDC80",
    "CKS2", "NUF2", "CKS1B", "MKI67", "TMPO", "CENPF", "TACC3", "PIMREG",
    "SMC4", "CCNB2", "CKAP2L", "CKAP2", "AURKB", "BUB1", "KIF11", "ANP32E",
    "TUBB4B", "GTSE1", "KIF20B", "HJURP", "CDCA3", "JPT1", "CDC20", "TTK",
    "CDC25C", "KIF2C", "RANGAP1", "NCAPD2", "DLGAP5", "CDCA2", "CDCA8",
    "ECT2", "KIF23", "HMMR", "AURKA", "PSRC1", "ANLN", "LBR", "CKAP5",
    "CENPE", "CTCF", "NEK2", "G2E3", "GAS2L3", "CBX5", "CENPA",
]

# The headline "is the cell dividing?" program: S + G2/M together, plus the
# canonical proliferation marker MKI67 / Ki-67 (already in G2M) and cyclins.
CELL_DIVISION_GENES = sorted(set(
    S_PHASE_GENES + G2M_PHASE_GENES + ["CCNB1", "CCNA2", "CCNE1", "CDK4", "CDK6"]
))

# --- Pan-neuronal identity markers ------------------------------------------
NEURON_IDENTITY_GENES = [
    "MAP2", "RBFOX3", "TUBB3", "NEFL", "NEFM", "NEFH", "SYN1", "SYP", "DCX",
    "STMN2", "ELAVL3", "ELAVL4", "GAP43", "SNAP25", "NRXN1", "NRXN3", "DLG4",
    "GRIA1", "GRIN1", "SYT1", "NCAM1", "ENO2", "MAPT", "GRIA2", "CAMK2A",
    "RELN", "DCLK1", "TUBB2B", "INA", "MEF2C",
]

# Pro-neural transcription factors: knocking these *down* should reduce the
# neuronal program; they are the master regulators of the neuron axis.
PRONEURAL_TFS = [
    "NEUROD1", "NEUROD2", "NEUROG1", "NEUROG2", "ASCL1", "POU3F2", "MYT1L",
    "INSM1", "EOMES", "TBR1",
]

# --- Pluripotency / stem-cell identity (the hESC ground state) --------------
PLURIPOTENCY_GENES = [
    "POU5F1", "NANOG", "SOX2", "LIN28A", "LIN28B", "ZFP42", "DPPA4",
    "DNMT3B", "SALL4", "TDGF1", "UTF1", "PRDM14", "KLF4", "FOXD3", "TERT",
]


# --- Endocrine / hormone signalling (all present in the VCC measured panel) ----
# hESCs are not endocrine tissue, so this program is deliberately "odd" for this
# cell type -- part of the research point that generality still lets the downstream
# model speak regardless of how unusual the modulation mixture is.
ENDOCRINE_GENES = [
    # peptide hormones
    "INS", "GCG", "SST", "GH1", "PRL", "POMC", "OXT", "AVP", "LEP", "NPY",
    "GIP", "TRH", "CRH", "GNRH1", "LHB", "FSHB",
    # hormone receptors
    "INSR", "IGF1R", "IGF2R", "GHR", "PRLR", "LEPR", "GLP1R", "GCGR", "TSHR",
    "MC4R", "CRHR1",
    # nuclear receptors
    "ESR1", "AR", "PGR", "THRA", "THRB", "NR3C1", "NR3C2", "PPARG", "PPARA",
    "RXRA", "RARA", "VDR", "NR5A1", "NR0B1",
    # steroidogenesis
    "STAR", "CYP11A1", "CYP17A1", "CYP19A1", "HSD11B1", "HSD11B2",
    # endocrine transcription factors + secretory granule
    "PDX1", "NEUROD1", "NKX2-1", "NKX6-1", "ISL1", "FOXA2", "HNF1A", "HNF1B",
    "HNF4A", "GATA4", "GATA6", "CHGA", "CHGB", "SCG2", "SCG3", "SCG5",
    "IGFBP3", "IGFBP5",
]

# --- Core metabolism (glycolysis, TCA, lipid/sterol, one-carbon) + OxPhos ------
METABOLISM_GENES = [
    # glycolysis / TCA
    "HK1", "HK2", "GPI", "PFKM", "PGK1", "ENO1", "PKM", "LDHB", "PDHA1", "CS",
    "IDH1", "IDH2", "FH", "SDHA", "MDH2", "G6PD",
    # lipid / sterol synthesis
    "ACLY", "ACACA", "FASN", "SCD", "HMGCR", "HMGCS1", "MVD", "FDPS", "SQLE",
    "DHCR24", "DHCR7", "INSIG1", "ACAT2", "SREBF1", "SREBF2",
    # amino-acid / one-carbon
    "MAT2A", "GLUD1", "GLS", "SHMT2", "PHGDH",
    # oxidative phosphorylation
    "NDUFB4", "NDUFB6", "NDUFA4", "SDHB", "UQCRB", "UQCRQ", "UQCRC1", "CYC1",
    "COX4I1", "COX6C", "COX7C", "ATP5F1A", "ATP5F1B", "ATP5MC1",
]

# --- Signalling / cell identity (RTK -> RAS/MAPK/PI3K, Wnt, TGFb, JAK/STAT) ----
SIGNALING_IDENTITY_GENES = [
    # receptor tyrosine kinases
    "IGF1R", "INSR", "EGFR", "ERBB2", "FGFR1",
    # RAS / MAPK
    "KRAS", "NRAS", "HRAS", "RAF1", "BRAF", "MAP2K1", "MAPK1", "MAPK3",
    # PI3K / AKT / mTOR
    "AKT1", "AKT2", "PIK3CA", "PTEN", "MTOR",
    # Wnt / beta-catenin
    "GSK3B", "CTNNB1", "TCF7L2", "MYC", "MAX",
    # immediate-early / JAK-STAT
    "JUN", "FOS", "STAT1", "STAT3", "STAT6", "SRC", "CSK",
    # Notch / TGFb
    "NOTCH1", "SMAD2", "SMAD3", "SMAD4", "TGFBR2",
    # adhesion / epithelial identity
    "THY1", "KRT18", "VCL", "ITGAV", "CXCL12", "CDH1", "EPCAM",
]


PROGRAMS: dict[str, list[str]] = {
    "cell_division": CELL_DIVISION_GENES,
    "neuron_identity": NEURON_IDENTITY_GENES,
    "pluripotency": PLURIPOTENCY_GENES,
    "endocrine": ENDOCRINE_GENES,
    "metabolism": METABOLISM_GENES,
    "signaling_identity": SIGNALING_IDENTITY_GENES,
}

# Regulators worth perturbing for each program (used by the synthetic data
# generator and as suggested perturbations in the explorer).
PROGRAM_REGULATORS: dict[str, list[str]] = {
    "cell_division": ["CDK1", "CCNB1", "AURKA", "AURKB", "PLK1", "MKI67", "TOP2A", "BUB1"],
    "neuron_identity": PRONEURAL_TFS,
    "pluripotency": ["POU5F1", "NANOG", "SOX2", "LIN28A", "PRDM14"],
    "endocrine": ["PDX1", "NEUROD1", "ISL1", "FOXA2", "NR5A1", "INSR", "IGF1R"],
    "metabolism": ["SREBF1", "SREBF2", "HMGCR", "ACLY", "MYC", "INSIG1"],
    "signaling_identity": ["MYC", "KRAS", "AKT1", "MTOR", "CTNNB1", "MAPK1", "EGFR"],
}


def program_indices(gene_names, program_genes) -> np.ndarray:
    """Indices into `gene_names` of the program genes that are present."""
    lookup = {g: i for i, g in enumerate(gene_names)}
    return np.array([lookup[g] for g in program_genes if g in lookup], dtype=np.int64)


def present_program_genes(gene_names, program_genes) -> list[str]:
    present = set(gene_names)
    return [g for g in program_genes if g in present]


def program_score(expr, gene_names, program_genes, ref_indices=None):
    """Score how strongly each profile expresses a program.

    Parameters
    ----------
    expr : array of shape (G,) or (N, G)
        Log-normalised expression (the same space the model is trained in).
    gene_names : sequence of str, length G
    program_genes : list of str
    ref_indices : optional array of gene indices used as a background. When
        given, the background mean is subtracted (Scanpy `score_genes` style),
        which removes profile-wide depth effects and isolates the program.

    Returns
    -------
    float (for a 1-D input) or np.ndarray of shape (N,)
    """
    expr = np.asarray(expr, dtype=np.float64)
    one_d = expr.ndim == 1
    if one_d:
        expr = expr[None, :]

    idx = program_indices(gene_names, program_genes)
    if idx.size == 0:
        raise ValueError("None of the program genes are present in gene_names.")

    score = expr[:, idx].mean(axis=1)
    if ref_indices is not None and len(ref_indices) > 0:
        score = score - expr[:, ref_indices].mean(axis=1)

    return float(score[0]) if one_d else score


def program_weight_vector(gene_names, weight=5.0, programs=PROGRAMS):
    """A per-gene weight vector that up-weights program genes for the loss.

    Baseline weight is 1.0; any gene belonging to one of `programs` gets
    `weight`. This is how the model is made to *focus* on the division /
    neuron biology instead of treating all 2000 HVGs equally.
    """
    w = np.ones(len(gene_names), dtype=np.float32)
    for genes in programs.values():
        idx = program_indices(gene_names, genes)
        w[idx] = weight
    return w


def summarize_coverage(gene_names) -> dict[str, tuple[int, int]]:
    """How many of each program's genes are present in `gene_names`."""
    return {
        name: (len(present_program_genes(gene_names, genes)), len(genes))
        for name, genes in PROGRAMS.items()
    }