Precompute inner graphs and drop duplicate edge rows

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

backend/app/protein/api/routes/explore.py

@@ -133,9 +133,11 @@ async def get_inner_graph(
     max_states: int = Query(7, ge=1, le=40),
 ):
     """Inner layer for one sequence: observations of its largest states and their pairwise similarities."""
-    data = _index_required(
-        await run_in_threadpool(catalog.inner_graph_nodes, sequence_id.upper(), max_nodes, max_states)
-    )
+    sid = sequence_id.upper()
+    pre = await run_in_threadpool(catalog.precomputed_inner_graph, sid, max_nodes, max_states)
+    if pre:
+        return pre
+    data = _index_required(await run_in_threadpool(catalog.inner_graph_nodes, sid, max_nodes, max_states))
     if not data:
         raise HTTPException(status_code=404, detail=f"Sequence {sequence_id} not found")
     chains = [(n["pdb_id"], n["auth_asym_id"]) for n in data["nodes"]]

backend/app/protein/catalog.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 import json
 import re
 import sqlite3
+import zlib
 from functools import lru_cache
 from typing import Any, Dict, List, Optional, Tuple
 
@@ -408,9 +409,51 @@ def outer_graph(hsets: List[str], per_group: int = 14, include: List[str] = ())
     return {"groups": groups, "names": names["homologies"]}
 
 
-def inner_graph_nodes(sequence_id: str, max_nodes: int = 60, max_states: int = 7) -> Optional[Dict[str, Any]]:
-    """Pick up to max_nodes observations from the sequence's max_states largest states."""
+FIDELITY_NAMES = ["identical", "low", "medium", "high", None]
+
+
+def precomputed_inner_graph(sequence_id: str, max_nodes: int, max_states: int) -> Optional[Dict[str, Any]]:
+    """Inner graph stored by build_index.py (one row read), or None if not precomputed."""
     conn = _connect()
+    if conn is None:
+        return None
+    try:
+        try:
+            row = conn.execute(
+                "SELECT nodes, edges FROM inner_graph WHERE sequence_id = ? AND max_nodes = ? AND max_states = ?",
+                (sequence_id, max_nodes, max_states),
+            ).fetchone()
+        except sqlite3.OperationalError:  # older index without the table
+            return None
+        if row is None:
+            return None
+        base = inner_graph_nodes(sequence_id, max_nodes, max_states, conn=conn, select=False)
+    finally:
+        conn.close()
+    by_key = {(m["pdb_id"], m["auth_asym_id"]): m for m in base["members"]}
+    nodes = [by_key[(p, c)] for p, c in json.loads(row["nodes"]) if (p, c) in by_key]
+    edges = [
+        {"s": i, "t": j, "sim": None if sim < 0 else sim / 1000, "fid": FIDELITY_NAMES[f]}
+        for i, j, sim, f in json.loads(zlib.decompress(row["edges"]))
+    ]
+    base.pop("members")
+    base.update(nodes=nodes, edges=edges)
+    return base
+
+
+def inner_graph_nodes(
+    sequence_id: str,
+    max_nodes: int = 60,
+    max_states: int = 7,
+    conn: Optional[sqlite3.Connection] = None,
+    select: bool = True,
+) -> Optional[Dict[str, Any]]:
+    """Pick up to max_nodes observations from the sequence's max_states largest states.
+
+    The selection rule is mirrored by build_index.select_inner_nodes; keep them in sync.
+    """
+    own = conn is None
+    conn = conn or _connect()
     if conn is None:
         return None
     try:
@@ -435,7 +478,18 @@ def inner_graph_nodes(sequence_id: str, max_nodes: int = 60, max_states: int = 7
             )
         ]
     finally:
-        conn.close()
+        if own:
+            conn.close()
+
+    shown_states = {s["state_id"] for s in states[:MAX_HEATMAP_STATES]}
+    base = {
+        "sequence": _sequence_dict(row),
+        "n_total": len(members),
+        "states": states,
+        "state_pairs": [p for p in pairs if p["a"] in shown_states and p["b"] in shown_states],
+    }
+    if not select:
+        return {**base, "members": members}
 
     order = [s["state_id"] for s in states[:max_states]]
     pool = [m for m in members if m["state_id"] in set(order)]
@@ -457,11 +511,4 @@ def inner_graph_nodes(sequence_id: str, max_nodes: int = 60, max_states: int = 7
             if not added:
                 break
             depth += 1
-    shown_states = {s["state_id"] for s in states[:MAX_HEATMAP_STATES]}
-    return {
-        "sequence": _sequence_dict(row),
-        "nodes": chosen,
-        "n_total": len(members),
-        "states": states,
-        "state_pairs": [p for p in pairs if p["a"] in shown_states and p["b"] in shown_states],
-    }
+    return {**base, "nodes": chosen}

backend/app/protein/records.py

@@ -186,7 +186,12 @@ def get_transitions(pdb_id: str, chain: str, limit: int = 20000) -> List[Dict[st
         conn.close()
 
     out = []
+    seen = set()
     for r in rows:
+        key = (r["pdb_id_B"].lower(), r["auth_asym_id_B"])
+        if key in seen:  # the edge table holds some exact duplicate rows
+            continue
+        seen.add(key)
         sim = num(r["pair_similarity"])
         out.append({
             "pdb_id": r["pdb_id_B"].lower(),
@@ -230,6 +235,7 @@ def pair_similarities(chains: List[tuple]) -> List[Dict[str, Any]]:
     """Undirected similarity edges among a set of (pdb_id, chain) observations of one sequence."""
     keep = {(p.lower(), c): i for i, (p, c) in enumerate(chains)}
     out = []
+    seen = set()
     conn = _connect()
     try:
         for i, (pdb, chain) in enumerate(chains):
@@ -239,7 +245,8 @@ def pair_similarities(chains: List[tuple]) -> List[Dict[str, Any]]:
                 (pdb, chain),
             ):
                 j = keep.get((r[0].lower(), r[1]))
-                if j is not None and j > i:
+                if j is not None and j > i and (i, j) not in seen:
+                    seen.add((i, j))
                     out.append({"s": i, "t": j, "sim": num(r[2]), "fid": text(r[3])})
     finally:
         conn.close()

backend/scripts/build_index.py

@@ -35,6 +35,9 @@ from pathlib import Path
 FUNCTION_TEXT_TOP_N = 3
 FUNCTION_TEXT_MAX_CHARS = 600
 MAX_PAIR_STATES = 40  # state_pair rows are kept only among a sequence's largest states
+INNER_MAX_NODES = 60  # precomputed inner graph: observations per sequence …
+INNER_MAX_STATES = 7  # … drawn from its largest states (must match app/protein/catalog.py)
+FIDELITY_CODE = {"identical": 0, "low": 1, "medium": 2, "high": 3}
 
 
 def log(msg: str) -> None:
@@ -512,6 +515,100 @@ def order_stats(src: sqlite3.Connection, limit: int | None) -> dict:
     }
 
 
+def select_inner_nodes(members: list, states: list, max_nodes: int, max_states: int) -> list:
+    """Round-robin over the largest states. Mirrors catalog.inner_graph_nodes exactly."""
+    order = [st for st, _ in states[:max_states]]
+    keep = set(order)
+    pool = [m for m in members if m[0] in keep]
+    if len(pool) <= max_nodes:
+        return pool
+    by_state: dict = {}
+    for m in pool:
+        by_state.setdefault(m[0], []).append(m)
+    chosen, depth = [], 0
+    while len(chosen) < max_nodes:
+        added = False
+        for st in order:
+            bucket = by_state.get(st, [])
+            if depth < len(bucket) and len(chosen) < max_nodes:
+                chosen.append(bucket[depth])
+                added = True
+        if not added:
+            break
+        depth += 1
+    return chosen
+
+
+def build_inner_graphs(src: sqlite3.Connection, dst: sqlite3.Connection) -> None:
+    """Precompute each sequence's inner-graph sample and its pairwise similarities.
+
+    Serving these live needs one edge-index lookup per node, which is slow when the
+    DB sits on a network mount; stored here it is a single row read.
+    """
+    import zlib
+
+    states: dict = defaultdict(list)
+    for seq, st, n in dst.execute(
+        "SELECT sequence_id, state_id, n_members FROM state ORDER BY sequence_id, n_members DESC, state_id"
+    ):
+        states[seq].append((st, n))
+    members: dict = defaultdict(list)
+    for seq, st, pdb, chain in dst.execute(
+        "SELECT sequence_id, state_id, pdb_id, auth_asym_id FROM member "
+        "ORDER BY sequence_id, state_id, resolution IS NULL, resolution"
+    ):
+        members[seq].append((st, pdb.lower(), chain))
+
+    src.execute("CREATE TEMP TABLE sel (pdb TEXT, chain TEXT, seq TEXT, idx INTEGER, PRIMARY KEY (pdb, chain))")
+    nodes_by_seq = {}
+    rows = []
+    for seq, mem in members.items():
+        chosen = select_inner_nodes(mem, states[seq], INNER_MAX_NODES, INNER_MAX_STATES)
+        nodes_by_seq[seq] = [[pdb, chain] for _, pdb, chain in chosen]
+        rows.extend((pdb, chain, seq, i) for i, (_, pdb, chain) in enumerate(chosen))
+    src.executemany("INSERT OR IGNORE INTO sel VALUES (?,?,?,?)", rows)
+    log(f"  {len(rows):,} sampled observations over {len(nodes_by_seq):,} sequences")
+
+    edges: dict = defaultdict(list)
+    seen: set = set()
+    n = 0
+    cur = src.execute("""
+        SELECT a.seq, a.idx, b.idx, e.pair_similarity, e.pair_fidelity
+        FROM temp.sel a
+        JOIN edge e ON e.pdb_id_A = a.pdb COLLATE NOCASE AND e.auth_asym_id_A = a.chain
+        JOIN temp.sel b ON b.pdb = lower(e.pdb_id_B) AND b.chain = e.auth_asym_id_B AND b.seq = a.seq
+        WHERE a.idx < b.idx
+    """)
+    while True:
+        batch = cur.fetchmany(200000)
+        if not batch:
+            break
+        for seq, i, j, sim, fid in batch:
+            if (seq, i, j) in seen:  # the edge table holds some exact duplicate rows
+                continue
+            seen.add((seq, i, j))
+            simf = to_float(sim)
+            edges[seq].append([i, j, -1 if simf is None else round(simf * 1000), FIDELITY_CODE.get(fid, 4)])
+        n += len(batch)
+        log(f"  inner edges: {n:,}")
+
+    dst.executescript(
+        """
+        DROP TABLE IF EXISTS inner_graph;
+        CREATE TABLE inner_graph (sequence_id TEXT PRIMARY KEY, max_nodes INTEGER, max_states INTEGER,
+                                  nodes TEXT, edges BLOB);
+        """
+    )
+    dst.executemany(
+        "INSERT INTO inner_graph VALUES (?,?,?,?,?)",
+        (
+            (seq, INNER_MAX_NODES, INNER_MAX_STATES, json.dumps(nodes),
+             zlib.compress(json.dumps(edges.get(seq, []), separators=(",", ":")).encode(), 6))
+            for seq, nodes in nodes_by_seq.items()
+        ),
+    )
+
+
 # ─────────────────────────────────────────────────────────────── overview
 def outer_graph_stats(dst: sqlite3.Connection, n_sequences: int) -> dict:
     rows = dst.execute(
@@ -607,6 +704,8 @@ def main():
     ap.add_argument("--out-dir", type=Path, required=True)
     ap.add_argument("--limit", type=int, default=None, help="debug: only read the first N rows")
     ap.add_argument("--ecod", type=Path, default=None, help="ecod.latest.domains.txt (for ECOD names)")
+    ap.add_argument("--inner-only", action="store_true",
+                    help="only (re)build the inner_graph table in an existing MuSProt-index.db")
     args = ap.parse_args()
 
     args.out_dir.mkdir(parents=True, exist_ok=True)
@@ -617,6 +716,17 @@ def main():
     src = sqlite3.connect(f"file:{args.db}?mode=ro", uri=True)
     src.execute("PRAGMA temp_store = FILE")
     src.execute("PRAGMA cache_size = -2000000")
+
+    if args.inner_only:
+        dst = sqlite3.connect(index_path)
+        log("inner graphs")
+        build_inner_graphs(src, dst)
+        dst.commit()
+        dst.execute("VACUUM")
+        dst.close()
+        log(f"done → {index_path} ({index_path.stat().st_size / 1e6:.1f} MB)")
+        return
+
     dst = sqlite3.connect(tmp_path)
     dst.execute("PRAGMA journal_mode = OFF")
     dst.execute("PRAGMA synchronous = OFF")
@@ -637,6 +747,8 @@ def main():
     edges = edge_histograms(src, args.limit)
     log("pass 4/4: order / disorder labels")
     order = order_stats(src, args.limit)
+    log("inner graphs")
+    build_inner_graphs(src, dst)
 
     overview = build_overview(seqs, ov, n_nodes, edges, state_fid)
     overview["outer_graph"] = outer_graph_stats(dst, len(seqs))

frontend/src/components/KnowledgeGraph.tsx

@@ -1,8 +1,8 @@
-import { useMemo, useState } from 'react';
+import { useEffect, useMemo, useState } from 'react';
 import { Link, useNavigate } from 'react-router-dom';
 import { ArrowRight } from 'lucide-react';
 import { api, type InnerGraph, type OuterGraph } from '../lib/api';
-import { useApi } from '../lib/hooks';
+import { prefetch, useApi } from '../lib/hooks';
 import { FIDELITY_INFO, fidelityColor } from '../lib/fidelity';
 import { chainLabel, chainPath, cleanFunction, fmtInt, fmtNum, methodShort, sequencePath } from '../lib/format';
 import { ForceGraph, type FGGroup, type FGLink, type FGNode } from './ForceGraph';
@@ -234,6 +234,20 @@ export function KnowledgeGraph({ featured }: { featured: Array<{ id: string; nam
   const inner = useApi(`inner:${selected}`, (sig) => api.innerGraph(selected, 60, sig));
   const group = outer.data?.groups[0];
   const g = inner.data;
+
+  // once the first graph is up, warm the other examples so switching is instant
+  const firstReady = !!inner.data;
+  useEffect(() => {
+    if (!firstReady) return;
+    const t = setTimeout(() => {
+      for (const f of featured) {
+        prefetch(`inner:${f.id}`, () => api.innerGraph(f.id, 60));
+        prefetch(`outer1:${f.id}`, () => api.outerGraph([f.id]));
+      }
+    }, 800);
+    return () => clearTimeout(t);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [firstReady]);
   const selectedSeq = group?.sequences.find((s) => s.sequence_id === selected) ?? g?.sequence;
 
   return (

frontend/src/lib/hooks.ts

@@ -49,6 +49,16 @@ export function useApi<T>(key: string | null, fetcher: (signal: AbortSignal) =>
   return state;
 }
 
+/** Warm the useApi cache for a key (e.g. the next thing a user is likely to open). */
+export function prefetch<T>(key: string, fetcher: () => Promise<T>): void {
+  if (cache.has(key)) return;
+  fetcher()
+    .then((data) => cache.set(key, data))
+    .catch(() => {
+      /* best effort */
+    });
+}
+
 /** Resolve a best-effort promise (e.g. an RCSB title) into state. */
 export function usePromise<T>(factory: () => Promise<T> | null, deps: unknown[]): T | null {
   const [value, setValue] = useState<T | null>(null);