""" Narada: Patient case generator. Builds episode cases from the knowledge graph + disease catalog. Each case is a dict matching PatientCase structure. Task types: monogenic — single causal gene, 3-4 phenotypes, 5-8 candidates oligogenic — 2 causal genes (one variant each), 5-7 phenotypes, 10-15 candidates phenotype_mismatch — cardiac patient + high-pathogenicity cancer decoy """ from __future__ import annotations import random import uuid from typing import Any, Dict, List, Optional, Set, Tuple from .graph import ( DISEASE_CATALOG, GENE_TO_DISEASES, PATHWAY_MAP, NaradaGraph, _clinsig_to_score, _slugify, ) from .models import GraphNode, Variant # ── BRCA1/BRCA2 decoy pool ───────────────────────────────────────────────────── # Frameshift/nonsense variants are maximally salient for LLMs — best decoys. _DECOY_GENES = ["BRCA1", "BRCA2", "TP53", "MLH1", "MSH2"] _DECOY_TYPES = {"frameshift", "deletion", "nonsense", "stop_gained", "indel"} def _is_high_impact(v: Dict[str, Any]) -> bool: vtype = v.get("variant_type", "").lower() name = v.get("name", "").lower() return ( any(t in vtype for t in _DECOY_TYPES) or "frameshift" in name or "stop" in name or "del" in vtype ) def _pick_variants( graph: NaradaGraph, genes: List[str], n: int, prefer_high_impact: bool = False, rng: Optional[random.Random] = None, ) -> List[Dict[str, Any]]: """Pick up to n variants from the given gene list.""" if rng is None: rng = random.Random() pool = graph.get_variants_for_genes(genes) if prefer_high_impact: high = [v for v in pool if _is_high_impact(v)] pool = high if high else pool if len(pool) <= n: return pool return rng.sample(pool, n) def _variant_to_model(v: Dict[str, Any], graph: NaradaGraph) -> Variant: var_id = graph.variant_node_id(v["allele_id"]) return Variant( id=var_id, allele_id=v["allele_id"], gene=v["gene"], name=v["name"][:150] if v["name"] else f"{v['gene']} variant", variant_type=v["variant_type"], clinical_significance=v["clnsig"], pathogenicity_score=_clinsig_to_score(v["clnsig"]), disease_associations=v["diseases"][:3], ) def _dict_to_graph_node(graph: NaradaGraph, node_id: str) -> GraphNode: nd = graph.get_node(node_id) if nd is None: return GraphNode( id=node_id, type="unknown", name=node_id, description="", connected_node_ids=[], ) neighbors = graph.get_neighbors(node_id) return GraphNode( id=nd["id"], type=nd["type"], name=nd["name"], description=nd["description"], connected_node_ids=neighbors[:30], # cap for observation size metadata=nd["metadata"], ) # ── Case structure ───────────────────────────────────────────────────────────── class PatientCase: """ A single patient episode definition. Immutable after construction — shared state lives in the environment. """ def __init__( self, case_id: str, task_type: str, disease_name: str, causal_genes: List[str], causal_allele_ids: List[str], # ground truth patient_hpo_ids: List[str], patient_phenotype_names: List[str], candidate_variants: List[Variant], starting_node_id: str, relevant_node_ids: Set[str], decoy_gene: Optional[str] = None, absent_hpo_ids: Optional[List[str]] = None, absent_phenotype_names: Optional[List[str]] = None, ) -> None: self.case_id = case_id self.task_type = task_type self.disease_name = disease_name self.causal_genes = causal_genes self.causal_allele_ids = causal_allele_ids self.patient_hpo_ids = patient_hpo_ids self.patient_phenotype_names = patient_phenotype_names self.absent_hpo_ids = absent_hpo_ids or [] self.absent_phenotype_names = absent_phenotype_names or [] self.candidate_variants = candidate_variants self.starting_node_id = starting_node_id self.relevant_node_ids = relevant_node_ids self.decoy_gene = decoy_gene @property def ground_truth_variant_ids(self) -> List[str]: return [f"VAR:{aid}" for aid in self.causal_allele_ids] # ── Generators ───────────────────────────────────────────────────────────────── def _pick_hpo_subset( hpo_ids: List[str], graph: NaradaGraph, n: int, rng: random.Random, ) -> Tuple[List[str], List[str]]: """Return (hpo_ids, names) for n terms, using only ones present in graph.""" present = [h for h in hpo_ids if h in graph.nodes] if not present: present = hpo_ids[:n] chosen = rng.sample(present, min(n, len(present))) names = [graph.get_hpo_name(h) for h in chosen] return chosen, names def _find_starting_node( graph: NaradaGraph, hpo_ids: List[str], rng: random.Random, ) -> str: """Find a good starting phenotype node in the graph.""" for h in rng.sample(hpo_ids, len(hpo_ids)): if h in graph.nodes: return h # Fallback: any phenotype node pheno_nodes = [nid for nid, nd in graph.nodes.items() if nd["type"] == "phenotype"] return rng.choice(pheno_nodes) if pheno_nodes else list(graph.nodes.keys())[0] def generate_monogenic_case( graph: NaradaGraph, rng: Optional[random.Random] = None, ) -> PatientCase: """Single causal gene, 3-4 phenotypes, 5-8 candidate variants.""" if rng is None: rng = random.Random() eligible = [d for d in DISEASE_CATALOG if "monogenic" in d["task_types"] and d["genes"]] disease = rng.choice(eligible) # Pick one primary gene that has variants for gene in rng.sample(disease["genes"], len(disease["genes"])): if graph.get_variants_for_gene(gene): causal_gene = gene break else: raise RuntimeError(f"No variants found for any gene in {disease['disease']}") # Ground truth: one causal variant for the single-gene tier. causal_raw = _pick_variants(graph, [causal_gene], n=1, prefer_high_impact=True, rng=rng) if not causal_raw: raise RuntimeError(f"No variants for {causal_gene}") causal_allele_ids = [v["allele_id"] for v in causal_raw] # Patient phenotypes: 3-4 terms n_pheno = rng.randint(3, 4) hpo_ids, hpo_names = _pick_hpo_subset(disease["hpo_ids"], graph, n_pheno, rng) # Absent phenotypes: disease HPO terms the patient does NOT have (diagnostic exclusions) chosen_set = set(hpo_ids) absent_candidates = [h for h in disease["hpo_ids"] if h not in chosen_set and h in graph.nodes] rng.shuffle(absent_candidates) absent_hpo_ids = absent_candidates[:3] absent_names = [graph.get_hpo_name(h) for h in absent_hpo_ids] # Candidate variants: causal + 3-6 distractors from same-pathway genes target_pathway = disease["pathway"] distractor_genes = [ g for g in graph.gene_variants.keys() if g != causal_gene and ( PATHWAY_MAP.get(g) == target_pathway or any(target_pathway == d["pathway"] for d in GENE_TO_DISEASES.get(g, [])) ) ] if len(distractor_genes) < 3: distractor_genes = [g for g in graph.gene_variants.keys() if g != causal_gene] n_distractors = rng.randint(3, 6) distractor_raw = _pick_variants( graph, rng.sample(distractor_genes, min(6, len(distractor_genes))), n=n_distractors, rng=rng, ) all_raw = causal_raw + distractor_raw rng.shuffle(all_raw) candidates = [_variant_to_model(v, graph) for v in all_raw] starting_node = _find_starting_node(graph, hpo_ids, rng) relevant = graph.relevant_nodes_for_case( causal_genes=[causal_gene], patient_hpo_ids=hpo_ids, causal_allele_ids=causal_allele_ids, ) return PatientCase( case_id=str(uuid.uuid4())[:8], task_type="monogenic", disease_name=disease["disease"], causal_genes=[causal_gene], causal_allele_ids=causal_allele_ids, patient_hpo_ids=hpo_ids, patient_phenotype_names=hpo_names, candidate_variants=candidates, starting_node_id=starting_node, relevant_node_ids=relevant, absent_hpo_ids=absent_hpo_ids, absent_phenotype_names=absent_names, ) def generate_oligogenic_case( graph: NaradaGraph, rng: Optional[random.Random] = None, ) -> PatientCase: """2 causal genes (one variant each), 5-7 phenotypes, 10-15 candidates.""" if rng is None: rng = random.Random() # Only accept diseases whose catalog lists >=2 genes AND where at least # two of those genes actually have ClinVar variants in the loaded data. # Without this filter the oligogenic tier can silently degrade to a # single-gene case and violate the "flag both variants" contract. eligible = [] for d in DISEASE_CATALOG: if "oligogenic" not in d["task_types"] or len(d["genes"]) < 2: continue with_variants = [g for g in d["genes"] if graph.get_variants_for_gene(g)] if len(with_variants) >= 2: eligible.append((d, with_variants)) if not eligible: raise RuntimeError("No oligogenic diseases have >=2 genes with variants") disease, genes_with_variants = rng.choice(eligible) causal_genes = rng.sample(genes_with_variants, 2) # One causal variant per contributing gene. causal_raw = [] causal_allele_ids = [] for gene in causal_genes: vs = _pick_variants(graph, [gene], n=1, prefer_high_impact=True, rng=rng) causal_raw.extend(vs) causal_allele_ids.extend(v["allele_id"] for v in vs) # Patient phenotypes: 5-7 terms n_pheno = rng.randint(5, min(7, len(disease["hpo_ids"]))) hpo_ids, hpo_names = _pick_hpo_subset(disease["hpo_ids"], graph, n_pheno, rng) # Absent phenotypes chosen_set = set(hpo_ids) absent_candidates = [h for h in disease["hpo_ids"] if h not in chosen_set and h in graph.nodes] rng.shuffle(absent_candidates) absent_hpo_ids = absent_candidates[:3] absent_names = [graph.get_hpo_name(h) for h in absent_hpo_ids] # Distractors: from same-pathway genes target_pathway = disease["pathway"] distractor_genes = [ g for g in graph.gene_variants.keys() if g not in causal_genes and ( PATHWAY_MAP.get(g) == target_pathway or any(target_pathway == d["pathway"] for d in GENE_TO_DISEASES.get(g, [])) ) ] if len(distractor_genes) < 4: distractor_genes = [g for g in graph.gene_variants.keys() if g not in causal_genes] n_distractors = rng.randint(6, 9) distractor_raw = _pick_variants( graph, rng.sample(distractor_genes, min(6, len(distractor_genes))), n=n_distractors, rng=rng, ) all_raw = causal_raw + distractor_raw rng.shuffle(all_raw) candidates = [_variant_to_model(v, graph) for v in all_raw[:15]] starting_node = _find_starting_node(graph, hpo_ids, rng) relevant = graph.relevant_nodes_for_case( causal_genes=causal_genes, patient_hpo_ids=hpo_ids, causal_allele_ids=causal_allele_ids, ) return PatientCase( case_id=str(uuid.uuid4())[:8], task_type="oligogenic", disease_name=disease["disease"], causal_genes=causal_genes, causal_allele_ids=causal_allele_ids, patient_hpo_ids=hpo_ids, patient_phenotype_names=hpo_names, candidate_variants=candidates, starting_node_id=starting_node, relevant_node_ids=relevant, absent_hpo_ids=absent_hpo_ids, absent_phenotype_names=absent_names, ) def generate_mismatch_case( graph: NaradaGraph, rng: Optional[random.Random] = None, ) -> PatientCase: """ Phenotype mismatch: cardiac/neurological patient with high-pathogenicity cancer decoy in the candidate pool. Tests causal discipline. """ if rng is None: rng = random.Random() eligible = [d for d in DISEASE_CATALOG if "phenotype_mismatch" in d["task_types"]] disease = rng.choice(eligible) # Causal gene from actual disease for gene in rng.sample(disease["genes"], len(disease["genes"])): if graph.get_variants_for_gene(gene): causal_gene = gene break else: raise RuntimeError(f"No variants for {disease['disease']}") # Causal variant causal_raw = _pick_variants(graph, [causal_gene], n=1, prefer_high_impact=True, rng=rng) causal_allele_ids = [v["allele_id"] for v in causal_raw] # Patient phenotypes: 4-6 terms n_pheno = rng.randint(4, min(6, len(disease["hpo_ids"]))) hpo_ids, hpo_names = _pick_hpo_subset(disease["hpo_ids"], graph, n_pheno, rng) # Absent phenotypes chosen_set = set(hpo_ids) absent_candidates = [h for h in disease["hpo_ids"] if h not in chosen_set and h in graph.nodes] rng.shuffle(absent_candidates) absent_hpo_ids = absent_candidates[:3] absent_names = [graph.get_hpo_name(h) for h in absent_hpo_ids] # DECOY: pick a high-pathogenicity BRCA1/BRCA2 frameshift decoy_gene = rng.choice([g for g in _DECOY_GENES if graph.get_variants_for_gene(g)]) decoy_raw = _pick_variants(graph, [decoy_gene], n=2, prefer_high_impact=True, rng=rng) # Boost decoy salience without mutating the graph cache. decoy_raw = [dict(v, clnsig="Pathogenic") for v in decoy_raw] # Same-pathway distractors target_pathway = disease["pathway"] distractor_genes = [ g for g in graph.gene_variants.keys() if g != causal_gene and g not in _DECOY_GENES and ( PATHWAY_MAP.get(g) == target_pathway or any(target_pathway == d["pathway"] for d in GENE_TO_DISEASES.get(g, [])) ) ] if len(distractor_genes) < 2: distractor_genes = [ g for g in graph.gene_variants.keys() if g != causal_gene and g not in _DECOY_GENES ] n_distractors = rng.randint(3, 5) distractor_raw = _pick_variants( graph, rng.sample(distractor_genes, min(4, len(distractor_genes))), n=n_distractors, rng=rng, ) all_raw = causal_raw + decoy_raw + distractor_raw rng.shuffle(all_raw) candidates = [_variant_to_model(v, graph) for v in all_raw[:15]] starting_node = _find_starting_node(graph, hpo_ids, rng) relevant = graph.relevant_nodes_for_case( causal_genes=[causal_gene], patient_hpo_ids=hpo_ids, causal_allele_ids=causal_allele_ids, ) return PatientCase( case_id=str(uuid.uuid4())[:8], task_type="phenotype_mismatch", disease_name=disease["disease"], causal_genes=[causal_gene], causal_allele_ids=causal_allele_ids, patient_hpo_ids=hpo_ids, patient_phenotype_names=hpo_names, candidate_variants=candidates, starting_node_id=starting_node, relevant_node_ids=relevant, decoy_gene=decoy_gene, absent_hpo_ids=absent_hpo_ids, absent_phenotype_names=absent_names, ) _GENERATORS = { "monogenic": generate_monogenic_case, "oligogenic": generate_oligogenic_case, "phenotype_mismatch": generate_mismatch_case, } MAX_STEPS = { "monogenic": 15, "oligogenic": 25, "phenotype_mismatch": 20, } def generate_case( graph: NaradaGraph, task_type: str, seed: Optional[int] = None, ) -> PatientCase: if task_type not in _GENERATORS: raise ValueError(f"Unknown task_type: {task_type!r}. Choose from {list(_GENERATORS)}") rng = random.Random(seed) return _GENERATORS[task_type](graph, rng=rng)