| """
|
| SQLite-based data loader for protein structure data.
|
| Reads from MuSProt.db: edge table (pairwise comparisons from CSV)
|
| and node table (chain annotations from TSV).
|
| """
|
| import json
|
| import pandas as pd
|
| from pathlib import Path
|
| from typing import Any, Dict, List, Optional, Set
|
| import logging
|
|
|
| from app.protein.config import get_database_path, get_summary_path
|
| from app.protein.database import connect_readonly
|
|
|
| logger = logging.getLogger(__name__)
|
|
|
|
|
| def _assign_cluster(tm: float) -> str:
|
| if tm >= 0.99:
|
| return "cluster_ultra_high"
|
| if tm >= 0.95:
|
| return "cluster_very_high"
|
| if tm >= 0.90:
|
| return "cluster_high"
|
| if tm >= 0.80:
|
| return "cluster_medium"
|
| return "cluster_low"
|
|
|
|
|
| class DataManager:
|
| """Manages SQL queries against MuSProt.db edge and node tables."""
|
|
|
| def __init__(self, db_path: Optional[Path] = None, **_kwargs):
|
| self.db_path = db_path or get_database_path()
|
| self._loaded = False
|
| self._total_records = 0
|
| self._summary: Dict[str, Any] = {}
|
|
|
| def load_data(self) -> None:
|
| """Verify DB is accessible and load optional precomputed metadata."""
|
| if self._loaded:
|
| return
|
| summary_path = get_summary_path()
|
| if summary_path:
|
| self._summary = json.loads(summary_path.read_text(encoding="utf-8"))
|
| self._total_records = int(self._summary.get("total_records", 0))
|
|
|
| conn = connect_readonly(self.db_path)
|
| try:
|
| cur = conn.cursor()
|
| cur.execute("SELECT 1 FROM edge LIMIT 1")
|
| if not self._total_records:
|
| logger.warning(
|
| "No musprot_summary.json found; counting edge rows during startup. "
|
| "Publish the summary sidecar to avoid this scan."
|
| )
|
| cur.execute("SELECT COUNT(*) FROM edge")
|
| self._total_records = cur.fetchone()[0]
|
| finally:
|
| conn.close()
|
| self._loaded = True
|
| logger.info(f"DataManager ready: {self._total_records:,} edge records in {self.db_path}")
|
|
|
| def _connect(self):
|
| if not self._loaded:
|
| raise RuntimeError("Call load_data() first.")
|
| return connect_readonly(self.db_path)
|
|
|
| def get_chain_sample(self, limit: int = 500) -> Set[str]:
|
| """Return a small chain sample without loading the full node index."""
|
| conn = self._connect()
|
| try:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT pdb_id, auth_asym_id FROM node "
|
| "WHERE pdb_id IS NOT NULL AND auth_asym_id IS NOT NULL LIMIT ?",
|
| (limit,),
|
| )
|
| return {
|
| f"{str(pdb).upper()}_{str(chain).upper()}"
|
| for pdb, chain in cur.fetchall()
|
| }
|
| finally:
|
| conn.close()
|
|
|
| def get_filters_data(self) -> Dict[str, Any]:
|
| """Return aggregated ranges for filter UI."""
|
| if self._summary:
|
| return {
|
| "tm_min": self._summary["tm_score_range"]["min"],
|
| "tm_max": self._summary["tm_score_range"]["max"],
|
| "rmsd_min": self._summary["rmsd_range"]["min"],
|
| "rmsd_max": self._summary["rmsd_range"]["max"],
|
| "length_min": self._summary["length_range"]["min"],
|
| "length_max": self._summary["length_range"]["max"],
|
| "total_records": self._total_records,
|
| }
|
|
|
| conn = self._connect()
|
| try:
|
| cur = conn.cursor()
|
| cur.execute(
|
| "SELECT MIN(CAST(TM1 AS REAL)), MAX(CAST(TM1 AS REAL)),"
|
| " MIN(CAST(RMSD AS REAL)), MAX(CAST(RMSD AS REAL))"
|
| " FROM edge"
|
| )
|
| tm_min, tm_max, rmsd_min, rmsd_max = cur.fetchone()
|
| finally:
|
| conn.close()
|
|
|
| from app.protein import tsv_loader
|
| lengths = []
|
| for row in tsv_loader._get_index().values():
|
| try:
|
| lengths.append(int(row["sequence_length"]))
|
| except (ValueError, TypeError):
|
| pass
|
|
|
| return {
|
| "tm_min": tm_min or 0.0,
|
| "tm_max": tm_max or 1.0,
|
| "rmsd_min": rmsd_min or 0.0,
|
| "rmsd_max": rmsd_max or 10.0,
|
| "length_min": min(lengths) if lengths else 0,
|
| "length_max": max(lengths) if lengths else 0,
|
| "total_records": self._total_records,
|
| }
|
|
|
| def query_edge(
|
| self,
|
| pdb_id_a: Optional[str] = None,
|
| auth_asym_id_a: Optional[str] = None,
|
| chain_ids: Optional[List[str]] = None,
|
| rmsd_min: Optional[float] = None,
|
| rmsd_max: Optional[float] = None,
|
| tm_min: Optional[float] = None,
|
| tm_max: Optional[float] = None,
|
| limit: int = 1000,
|
| ) -> pd.DataFrame:
|
| """Query edge table with filters. Returns DataFrame with cluster_id and lengths."""
|
| where_clauses: List[str] = []
|
| params: List[Any] = []
|
|
|
| if pdb_id_a and auth_asym_id_a:
|
| where_clauses.append("LOWER(pdb_id_A) = LOWER(?) AND LOWER(auth_asym_id_A) = LOWER(?)")
|
| params.extend([pdb_id_a, auth_asym_id_a])
|
| elif chain_ids:
|
| sub: List[str] = []
|
| for cid in chain_ids:
|
| parts = cid.split("_", 1)
|
| if len(parts) == 2:
|
| sub.append("(LOWER(pdb_id_A) = LOWER(?) AND LOWER(auth_asym_id_A) = LOWER(?))")
|
| params.extend(parts)
|
| if sub:
|
| where_clauses.append("(" + " OR ".join(sub) + ")")
|
|
|
| if rmsd_min is not None:
|
| where_clauses.append("CAST(RMSD AS REAL) >= ?")
|
| params.append(rmsd_min)
|
| if rmsd_max is not None:
|
| where_clauses.append("CAST(RMSD AS REAL) <= ?")
|
| params.append(rmsd_max)
|
| if tm_min is not None:
|
| where_clauses.append("CAST(TM1 AS REAL) >= ?")
|
| params.append(tm_min)
|
| if tm_max is not None:
|
| where_clauses.append("CAST(TM1 AS REAL) <= ?")
|
| params.append(tm_max)
|
|
|
| where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""
|
| sql = f"""
|
| SELECT e.pdb_id_A, e.auth_asym_id_A, e.pdb_id_B, e.auth_asym_id_B,
|
| CAST(e.TM1 AS REAL) AS TM1,
|
| CAST(e.RMSD AS REAL) AS RMSD,
|
| CAST(e.structure_sim AS REAL) AS structure_sim,
|
| e."delta_Rosetta", e."delta_FoldX", e."delta_EvoEF2",
|
| e."delta_RW", e."delta_RW+",
|
| e.state_id_B, e.state_fidelity, e.avg_sim, e.observation_fidelity
|
| FROM edge e
|
| {where_sql}
|
| LIMIT ?
|
| """
|
| params.append(limit)
|
|
|
| conn = self._connect()
|
| try:
|
| df = pd.read_sql_query(sql, conn, params=params)
|
| finally:
|
| conn.close()
|
|
|
|
|
| df = df.rename(columns={"delta_RW+": "delta_RW_plus"})
|
|
|
| df["cluster_id"] = df["TM1"].apply(
|
| lambda t: _assign_cluster(t) if pd.notna(t) else "cluster_low"
|
| )
|
|
|
| from app.protein import tsv_loader
|
| df["length_a"] = [
|
| tsv_loader.get_sequence(r.pdb_id_A, r.auth_asym_id_A)[1]
|
| for r in df.itertuples(index=False)
|
| ]
|
| df["length_b"] = [
|
| tsv_loader.get_sequence(r.pdb_id_B, r.auth_asym_id_B)[1]
|
| for r in df.itertuples(index=False)
|
| ]
|
| node_rows = [
|
| tsv_loader.get_node_row(r.pdb_id_B, r.auth_asym_id_B)
|
| for r in df.itertuples(index=False)
|
| ]
|
| df["experimental_method"] = [row.get("experimental_method") if row else None for row in node_rows]
|
| df["pH"] = [row.get("pH") if row else None for row in node_rows]
|
| df["temp_K"] = [row.get("temp_K") if row else None for row in node_rows]
|
| df["base_label"] = [row.get("base_label") if row else None for row in node_rows]
|
| df["chain_composition"] = [row.get("chain_composition") if row else None for row in node_rows]
|
|
|
| return df
|
|
|
| def get_summary_stats(
|
| self,
|
| chain_ids: Optional[List[str]] = None,
|
| rmsd_min: Optional[float] = None,
|
| rmsd_max: Optional[float] = None,
|
| tm_min: Optional[float] = None,
|
| tm_max: Optional[float] = None,
|
| ) -> Dict[str, Any]:
|
| """Compute aggregate statistics via SQL, with optional filters."""
|
| if not any([chain_ids, rmsd_min, rmsd_max, tm_min, tm_max]) and self._summary:
|
| return self._summary
|
|
|
| where_clauses: List[str] = []
|
| params: List[Any] = []
|
|
|
| if chain_ids:
|
| sub: List[str] = []
|
| for cid in chain_ids:
|
| parts = cid.split("_", 1)
|
| if len(parts) == 2:
|
| sub.append("(LOWER(pdb_id_A) = LOWER(?) AND LOWER(auth_asym_id_A) = LOWER(?))")
|
| params.extend(parts)
|
| if sub:
|
| where_clauses.append("(" + " OR ".join(sub) + ")")
|
|
|
| if rmsd_min is not None:
|
| where_clauses.append("CAST(RMSD AS REAL) >= ?")
|
| params.append(rmsd_min)
|
| if rmsd_max is not None:
|
| where_clauses.append("CAST(RMSD AS REAL) <= ?")
|
| params.append(rmsd_max)
|
| if tm_min is not None:
|
| where_clauses.append("CAST(TM1 AS REAL) >= ?")
|
| params.append(tm_min)
|
| if tm_max is not None:
|
| where_clauses.append("CAST(TM1 AS REAL) <= ?")
|
| params.append(tm_max)
|
|
|
| where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""
|
|
|
| conn = self._connect()
|
| try:
|
| cur = conn.cursor()
|
| cur.execute(
|
| f"SELECT COUNT(*), AVG(CAST(RMSD AS REAL)), AVG(CAST(TM1 AS REAL))"
|
| f" FROM edge {where_sql}",
|
| params,
|
| )
|
| total, avg_rmsd, avg_tm = cur.fetchone()
|
|
|
| cur.execute(
|
| f"SELECT COUNT(DISTINCT LOWER(pdb_id_A) || '_' || LOWER(auth_asym_id_A))"
|
| f" FROM edge {where_sql}",
|
| params,
|
| )
|
| unique_chains = cur.fetchone()[0]
|
|
|
| def _dist(col, bins, labels):
|
| cases = " ".join(
|
| f"WHEN CAST({col} AS REAL) >= {lo} AND CAST({col} AS REAL) < {hi} THEN '{lbl}'"
|
| for (lo, hi), lbl in zip(zip(bins, bins[1:]), labels)
|
| )
|
| sql = (
|
| f"SELECT CASE {cases} ELSE '{labels[-1]}' END AS bucket, COUNT(*)"
|
| f" FROM edge {where_sql} GROUP BY bucket"
|
| )
|
| cur.execute(sql, params)
|
| return {r[0]: r[1] for r in cur.fetchall()}
|
|
|
| rmsd_dist = _dist(
|
| "RMSD",
|
| [0, 0.5, 1.0, 1.5, 2.0, 5.0, 1e9],
|
| ["0-0.5", "0.5-1.0", "1.0-1.5", "1.5-2.0", "2.0-5.0", ">5.0"],
|
| )
|
| tm_dist = _dist(
|
| "TM1",
|
| [0, 0.5, 0.7, 0.85, 0.95, 1.01],
|
| ["0-0.5", "0.5-0.7", "0.7-0.85", "0.85-0.95", "0.95-1.0"],
|
| )
|
| finally:
|
| conn.close()
|
|
|
| from app.protein import tsv_loader
|
| lengths = []
|
| for row in tsv_loader._get_index().values():
|
| try:
|
| lengths.append(int(row["sequence_length"]))
|
| except (ValueError, TypeError):
|
| pass
|
| avg_len = sum(lengths) / len(lengths) if lengths else 0.0
|
| len_dist = {}
|
| for length in lengths:
|
| if length <= 100:
|
| bucket = "0-100"
|
| elif length <= 200:
|
| bucket = "100-200"
|
| elif length <= 300:
|
| bucket = "200-300"
|
| elif length <= 500:
|
| bucket = "300-500"
|
| elif length <= 1000:
|
| bucket = "500-1000"
|
| else:
|
| bucket = ">1000"
|
| len_dist[bucket] = len_dist.get(bucket, 0) + 1
|
|
|
| return {
|
| "total_records": total or 0,
|
| "unique_chains": unique_chains or 0,
|
| "avg_rmsd": avg_rmsd or 0.0,
|
| "avg_tm_score": avg_tm or 0.0,
|
| "avg_sequence_length": avg_len,
|
| "rmsd_distribution": rmsd_dist,
|
| "tm_score_distribution": tm_dist,
|
| "length_distribution": len_dist,
|
| }
|
|
|
| def get_total(self) -> int:
|
| return self._total_records
|
|
|
