{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":116062,"databundleVersionId":14875579}],"dockerImageVersionId":31259,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Super fast DAG score propagation\n\nThis notebook shows a very fast implementation of max score propagation over the Gene Ontology DAG using CSR adjacency and Numba JIT.\n\nThe point is speed. After compilation, propagation on full CAFA6-scale data runs\nin roughly well **under 1 second per ontology** in the usual sparse setting, and stays\n**practical even in a much heavier stress test** with 500 seed terms per protein.\n\nWhy Numba is used here is simple: this operation is dominated by many small graph\ntraversals. In pure Python, that becomes slow because of interpreter overhead. With\nCSR arrays and Numba, the core loop runs as compiled code, so ontology-consistent\nancestor closure becomes a cheap post-processing step instead of a bottleneck.\n\nThis notebook uses CAFA6 data: `data6/Train/train_terms.tsv` and ontology (OBO). It measures propagation time and correctness.","metadata":{}},{"cell_type":"code","source":"import time\nimport numpy as np\nimport numba as nb\nimport pandas as pd\n\nfrom typing import Tuple","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:22:28.768756Z","iopub.execute_input":"2026-04-01T12:22:28.769011Z","iopub.status.idle":"2026-04-01T12:22:31.892161Z","shell.execute_reply.started":"2026-04-01T12:22:28.768983Z","shell.execute_reply":"2026-04-01T12:22:31.891029Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"In CAFA competitions we have a directed acyclic graph (DAG) of ontology terms (Gene Ontology). Edges are directed from child to parent (more specific term to more general term).\n\nIn CAFA-style datasets, GO annotations may be provided without full ancestor closure; training on these raw labels can yield predictions where a child term is present but some (or all) parents are missing. Also, unless ontology consistency is enforced explicitly, a model can assign a higher score to a child than to one or more of its parents.\n\nThe goal is max-propagation upward, with adding missing ancestors: for every protein and every ancestor term, set score to the maximum among all its predicted descendants, and the ancestor’s own original score if it existed.","metadata":{}},{"cell_type":"markdown","source":"## 1. Paths and build ontology from raw OBO","metadata":{}},{"cell_type":"markdown","source":"`_iter_obo_terms` is a streaming OBO parser. For each non-obsolete [Term] stanza it yields:\n- term_id\n- namespace (BP/CC/MF in GO)\n- a list of parents collected from is_a: lines and relationship: part_of lines\n- alt_id values\n\n`parse_obo` constructs for each GO aspect\n- a mapping from each term identifier to the list of its parent term identifiers as defined in the OBO file\n- a lookup table that maps every alternative GO identifier (alt_id) to its corresponding canonical (primary) GO identifier\n\nThen `build_graph` canonizes both children and parents via `_canonize`, unions everything into term_set, creates a stable term_list and term_to_idx, and materializes the parent adjacency as a CSR matrix - the core data structure that makes the Numba kernel fast, because it allows iterating parents of a node with two integer lookups and a contiguous slice.","metadata":{}},{"cell_type":"code","source":"TRAIN_DIR = \"/kaggle/input/cafa-6-protein-function-prediction/Train\"\nTRAIN_TERMS_TSV = f\"{TRAIN_DIR}/train_terms.tsv\"\nONTOLOGY_OBO = f\"{TRAIN_DIR}/go-basic.obo\"\n\nASPECT_TO_NS = {\n \"BP\": \"biological_process\",\n \"CC\": \"cellular_component\",\n \"MF\": \"molecular_function\",\n}\n\nASPECT_CODE_TO_LONG = {\n \"P\": \"BP\",\n \"C\": \"CC\",\n \"F\": \"MF\",\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:22:31.894421Z","iopub.execute_input":"2026-04-01T12:22:31.895026Z","iopub.status.idle":"2026-04-01T12:22:31.901506Z","shell.execute_reply.started":"2026-04-01T12:22:31.894993Z","shell.execute_reply":"2026-04-01T12:22:31.900386Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def _iter_obo_terms(obo_file: str):\n \"\"\"Yield (term_id, namespace, parents, alt_ids) for each non-obsolete [Term] in OBO.\"\"\"\n cur_id, cur_ns, cur_parents, cur_alt, cur_obsolete, in_term = (\n None, None, [], [], False, False\n )\n with open(obo_file, \"r\", encoding=\"utf-8\") as h:\n for raw in h:\n line = raw.strip()\n if not line:\n continue\n if line == \"[Term]\":\n if in_term and cur_id and cur_ns is not None and not cur_obsolete:\n yield cur_id, cur_ns, cur_parents, cur_alt\n cur_id, cur_ns, cur_parents, cur_alt, cur_obsolete, in_term = (\n None, None, [], [], False, True\n )\n continue\n if line.startswith(\"[\") and line.endswith(\"]\") and line != \"[Term]\":\n if in_term and cur_id and cur_ns is not None and not cur_obsolete:\n yield cur_id, cur_ns, cur_parents, cur_alt\n cur_id, cur_ns, cur_parents, cur_alt, cur_obsolete, in_term = (\n None, None, [], [], False, False\n )\n continue\n if not in_term:\n continue\n if line.startswith(\"id: \"):\n cur_id = line[4:].strip()\n elif line.startswith(\"namespace: \"):\n cur_ns = line[11:].strip()\n elif line.startswith(\"alt_id: \"):\n cur_alt.append(line[8:].strip())\n elif line.startswith(\"is_obsolete: \"):\n cur_obsolete = cur_obsolete or (line[13:].strip() == \"true\")\n elif line.startswith(\"is_a: \"):\n cur_parents.append(line[6:].split(\"!\")[0].strip())\n elif line.startswith(\"relationship: part_of \"):\n parts = line.split()\n if len(parts) >= 3:\n cur_parents.append(parts[2].strip())\n if in_term and cur_id and cur_ns is not None and not cur_obsolete:\n yield cur_id, cur_ns, cur_parents, cur_alt\n\n\ndef parse_obo(obo_file: str, namespaces: list[str]) -> Tuple[dict, dict]:\n \"\"\"Parse OBO file into edges per namespace and alt_id->canonical mapping.\"\"\"\n edges_by_ns = {ns: {} for ns in namespaces}\n alt_to_canon = {}\n for term_id, term_ns, parents, alt_ids in _iter_obo_terms(obo_file):\n if term_ns not in namespaces:\n continue\n edges_by_ns[term_ns][term_id] = parents\n for a in alt_ids:\n alt_to_canon[a] = term_id\n return edges_by_ns, alt_to_canon\n\n\ndef _canonize(t: str, alt: dict[str, str]) -> str:\n \"\"\"Return canonical term id; if not in alt map, return t.\"\"\"\n return alt.get(t, t)\n\ndef _pack_csr_from_pairs(\n row_index: np.ndarray,\n col_index: np.ndarray,\n n_prots: int,\n) -> Tuple[np.ndarray, np.ndarray]:\n \"\"\"Build CSR (indptr, indices) from (row, col) pairs; dedup and sort by row then col.\"\"\"\n order = np.lexsort((col_index, row_index))\n row_sorted = row_index[order]\n col_sorted = col_index[order]\n if row_sorted.size > 1:\n keep = np.ones(row_sorted.size, dtype=bool)\n keep[1:] = ~((row_sorted[1:] == row_sorted[:-1]) &\n (col_sorted[1:] == col_sorted[:-1]))\n row_sorted = row_sorted[keep]\n col_sorted = col_sorted[keep]\n counts = np.bincount(row_sorted, minlength=n_prots)\n indptr = np.empty(n_prots + 1, dtype=np.int64)\n indptr[0] = 0\n np.cumsum(counts, out=indptr[1:])\n return indptr, col_sorted.astype(np.int32)\n\ndef build_graph(\n edges: dict[str, list[str]],\n alt_to_canon: dict[str, str],\n) -> Tuple[int, np.ndarray, np.ndarray, list]:\n \"\"\"Build parent CSR (indptr, indices) and term list from edge dict.\"\"\"\n term_set = set()\n child_to_parents = {}\n for child_id, parent_ids in edges.items():\n c = _canonize(child_id, alt_to_canon)\n term_set.add(c)\n child_to_parents.setdefault(c, set()).update(\n _canonize(p, alt_to_canon) for p in parent_ids\n )\n term_set.update(child_to_parents[c])\n term_list = sorted(term_set)\n term_to_idx = {t: i for i, t in enumerate(term_list)}\n row_list, col_list = [], []\n for child, parents in child_to_parents.items():\n if not parents:\n continue\n ci = term_to_idx[child]\n for p in parents:\n row_list.append(ci)\n col_list.append(term_to_idx[p])\n if not row_list:\n raise ValueError(\"No parent edges.\")\n n = len(term_list)\n p_indptr, p_indices = _pack_csr_from_pairs(\n np.array(row_list, dtype=np.int64),\n np.array(col_list, dtype=np.int32),\n n,\n )\n return n, p_indptr, p_indices, term_list","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:22:31.906685Z","iopub.execute_input":"2026-04-01T12:22:31.907053Z","iopub.status.idle":"2026-04-01T12:22:31.940440Z","shell.execute_reply.started":"2026-04-01T12:22:31.907016Z","shell.execute_reply":"2026-04-01T12:22:31.939451Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\naspect_graphs = {}\nnamespaces = list(ASPECT_TO_NS.values())\nedges_by_ns, alt_to_canon = parse_obo(ONTOLOGY_OBO, namespaces)\nfor aspect, ns in ASPECT_TO_NS.items():\n n, p_indptr, p_indices, term_list = build_graph(edges_by_ns[ns], alt_to_canon)\n term_to_idx = {str(t): i for i, t in enumerate(term_list)}\n aspect_graphs[aspect] = {\n \"n_terms\": n,\n \"parents_indptr\": p_indptr,\n \"parents_indices\": p_indices,\n \"term_to_idx\": term_to_idx,\n }\nprint(\"Ontology from\", ONTOLOGY_OBO)\nfor aspect in ASPECT_TO_NS:\n print(f\" {aspect}: n_terms = {aspect_graphs[aspect]['n_terms']}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:22:31.941685Z","iopub.execute_input":"2026-04-01T12:22:31.942037Z","iopub.status.idle":"2026-04-01T12:22:33.106681Z","shell.execute_reply.started":"2026-04-01T12:22:31.941999Z","shell.execute_reply":"2026-04-01T12:22:33.105540Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"
\n 💡 Note: \n The important part here is not only parsing the OBO file, but converting the\nontology into a compact CSR graph once. After that, parent traversal is just array\nindexing over contiguous memory, which is exactly the layout Numba can process very\nfast.\n\nIn other words, this step prepares the data structure that makes the later\npropagation really fast at CAFA scale.\n
","metadata":{}},{"cell_type":"markdown","source":"## 2. Propagation time per ontology\n\nThe main question is whether ontology-consistent propagation can be made cheap at CAFA6 scale. The benchmark below shows that, after JIT compilation, the answer is yes.\n\nHere I build CSR of protein row to term indices using `build_input_csr_for_aspect` (term indices from ontology built from OBO). This CSR means: for protein i, its initial terms are in indices[indptr[i]:indptr[i+1]] with scores in the corresponding slice of scores. Then, I add random scores for the demo. \n\nPer aspect: run propagation, measure time.\n\nThe propagation kernel `propagate_max_up_csr` is a two-pass CSR builder. For each protein row, it first sorts the row’s seed terms by score in descending order. It then performs a parent traversal where each ontology node is assigned at most once: when a node is first reached from a seed of score s, its propagated score is fixed to s, and the node is pushed to expand to its parents. Because seeds are processed from high to low score, the first assignment to any node equals the required maximum over all reachable seeds.\n\nAfter this completes, the set of reached terms (the original terms plus all ancestors) have their propagated max scores.","metadata":{}},{"cell_type":"code","source":"@nb.njit(cache=True)\ndef _argsort_desc(values: np.ndarray) -> np.ndarray:\n order = np.argsort(values)\n return order[::-1]\n\n\n@nb.njit(cache=True)\ndef propagate_max_up_csr(\n indptr: np.ndarray,\n indices: np.ndarray,\n scores: np.ndarray,\n parents_indptr: np.ndarray,\n parents_indices: np.ndarray,\n n_terms: int,\n) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n n_rows = int(indptr.size - 1)\n best = np.full(n_terms, np.float32(-np.inf), dtype=np.float32)\n touched = np.empty(n_terms, dtype=np.int32)\n stack = np.empty(n_terms, dtype=np.int32)\n\n out_counts = np.empty(n_rows, dtype=np.int64)\n\n for row_index in range(n_rows):\n touched_len = 0\n start = int(indptr[row_index])\n end = int(indptr[row_index + 1])\n\n row_terms = indices[start:end].astype(np.int32)\n row_scores = scores[start:end].astype(np.float32)\n order = _argsort_desc(row_scores)\n\n for oi in range(order.size):\n seed = int(row_terms[int(order[oi])])\n seed_score = np.float32(row_scores[int(order[oi])])\n if seed_score <= best[seed]:\n continue\n\n if best[seed] == np.float32(-np.inf):\n touched[touched_len] = np.int32(seed)\n touched_len += 1\n best[seed] = seed_score\n\n stack_len = 0\n stack[stack_len] = np.int32(seed)\n stack_len += 1\n\n sp = 0\n while sp < stack_len:\n term = int(stack[sp])\n sp += 1\n ps = int(parents_indptr[term])\n pe = int(parents_indptr[term + 1])\n for p in range(ps, pe):\n parent = int(parents_indices[p])\n if best[parent] == np.float32(-np.inf):\n best[parent] = seed_score\n touched[touched_len] = np.int32(parent)\n touched_len += 1\n stack[stack_len] = np.int32(parent)\n stack_len += 1\n\n out_counts[row_index] = touched_len\n for i in range(touched_len):\n best[int(touched[i])] = np.float32(-np.inf)\n\n out_indptr = np.empty(n_rows + 1, dtype=np.int64)\n out_indptr[0] = 0\n for i in range(n_rows):\n out_indptr[i + 1] = out_indptr[i] + out_counts[i]\n total = int(out_indptr[n_rows])\n\n out_indices = np.empty(total, dtype=np.int32)\n out_scores = np.empty(total, dtype=np.float32)\n\n for row_index in range(n_rows):\n touched_len = 0\n start = int(indptr[row_index])\n end = int(indptr[row_index + 1])\n\n row_terms = indices[start:end].astype(np.int32)\n row_scores = scores[start:end].astype(np.float32)\n order = _argsort_desc(row_scores)\n\n for oi in range(order.size):\n seed = int(row_terms[int(order[oi])])\n seed_score = np.float32(row_scores[int(order[oi])])\n if seed_score <= best[seed]:\n continue\n\n if best[seed] == np.float32(-np.inf):\n touched[touched_len] = np.int32(seed)\n touched_len += 1\n best[seed] = seed_score\n\n stack_len = 0\n stack[stack_len] = np.int32(seed)\n stack_len += 1\n\n sp = 0\n while sp < stack_len:\n term = int(stack[sp])\n sp += 1\n ps = int(parents_indptr[term])\n pe = int(parents_indptr[term + 1])\n for p in range(ps, pe):\n parent = int(parents_indices[p])\n if best[parent] == np.float32(-np.inf):\n best[parent] = seed_score\n touched[touched_len] = np.int32(parent)\n touched_len += 1\n stack[stack_len] = np.int32(parent)\n stack_len += 1\n\n tmp_terms = touched[:touched_len].copy()\n tmp_terms.sort()\n\n out_pos = int(out_indptr[row_index])\n for j in range(tmp_terms.size):\n term = int(tmp_terms[j])\n out_indices[out_pos + j] = np.int32(term)\n out_scores[out_pos + j] = np.float32(best[term])\n\n for i in range(touched_len):\n best[int(touched[i])] = np.float32(-np.inf)\n\n return out_indptr, out_indices, out_scores","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:22:33.107868Z","iopub.execute_input":"2026-04-01T12:22:33.108348Z","iopub.status.idle":"2026-04-01T12:22:33.210047Z","shell.execute_reply.started":"2026-04-01T12:22:33.108292Z","shell.execute_reply":"2026-04-01T12:22:33.208866Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_df = pd.read_csv(TRAIN_TERMS_TSV, sep=\"\\t\")\ntrain_df = train_df[[\"EntryID\", \"term\", \"aspect\"]].drop_duplicates()\ntrain_df[\"aspect\"] = train_df[\"aspect\"].replace(ASPECT_CODE_TO_LONG)\n\ndef build_input_csr_for_aspect(aspect: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n \"\"\"Build CSR (indptr, indices, scores) from train_terms for one aspect.\"\"\"\n adf = train_df[train_df[\"aspect\"] == aspect]\n term_to_idx = aspect_graphs[aspect][\"term_to_idx\"]\n _, inv = np.unique(adf[\"EntryID\"].values, return_inverse=True)\n term_idx = adf[\"term\"].astype(str).map(term_to_idx)\n valid = term_idx.notna().values\n \n row_idx = np.asarray(inv, dtype=np.int64)[valid]\n col_idx = term_idx.values[valid].astype(np.int32)\n scores = np.random.rand(col_idx.size).astype(np.float32)\n\n order = np.lexsort((col_idx, row_idx))\n row_idx = row_idx[order]\n col_idx = col_idx[order]\n scores = scores[order]\n\n n_prots = int(row_idx.max()) + 1\n counts = np.bincount(row_idx, minlength=n_prots)\n indptr = np.empty(n_prots + 1, dtype=np.int64)\n indptr[0] = 0\n np.cumsum(counts, out=indptr[1:])\n return indptr, col_idx, scores\n\ndef min_med_max(counts, n_prots):\n return (counts.min(), np.median(counts), counts.max()) if n_prots else (\n 0, 0.0, 0\n )\n\nn_rep = 10\nprint(f\"Propagation time per ontology (ms), with {n_rep} repetitions:\\n\")\n\nrows = []\nt0 = time.perf_counter()\nfor aspect in ASPECT_TO_NS:\n graph = aspect_graphs[aspect]\n indptr, indices, scores = build_input_csr_for_aspect(aspect)\n n_prots = indptr.size - 1\n\n in_min, in_med, in_max = min_med_max(np.diff(indptr), n_prots)\n\n _ = propagate_max_up_csr(\n indptr[:2], indices[:2], scores[:2],\n graph[\"parents_indptr\"], graph[\"parents_indices\"], graph[\"n_terms\"]\n )\n\n t0_aspect = time.perf_counter()\n for _ in range(n_rep):\n o_iptr, _, _ = propagate_max_up_csr(\n indptr, indices, scores,\n graph[\"parents_indptr\"], graph[\"parents_indices\"], graph[\"n_terms\"]\n )\n elapsed_s = (time.perf_counter() - t0_aspect) / n_rep\n\n out_min, out_med, out_max = min_med_max(np.diff(o_iptr), n_prots)\n\n rows.append((\n aspect, graph[\"n_terms\"], n_prots, elapsed_s,\n in_min, in_med, in_max, out_min, out_med, out_max\n ))\n\n\ncols = (\n \"aspect\", \"n_terms\", \"n_prots\", \"time_s\",\n \"in_terms_min\", \"in_terms_med\", \"in_terms_max\",\n \"out_terms_min\", \"out_terms_med\", \"out_terms_max\"\n)\n\nresult_df = pd.DataFrame.from_records(rows, columns=cols)\ndisplay(result_df)\n\ntotal_one_run_s = (time.perf_counter() - t0) / n_rep\nprint(f\"\\nTotal time (one run, 3 ontologies): {total_one_run_s:.3f} s\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:22:33.212400Z","iopub.execute_input":"2026-04-01T12:22:33.212706Z","iopub.status.idle":"2026-04-01T12:22:49.539619Z","shell.execute_reply.started":"2026-04-01T12:22:33.212679Z","shell.execute_reply":"2026-04-01T12:22:49.538782Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"
\n 💡 Note: \nThe input per protein is sparse (see `in_terms_*`), but the output is much larger (`out_terms_*`) because full\nancestor closure is added. So the method is computing a real ontology expansion and still finishes in roughly under a second per aspect.\n\nThat is the practical value here: ontology-consistent score propagation is cheap enough to use routinely at inference or post-processing time.\n
","metadata":{}},{"cell_type":"markdown","source":"The next benchmark is a near worst-case scenario: 500 seed terms per protein, which is the CAFA6 submission limit. This tests whether the implementation remains usable when the seed set is very dense.","metadata":{}},{"cell_type":"code","source":"def build_fixed_k_csr_like(\n indptr: np.ndarray,\n n_terms: int,\n k_terms_per_prot: int,\n seed: int,\n) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n n_prots = int(indptr.size - 1)\n if k_terms_per_prot > n_terms:\n raise ValueError(\"k_terms_per_prot must be <= n_terms\")\n\n rng = np.random.default_rng(seed)\n new_indptr = np.empty(n_prots + 1, dtype=np.int64)\n new_indptr[0] = 0\n new_indptr[1:] = np.arange(1, n_prots + 1, dtype=np.int64) * k_terms_per_prot\n\n new_indices = np.empty(n_prots * k_terms_per_prot, dtype=np.int32)\n new_scores = rng.random(n_prots * k_terms_per_prot, dtype=np.float32)\n\n for row_index in range(n_prots):\n row_start = row_index * k_terms_per_prot\n row_end = row_start + k_terms_per_prot\n new_indices[row_start:row_end] = rng.choice(\n n_terms, size=k_terms_per_prot, replace=False\n ).astype(np.int32)\n\n return new_indptr, new_indices, new_scores\n\n\nk_terms_per_prot = 500\nprint(f\"Single-call timing with seeds/prot={k_terms_per_prot}:\\n\")\n\nrows = []\nfor aspect in ASPECT_TO_NS:\n g = aspect_graphs[aspect]\n indptr, indices, scores = build_input_csr_for_aspect(aspect)\n n_prots = indptr.size - 1\n seed = 0 if aspect == \"BP\" else (1 if aspect == \"CC\" else 2)\n\n indptr2, indices2, scores2 = build_fixed_k_csr_like(\n indptr=indptr, n_terms=g[\"n_terms\"], k_terms_per_prot=k_terms_per_prot,\n seed=seed\n )\n\n _ = propagate_max_up_csr(\n indptr2[:2], indices2[:k_terms_per_prot], scores2[:k_terms_per_prot],\n g[\"parents_indptr\"], g[\"parents_indices\"], g[\"n_terms\"]\n )\n\n t0_aspect = time.perf_counter()\n o_iptr, _, _ = propagate_max_up_csr(\n indptr2, indices2, scores2,\n g[\"parents_indptr\"], g[\"parents_indices\"], g[\"n_terms\"]\n )\n elapsed_s = time.perf_counter() - t0_aspect\n\n out_counts = np.diff(o_iptr)\n rows.append((\n aspect, g[\"n_terms\"], n_prots, elapsed_s, k_terms_per_prot,\n out_counts.min(), np.median(out_counts), out_counts.max()\n ))\n\ncols = (\n \"aspect\", \"n_terms\", \"n_prots\", \"time_s\", \"in_terms\",\n \"out_terms_min\", \"out_terms_med\", \"out_terms_max\"\n)\nresult_df = pd.DataFrame.from_records(rows, columns=cols).sort_values(\"aspect\")\nresult_df = result_df.reset_index(drop=True)\ndisplay(result_df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:22:49.540821Z","iopub.execute_input":"2026-04-01T12:22:49.541376Z","iopub.status.idle":"2026-04-01T12:23:57.745461Z","shell.execute_reply.started":"2026-04-01T12:22:49.541304Z","shell.execute_reply":"2026-04-01T12:23:57.744552Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"\n\n
\n 💡 Note: \nThis is a much harder setting: each protein starts with 500 seed terms, so both the graph traversal and the final closure become much larger. Even then, the method stays practical, with runtimes on the order of tens of seconds rather than minutes or hours.\n\nThis shows that the speed is not limited to the easy sparse case. The implementation still scales reasonably when the seed set becomes very dense.\n\nI do not include the Numba compilation time in this benchmark, because compilation happens only once on the first call and is then amortized over thousands of subsequent calls, whereas the benchmark is intended to measure the steady-state runtime of the propagation operator itself.\n
","metadata":{}},{"cell_type":"markdown","source":"## 3. Correctness\n\n- idempotency: propagating twice must not change the result:","metadata":{}},{"cell_type":"code","source":"for aspect in ASPECT_TO_NS:\n g = aspect_graphs[aspect]\n n_terms = g[\"n_terms\"]\n indptr, indices, scores = build_input_csr_for_aspect(aspect)\n o1 = propagate_max_up_csr(\n indptr, indices, scores, g[\"parents_indptr\"], g[\"parents_indices\"], n_terms\n )\n o2 = propagate_max_up_csr(\n o1[0], o1[1], o1[2], g[\"parents_indptr\"], g[\"parents_indices\"], n_terms\n )\n same = (\n np.array_equal(o1[0], o2[0])\n and np.array_equal(o1[1], o2[1])\n and np.allclose(o1[2], o2[2])\n )\n print(f\"{aspect}: {same}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:23:57.746722Z","iopub.execute_input":"2026-04-01T12:23:57.747798Z","iopub.status.idle":"2026-04-01T12:24:00.465486Z","shell.execute_reply.started":"2026-04-01T12:23:57.747752Z","shell.execute_reply":"2026-04-01T12:24:00.464354Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"
\n 💡 Note: \nThis means the operator is stable. Once scores have been propagated to all ancestors, a second pass changes nothing.\n
","metadata":{}},{"cell_type":"markdown","source":"- Edge-consistency: for each protein, for each reached child term, assert that every parent term has a score at least as large.\n- \"Bellman optimality\" check: for each reached term compute the maximum of the term’s original score (if any) and the propagated scores of its children (if any), and then assert their equality.","metadata":{}},{"cell_type":"code","source":"def invert_csr(\n indptr: np.ndarray,\n indices: np.ndarray,\n n_cols: int,\n) -> Tuple[np.ndarray, np.ndarray]:\n \"\"\"Transpose CSR: row->col adjacency becomes col->row (indptr_t, indices_t).\"\"\"\n if indices.size == 0 or n_cols == 0:\n return np.zeros(n_cols + 1, dtype=np.int64), np.empty(0, dtype=np.int32)\n n_prots = int(indptr.size - 1)\n col_counts = np.bincount(indices, minlength=n_cols)\n indptr_t = np.empty(n_cols + 1, dtype=np.int64)\n indptr_t[0] = 0\n np.cumsum(col_counts, out=indptr_t[1:])\n indices_t = np.empty(indices.size, dtype=np.int32)\n write_ptr = indptr_t[:-1].copy()\n for row in range(n_prots):\n start, end = int(indptr[row]), int(indptr[row + 1])\n for pos in range(start, end):\n col = int(indices[pos])\n out_pos = int(write_ptr[col])\n indices_t[out_pos] = np.int32(row)\n write_ptr[col] = out_pos + 1\n return indptr_t, indices_t","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:24:00.466740Z","iopub.execute_input":"2026-04-01T12:24:00.467843Z","iopub.status.idle":"2026-04-01T12:24:00.476276Z","shell.execute_reply.started":"2026-04-01T12:24:00.467796Z","shell.execute_reply":"2026-04-01T12:24:00.475376Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\naspect = \"CC\"\ng = aspect_graphs[aspect]\nn_terms = int(g[\"n_terms\"])\n\nindptr, indices, scores = build_input_csr_for_aspect(aspect)\nout_indptr, out_indices, out_scores = propagate_max_up_csr(\n indptr, indices, scores, g[\"parents_indptr\"], g[\"parents_indices\"], n_terms\n)\n\nchildren_indptr, children_indices = invert_csr(\n g[\"parents_indptr\"], g[\"parents_indices\"], n_terms\n)\n\nn_prots = int(indptr.size - 1)\nrows = np.arange(n_prots)\n\nedge_ok = True\nbellman_ok = True\n\nfor row in rows:\n row = int(row)\n\n in_start, in_end = int(indptr[row]), int(indptr[row + 1])\n out_start, out_end = int(out_indptr[row]), int(out_indptr[row + 1])\n\n in_terms = indices[in_start:in_end]\n in_vals = scores[in_start:in_end]\n\n out_terms = out_indices[out_start:out_end]\n out_vals = out_scores[out_start:out_end]\n\n in_map = {}\n for j in range(int(in_terms.size)):\n t = int(in_terms[j])\n v = np.float32(in_vals[j])\n old = in_map.get(t, np.float32(-np.inf))\n if v > old:\n in_map[t] = v\n\n out_map = {}\n for j in range(int(out_terms.size)):\n t = int(out_terms[j])\n out_map[t] = out_vals[j]\n\n for child in out_map.keys():\n child_score = out_map[child]\n ps = int(g[\"parents_indptr\"][child])\n pe = int(g[\"parents_indptr\"][child + 1])\n for pos in range(ps, pe):\n parent = int(g[\"parents_indices\"][pos])\n parent_score = out_map.get(parent, np.float32(-np.inf))\n if parent_score < child_score:\n edge_ok = False\n break\n if not edge_ok:\n break\n if not edge_ok:\n break\n\n for term, term_out in out_map.items():\n rhs = in_map.get(term, np.float32(-np.inf))\n cs = int(children_indptr[term])\n ce = int(children_indptr[term + 1])\n for pos in range(cs, ce):\n child = int(children_indices[pos])\n child_out = out_map.get(child, np.float32(-np.inf))\n if child_out > rhs:\n rhs = child_out\n if not np.isclose(term_out, rhs, rtol=0, atol=1e-6):\n bellman_ok = False\n break\n if not bellman_ok:\n break\n\nprint(\n f\"{aspect} correctness on {n_prots} proteins: \"\n f\"edge_monotone={edge_ok}, bellman_optimality={bellman_ok}\"\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-01T12:24:00.477407Z","iopub.execute_input":"2026-04-01T12:24:00.477873Z","iopub.status.idle":"2026-04-01T12:25:24.066446Z","shell.execute_reply.started":"2026-04-01T12:24:00.477832Z","shell.execute_reply":"2026-04-01T12:25:24.065371Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"
\n 💡 Note: \n \nSo this check says the method is not only fast, but exact.\n
","metadata":{}}]}