"""Natural-language request to ranked CIPHER phenotypes. Dense, two-surface retrieval over prebuilt Vidul BGE_FT_VA embeddings: - phenotype_metadata surface -> phenotype-level semantic similarity - code_evidence surface -> code-level supporting evidence (max over chunks) A phenotype is a document, not a bag of codes. Each candidate carries both channel scores, a summary copied from source metadata, and the resolved code set. Ranking uses the fine-tuned dense retriever. Startup requires the matching embedding index and model. """ from __future__ import annotations import json import os import re from pathlib import Path import numpy as np from . import paths from .retriever import DenseEmbedder from .spelling import SpellSuggester # The config beside the vectors records the query encoding convention. EMB = paths.PARTA_EMB_DIR # Optional detail sidecars. CODE_DESC = paths.CODE_DESCRIPTIONS ALGO_COMPONENTS = paths.ALGORITHM_COMPONENTS FACETS = paths.PHENOTYPE_FACETS SUMMARIES = paths.PHENOTYPE_SUMMARIES BUNDLES = paths.RELATED_BUNDLES # Algorithm-component fields to surface, in display order (professor 1.4: # "refer to CIPHER algorithm components"). Only fields present for a phenotype # are shown; absent fields are omitted. _ALGO_FIELDS = [ ("methods_used", "Methods"), ("related_diseases", "Related diseases"), ("data_period", "Data period"), ("validations", "Validation"), ("adjudication_method", "Adjudication"), ] def _embedding_convention(config: dict) -> dict: """Return the build-time query contract; never invent missing defaults.""" required = ("model", "pooling", "query_prefix", "max_seq_length") missing = [key for key in required if key not in config] if missing: raise RuntimeError( "Phenotype embedding config is missing its encoding convention: " + ", ".join(missing)) return {key: config[key] for key in required} def _method_label(value: str) -> str: return "Rule-based" if value.lower().startswith("rules-based") else value # CIPHER public phenotype library. Detail page: # /web/cipher/phenotype-viewer?uqid={uqid}&name={slug} # where slug = the full name with every non-alphanumeric char replaced by '_'. # uqid comes from the cipher_links sidecar (present for ~79% of phenotypes); # the rest fall back to the library root. Domain overridable for other deployments. CIPHER_BASE = os.environ.get("ENCODE_CIPHER_BASE", "https://phenomics.va.ornl.gov").rstrip("/") CIPHER_LINKS = paths.CIPHER_LINKS # Friendly stand-ins so a still-unresolved code never shows the raw # 'unsupported_code_system' token (professor feedback 1.1). (label, source) _UNRESOLVED_NOTE = { 466: ("Medication code", "needs_vocab:RxNorm/NDC"), 467: ("LOINC lab code", "needs_vocab:LOINC"), 471: ("OMOP concept ID", "needs_vocab:OMOP"), 465: ("VA clinic stop code", "needs_vocab:VA_STOP"), 468: ("Study-defined variable", "study_specific"), 519: ("Study-defined variable", "study_specific"), } class EncodeEngine: def __init__(self, embedder: DenseEmbedder | None = None, emb_dir: Path | str | None = None): emb = Path(emb_dir) if emb_dir else EMB self.emb_dir = emb if not (emb / "config.json").exists(): raise RuntimeError(f"Missing runtime phenotype embeddings: {emb}") self.config = json.loads((emb / "config.json").read_text()) convention = _embedding_convention(self.config) self.phenotypes = {} for line in paths.CANONICAL_PHENOTYPES.open(): r = json.loads(line) self.phenotypes[r["phenotype_id"]] = r # "Did you mean" vocabulary from the phenotype text itself. self._suggester = SpellSuggester( text for p in self.phenotypes.values() for text in (p.get("title"), p.get("description"), *(p.get("keywords") or []))) # Metadata surface: one row per phenotype, aligned to meta_ids. self.meta_ids: list[int] = json.loads((emb / "metadata_ids.json").read_text()) self.meta_mat = np.load(emb / "metadata.npy").astype(np.float32) self._row = {pid: i for i, pid in enumerate(self.meta_ids)} # Code surface: many rows per phenotype; map each to its phenotype row. code_meta = json.loads((emb / "code_meta.json").read_text()) self.code_mat = np.load(emb / "code.npy").astype(np.float32) self._code_row = np.array([self._row.get(m["phenotype_id"], -1) for m in code_meta]) # Per-phenotype filters aligned to meta_ids. self._cat = np.array([self.phenotypes[p].get("category") or "(none)" for p in self.meta_ids]) self._validated = np.array([bool(self.phenotypes[p].get("validated")) for p in self.meta_ids]) # Part B: resolved code descriptions + algorithm components (optional sidecars). self.code_desc = {} if CODE_DESC.exists(): for line in CODE_DESC.open(): r = json.loads(line) self.code_desc[f"{r['code_type_id']}|{r['sub_type_id']}|{(r['code'] or '').upper()}"] = r self.cipher_uqid = {} if CIPHER_LINKS.exists(): for line in CIPHER_LINKS.open(): r = json.loads(line) self.cipher_uqid[r["phenotype_id"]] = r["uqid"] # Extracted CIPHER algorithm components (methods / related diseases / # validation / adjudication / data period); one record per phenotype. self.algo_components = {} if ALGO_COMPONENTS.exists(): for line in ALGO_COMPONENTS.open(): r = json.loads(line) self.algo_components[r["phenotype_id"]] = r # Offline facet extraction (evidence-or-abstain); enum values reach the # payload only when the source stated them, never inferred. self.facets = {} if FACETS.exists(): for line in FACETS.open(): r = json.loads(line) self.facets[r["phenotype_id"]] = r.get("facets") or {} # Optional LLM sidecars: grounded plain-language summaries and # related-concept bundles (search seeds, never codes). self.llm_summaries = {} if SUMMARIES.exists(): for line in SUMMARIES.open(): r = json.loads(line) self.llm_summaries[r["phenotype_id"]] = r self.related_bundles = {} if BUNDLES.exists(): for line in BUNDLES.open(): r = json.loads(line) self.related_bundles[r["phenotype_id"]] = r # Query embedder must match how the documents were encoded (model + # pooling + prefix + max-seq all travel in config.json). # # ENCODE_PARTA_MODEL repoints the query encoder for the default vector # set (the Space image uses it to load local weights instead of the # hub). It applies only to the default vectors because queries and # documents must use the same model. override = os.environ.get("ENCODE_PARTA_MODEL") if emb.resolve() == EMB.resolve() else None model_name = override or convention["model"] self.model_name = model_name self.embedder = embedder or DenseEmbedder( model_name, pooling=convention["pooling"], query_prefix=convention["query_prefix"], max_seq_length=convention["max_seq_length"]) def categories(self) -> list[str]: return sorted(set(self._cat.tolist())) # -- retrieval --------------------------------------------------------- def search(self, query: str, k: int = 10, categories: set[str] | None = None, validated_only: bool = False) -> dict: q = self.embedder.encode([query])[0] meta_score = self.meta_mat @ q # (n_pheno,) # Max code-evidence similarity aggregated to phenotype level. code_sims = self.code_mat @ q # (n_code,) code_score = np.zeros(len(self.meta_ids), dtype=np.float32) valid = self._code_row >= 0 np.maximum.at(code_score, self._code_row[valid], code_sims[valid]) retrieval = 0.5 * meta_score + 0.5 * code_score # bi-encoder candidate score mask = np.ones(len(self.meta_ids), dtype=bool) if categories: mask &= np.isin(self._cat, list(categories)) if validated_only: mask &= self._validated candidates = np.flatnonzero(mask) results = [] if len(candidates): candidate_scores = retrieval[candidates] for j in np.argsort(candidate_scores)[::-1][:k]: i = candidates[j] results.append(self._result(self.meta_ids[i], float(meta_score[i]), float(code_score[i]), float(candidate_scores[j]))) return { "query": query, "model": self.model_name, "count": len(results), "suggestion": self._suggester.suggest(query.strip()), "results": results, } def _result(self, pid, m, c, relevance) -> dict: p = self.phenotypes[pid] evidence = [ {"code_system": g.get("code_type_label"), "code_count": g.get("code_count") or len(g.get("codes", [])), "samples": [x.get("code") for x in g.get("codes", [])[:6]]} for g in p.get("associated_code_groups", []) ] return { "phenotype_id": pid, "title": p.get("title"), "category": p.get("category"), "validated": p.get("validated"), "summary": self._summary(p), "keywords": p.get("keywords") or [], "scores": {"relevance": round(relevance, 4), "metadata": round(m, 4), "code_evidence": round(c, 4)}, "code_systems": [e["code_system"] for e in evidence], "code_evidence": evidence, "facets": self._facet_payload(pid), "warnings": [], } @staticmethod def _summary(p: dict) -> str: """Return a shortened source description.""" desc = (p.get("description") or p.get("algorithm_description") or "").strip() if len(desc) > 360: desc = desc[:360].rsplit(" ", 1)[0] + "…" return desc def _algo_component_rows(self, pid: int) -> list[dict]: """Format available algorithm components as ordered label/value rows.""" rec = self.algo_components.get(pid) if not rec: return [] rows = [] for key, label in _ALGO_FIELDS: val = rec.get(key) if not val: continue if isinstance(val, list): values = [str(item) for item in val] if key == "methods_used": values = [_method_label(item) for item in values] text = ", ".join(values) elif isinstance(val, dict): text = (f"{val.get('start', '?')} to {val.get('end', '?')}" if ("start" in val or "end" in val) else ", ".join(f"{k}: {v}" for k, v in val.items())) else: text = str(val) rows.append({"label": label, "value": text}) return rows def _cipher_url(self, pid: int, title: str | None) -> str: uqid = self.cipher_uqid.get(pid) if uqid: slug = re.sub(r"[^A-Za-z0-9]", "_", title or "") return f"{CIPHER_BASE}/web/cipher/phenotype-viewer?uqid={uqid}&name={slug}" return CIPHER_BASE + "/" def _code_display(self, code_type_id, sub_type_id, entry: dict) -> tuple[str | None, str, dict]: """Resolve a code to (description, source), replacing 'unsupported_code_system' when a description exists in the sidecar.""" status = entry.get("label_status") labels = entry.get("labels") or [] if status in ("exact", "prefix_expanded") and labels: return labels[0], status, {} # Both gaps are looked up: a code whose system we had no dictionary for, # and a code whose system we support but whose label our tables lacked. if status in ("unsupported_code_system", "label_missing"): hit = self.code_desc.get( f"{code_type_id}|{sub_type_id}|{(entry.get('code') or '').upper()}") if hit: # Return the recorded match method and release with the label. extra = {k: hit[k] for k in ("source_version", "match", "concept") if hit.get(k)} return hit["description"], hit["source"], extra if status == "unsupported_code_system": note = _UNRESOLVED_NOTE.get(code_type_id, ("No standard description", "unresolved")) return note[0], note[1], {} return (labels[0] if labels else None), status, {} # -- detail ------------------------------------------------------------ def _code_graphable(self, ct: int | None, code: str | None) -> bool: """Does this specific code have a knowledge-graph? ICD diagnosis (460/461) resolves to an ICD family; medication terms (466) use a local map or an exact/normalized RxNAV lookup.""" if ct in (460, 461): return True if ct == 466: return bool((code or "").strip()) return False def _code_group_detail(self, g: dict) -> dict: ct, sub = g.get("code_type_id"), g.get("sub_type_id") codes, resolved = [], 0 for x in g.get("codes", [])[:200]: desc, source, extra = self._code_display(ct, sub, x) if desc and source not in ("study_specific", "unresolved", "label_missing") \ and not source.startswith("needs_vocab"): resolved += 1 codes.append({"code": x.get("code"), "description": desc, "description_source": source, "label_status": x.get("label_status"), "graphable": self._code_graphable(ct, x.get("code")), **extra}) return {"code_system": g.get("code_type_label"), "sub_type": g.get("sub_type_label"), "code_count": g.get("code_count") or len(g.get("codes", [])), "resolved_count": resolved, "codes": codes, "graphable": any(c["graphable"] for c in codes)} def _facet_payload(self, pid: int) -> dict | None: """Compact facet view for the UI: stated values only, criteria as text.""" f = self.facets.get(pid) if not f: return None def stated(name): v = ((f.get(name) or {}).get("value") or "").strip() return v if v and v != "not_stated" else None out = { "age_group": stated("age_group"), "care_setting": stated("care_setting"), "incident_vs_prevalent": stated("incident_vs_prevalent"), "intended_use": stated("intended_use"), "inclusion": [c.get("criterion") for c in f.get("inclusion") or [] if c.get("criterion")], "exclusion": [c.get("criterion") for c in f.get("exclusion") or [] if c.get("criterion")], } return out if any(out.values()) else None def phenotype(self, pid: int) -> dict | None: p = self.phenotypes.get(pid) if not p: return None return { "phenotype_id": pid, "algorithm_id": p.get("algorithm_id"), "title": p.get("title"), "category": p.get("category"), "validated": p.get("validated"), "validation_description": p.get("validation_description"), "description": p.get("description"), "algorithm_description": p.get("algorithm_description"), "keywords": p.get("keywords") or [], "authors": p.get("authors") or [], "publications": p.get("publications") or [], "population_description": p.get("population_description"), "last_modified": p.get("last_modified"), "cipher_url": self._cipher_url(pid, p.get("title")), "algorithm_components": self._algo_component_rows(pid), "facets": self._facet_payload(pid), "llm_summary": self.llm_summaries.get(pid), "related_bundle": self.related_bundles.get(pid), "code_groups": [self._code_group_detail(g) for g in p.get("associated_code_groups", [])], } # -- phenotype code hierarchy (phecode -> ICD main -> ICD sub) ---------- def phenotype_code_graph(self, pid: int, focus: str | None = None, cap: int | None = None) -> dict | None: """A node+edge tree of the phenotype's diagnosis codes: the phenotype (as the phecode-level concept) -> each ICD main category (3-char stem) -> the specific ICD sub-codes under it (460 ICD-9, 461 ICD-10). When a code is clicked (`focus`), only its main branch is expanded and sub-codes are windowed around it. Other branches collapse to a count. Overflow within a branch is shown as a single "+N more" stub (an edge to a count, not every node). `path` = the root→main→code chain to highlight.""" p = self.phenotypes.get(pid) if not p: return None # cap=0 disables every window: all categories expand with all codes. unlimited = cap == 0 BIG = 10 ** 9 FOCUS_SUBS, MAIN_SUBS, MAX_MAINS = (BIG, BIG, BIG) if unlimited else (14, 6, 20) focus = (focus or "").strip() or None focus_main = focus.split(".")[0] if focus else None root = "pheno" title = p.get("title") or f"Phenotype {pid}" nodes = [{"id": root, "label": title[:60], "sub": "CIPHER phenotype (phecode)", "tier": 0, "current": focus is None, "path": True}] edges, path = [], [root] note = None for g in p.get("associated_code_groups", []): ct = g.get("code_type_id") if ct not in (460, 461): continue system = "ICD-9" if ct == 460 else "ICD-10" sub_type = g.get("sub_type_id") by_main: dict[str, list[dict]] = {} for x in g.get("codes", []): code = (x.get("code") or "").strip() if code: by_main.setdefault(code.split(".")[0], []).append(x) # focus can be a leaf code (clicked in the drawer) or a main category # (clicked to expand it) — both expand that main; a leaf also highlights. group_has_focus = focus is not None and focus_main in by_main focus_is_leaf = group_has_focus and any(x.get("code") == focus for x in by_main[focus_main]) mains = sorted(by_main) sel_mains = mains[:MAX_MAINS] if group_has_focus and focus_main not in sel_mains: sel_mains = sel_mains[:MAX_MAINS - 1] + [focus_main] if len(mains) > len(sel_mains): note = "Some categories are truncated." for cat in sel_mains: # Namespaced, because a category and one of its codes can be the # same string: a phenotype whose code list holds the 3-character # stem "F32" produced a category node and a code node with the # id "ICD-10:F32", and the two drew on top of each other in the # same column. The category rung is its own kind of thing and # gets its own id space. main_id = f"cat:{system}:{cat}" subs = by_main[cat] is_focus_main = group_has_focus and cat == focus_main # When focused, only the focused branch expands; others collapse. collapse = focus is not None and not is_focus_main and not unlimited nodes.append({"id": main_id, "label": cat, "tier": 1, "current": False, "sub": (f"{system} · {len(subs)} code" + ("s" if len(subs) != 1 else "") if collapse else f"{system} category"), "path": is_focus_main, "nav": {"code": cat, "code_type": system}}) edges.append([root, main_id]) if is_focus_main: path.append(main_id) if collapse: continue cap = FOCUS_SUBS if is_focus_main else MAIN_SUBS if is_focus_main and focus_is_leaf and len(subs) > cap: # window around the focus idx = next(i for i, x in enumerate(subs) if x.get("code") == focus) start = max(0, min(idx - cap // 2, len(subs) - cap)) sel = subs[start:start + cap] else: sel = subs[:cap] for x in sel: desc, _, _extra = self._code_display(ct, sub_type, x) code = x["code"] sid = f"{system}:{code}" is_cur = code == focus nodes.append({"id": sid, "label": code, "sub": desc or "", "tier": 2, "current": is_cur, "path": is_cur, "nav": {"code": code, "code_type": system}}) edges.append([main_id, sid]) if is_cur: path.append(sid) if len(subs) > len(sel): # overflow stub, not every node more_id = f"more:{main_id}" nodes.append({"id": more_id, "label": f"+{len(subs) - len(sel)} more", "sub": "", "tier": 2, "more": True}) edges.append([main_id, more_id]) if len(nodes) == 1: return {"available": False, "kind": "phenotype", "code": str(pid), "reason": "This phenotype has no ICD diagnosis codes to chart."} return {"available": True, "kind": "phenotype", "code": str(pid), "title": title, "subtitle": "Phecode → ICD main categories → ICD sub-codes" + (" (click a category to expand it)" if focus else ""), "nodes": nodes, "edges": edges, "path": path, "note": note}