Upload rare_disease_kg.py with huggingface_hub

rare_disease_kg.py

@@ -0,0 +1,160 @@
+"""Shared heterogeneous rare-disease knowledge graph for Pillars A & C.
+
+Loads the full PrimeKG-derived graph if it is present locally (the production
+path), otherwise falls back to a compact, biologically-real fixture covering the
+GEMEO/SUS cohort so the pillars run out-of-the-box with `python pillar_a_*.py`.
+
+The fixture encodes *real* rare-disease biology (causal genes, key phenotypes,
+gene–gene interactions, mechanistically-related diseases). It is small on
+purpose — enough to demonstrate the algorithms and reproduce Marfan→FBN1 — and
+is transparently replaced by PrimeKG when `GEMEO_KG_DIR` (or the default
+~/rarasnet-swarm-py/gemeo/data) contains hetero_graph.json + node_ids.json.
+
+Node id convention: "ORPHA:<n>" diseases, "HGNC:<sym>" genes, "HP:<n>" phenotypes.
+"""
+from __future__ import annotations
+import json, os
+from dataclasses import dataclass, field
+
+DEFAULT_KG_DIR = os.environ.get(
+    "GEMEO_KG_DIR", os.path.expanduser("~/rarasnet-swarm-py/gemeo/data"))
+
+
+# ── Compact biologically-real fixture (ORPHA / HGNC / HP) ────────────────────
+# disease -> {name, genes:[(gene, "causal"|"modifier")], phenotypes:[hp...]}
+_DISEASES = {
+    "ORPHA:558":  ("Marfan syndrome",            [("FBN1", "causal"), ("TGFBR2", "modifier")],
+                   ["HP:0004942", "HP:0001166", "HP:0000545", "HP:0001763"]),  # aortic dilat, arachnodactyly, myopia, pes planus
+    "ORPHA:98896":("Duchenne muscular dystrophy",[("DMD", "causal")],
+                   ["HP:0003560", "HP:0003707", "HP:0001644", "HP:0002650"]),
+    "ORPHA:70":   ("Spinal muscular atrophy 1",  [("SMN1", "causal"), ("SMN2", "modifier")],
+                   ["HP:0003457", "HP:0001324", "HP:0002093"]),
+    "ORPHA:586":  ("Cystic fibrosis",            [("CFTR", "causal")],
+                   ["HP:0006538", "HP:0002202", "HP:0001733", "HP:0002595"]),
+    "ORPHA:355":  ("Gaucher disease",            [("GBA", "causal")],
+                   ["HP:0001433", "HP:0001744", "HP:0001873", "HP:0002240"]),
+    "ORPHA:579":  ("Mucopolysaccharidosis I",    [("IDUA", "causal")],
+                   ["HP:0001407", "HP:0000256", "HP:0001182"]),
+    "ORPHA:580":  ("Mucopolysaccharidosis II",   [("IDS", "causal")],
+                   ["HP:0001407", "HP:0000256", "HP:0001371"]),
+    "ORPHA:905":  ("Wilson disease",             [("ATP7B", "causal")],
+                   ["HP:0001399", "HP:0002171", "HP:0001337"]),
+    "ORPHA:95":   ("Friedreich ataxia",          [("FXN", "causal")],
+                   ["HP:0002066", "HP:0001251", "HP:0001638"]),
+    "ORPHA:636":  ("Neurofibromatosis type 1",   [("NF1", "causal")],
+                   ["HP:0009732", "HP:0000957", "HP:0009735"]),
+    "ORPHA:778":  ("Rett syndrome",              [("MECP2", "causal")],
+                   ["HP:0001250", "HP:0000252", "HP:0001260"]),
+    "ORPHA:100":  ("Ataxia-telangiectasia",      [("ATM", "causal")],
+                   ["HP:0001251", "HP:0001009", "HP:0002205", "HP:0002584"]),
+}
+# gene–gene interactions (PPI / pathway) — enables 2-hop guilt-by-association
+_GENE_GENE = [("FBN1", "TGFBR2"), ("FBN1", "TGFBR1"), ("SMN1", "SMN2"),
+              ("ATM", "TP53"), ("GBA", "SNCA"), ("DMD", "UTRN")]
+# mechanistically-related diseases (storage disorders, neuromuscular, …)
+_DIS_DIS = [("ORPHA:579", "ORPHA:580"), ("ORPHA:355", "ORPHA:579"),
+            ("ORPHA:70", "ORPHA:98896"), ("ORPHA:95", "ORPHA:100")]
+_HP_NAME = {  # a few labels for readable evidence paths
+    "HP:0004942": "ascending aortic aneurysm", "HP:0001166": "arachnodactyly",
+    "HP:0003560": "muscular dystrophy", "HP:0003457": "EMG: neuropathic changes",
+    "HP:0006538": "recurrent pneumonia", "HP:0001433": "hepatosplenomegaly",
+    "HP:0002066": "gait ataxia", "HP:0001251": "ataxia", "HP:0001250": "seizure",
+}
+
+
+@dataclass
+class KG:
+    """Heterogeneous KG with typed adjacency. Nodes are typed string ids."""
+    nodes: list                      # ordered node ids
+    idx: dict                        # node id -> int index
+    ntype: dict                      # node id -> "disease"|"gene"|"phenotype"
+    adj: dict = field(default_factory=dict)   # node id -> set(neighbor ids)
+    edge_label: dict = field(default_factory=dict)  # (u,v) -> relation label
+    names: dict = field(default_factory=dict)        # node id -> human label
+
+    def neighbors(self, u):
+        return self.adj.get(u, set())
+
+    def of_type(self, t):
+        return [n for n in self.nodes if self.ntype[n] == t]
+
+
+def _add_edge(adj, lab, u, v, relation):
+    adj.setdefault(u, set()).add(v)
+    adj.setdefault(v, set()).add(u)
+    lab[(u, v)] = relation
+    lab[(v, u)] = relation
+
+
+def _build_fixture() -> KG:
+    nodes, ntype, names = [], {}, {}
+    adj, lab = {}, {}
+    def ensure(nid, t, nm=None):
+        if nid not in ntype:
+            ntype[nid] = t; nodes.append(nid)
+            if nm: names[nid] = nm
+    for did, (dname, genes, phens) in _DISEASES.items():
+        ensure(did, "disease", dname)
+        for g, kind in genes:
+            ensure(g, "gene", g)
+            _add_edge(adj, lab, did, g, f"associated_with:{kind}")
+        for hp in phens:
+            ensure(hp, "phenotype", _HP_NAME.get(hp, hp))
+            _add_edge(adj, lab, did, hp, "has_phenotype")
+    for g1, g2 in _GENE_GENE:
+        ensure(g1, "gene", g1); ensure(g2, "gene", g2)
+        _add_edge(adj, lab, g1, g2, "interacts_with")
+    for d1, d2 in _DIS_DIS:
+        _add_edge(adj, lab, d1, d2, "related_to")
+    idx = {n: i for i, n in enumerate(nodes)}
+    return KG(nodes=nodes, idx=idx, ntype=ntype, adj=adj, edge_label=lab, names=names)
+
+
+def _build_from_primekg(kg_dir: str) -> KG:
+    """Load the production PrimeKG-derived graph (hetero_graph.json + node_ids.json)."""
+    g = json.load(open(f"{kg_dir}/hetero_graph.json"))
+    node_ids = json.load(open(f"{kg_dir}/node_ids.json"))
+    nodes, ntype, names = [], {}, {}
+    adj, lab = {}, {}
+    row2id = {t: {int(k): (f"{ '' if t=='disease' else ''}{v}") for k, v in m.items()}
+              for t, m in node_ids.items()}
+    # canonical typed ids
+    def cid(t, raw):
+        pfx = {"disease": "ORPHA:", "gene": "HGNC:", "phenotype": "HP:"}.get(t, t + ":")
+        s = str(raw)
+        return s if s.startswith(pfx) else pfx + s
+    for t, m in row2id.items():
+        for _, raw in m.items():
+            nid = cid(t, raw)
+            if nid not in ntype:
+                ntype[nid] = t; nodes.append(nid)
+    REL = {"disease__associated_with__gene": ("disease", "gene", "associated_with"),
+           "disease__has_phenotype__phenotype": ("disease", "phenotype", "has_phenotype")}
+    for et, e in g.get("edges", {}).items():
+        if et not in REL:
+            continue
+        st, dt, rel = REL[et]
+        for s, d in zip(e["src"], e["dst"]):
+            u = cid(st, row2id[st].get(int(s), s)); v = cid(dt, row2id[dt].get(int(d), d))
+            if u in ntype and v in ntype:
+                _add_edge(adj, lab, u, v, rel)
+    idx = {n: i for i, n in enumerate(nodes)}
+    return KG(nodes=nodes, idx=idx, ntype=ntype, adj=adj, edge_label=lab, names=names)
+
+
+def load_kg(kg_dir: str | None = None) -> tuple[KG, str]:
+    """Return (KG, source). Prefer PrimeKG on disk; else the built-in fixture."""
+    kg_dir = kg_dir or DEFAULT_KG_DIR
+    try:
+        if os.path.exists(f"{kg_dir}/hetero_graph.json"):
+            return _build_from_primekg(kg_dir), "primekg"
+    except Exception as e:  # pragma: no cover
+        print(f"[kg] PrimeKG load failed ({e}); using fixture")
+    return _build_fixture(), "fixture"
+
+
+if __name__ == "__main__":
+    kg, src = load_kg()
+    print(f"KG source: {src} | {len(kg.nodes)} nodes "
+          f"(diseases={len(kg.of_type('disease'))}, genes={len(kg.of_type('gene'))}, "
+          f"phenotypes={len(kg.of_type('phenotype'))})")