if-split 0.1.0__py3-none-any.whl

ifsplit/parse.py

@@ -0,0 +1,111 @@
+"""Stage 3 - Filter candidates on metadata (no coordinates).
+
+Operates on the records in ``candidates.jsonl``. Drops entries that:
+
+- have no protein polymer entity (``no_protein_entity``),
+- have a protein entity but no usable sequence (``no_protein_sequence``),
+- exceed the residue cap (``too_large``): ``total_residues >= max_total_residues``,
+- violate an (optional) wwPDB validation-report quality cap — clashscore, R-free,
+  Ramachandran/rotamer/RSRZ outliers — or lack a report when one is required.
+
+When ``use_biological_assembly`` the residue count is taken from assembly 1
+(``<entry>-1``); otherwise the deposited polymer monomer count is used. Quality
+metrics come from the metadata API (no coordinates); a cap fires only when both
+the cap and the metric are present, so a missing metric never drops an entry.
+Every drop is recorded with its reason so the build is auditable.
+"""
+
+from __future__ import annotations
+
+from .config import Config
+from .schema import CandidateRecord
+
+DROP_NO_PROTEIN = "no_protein_entity"
+DROP_NO_SEQUENCE = "no_protein_sequence"
+DROP_TOO_LARGE = "too_large"
+DROP_CLASHSCORE = "clashscore_too_high"
+DROP_RFREE = "rfree_too_high"
+DROP_RAMACHANDRAN = "ramachandran_outliers_too_high"
+DROP_ROTAMER = "rotamer_outliers_too_high"
+DROP_RSRZ = "rsrz_outliers_too_high"
+DROP_NO_VALIDATION = "no_validation_report"
+
+
+def assembly1_residue_count(record: CandidateRecord) -> int | None:
+    """Residue count of biological assembly 1 (id ending ``-1``), else smallest."""
+    if not record.assemblies:
+        return None
+    for aid in sorted(record.assemblies):
+        if aid.endswith("-1"):
+            return record.assemblies[aid]
+    return record.assemblies[sorted(record.assemblies)[0]]
+
+
+def total_residues(record: CandidateRecord, cfg: Config) -> int | None:
+    """Residue count used for the size filter, per the assembly config."""
+    if cfg.use_biological_assembly:
+        count = assembly1_residue_count(record)
+        if count is not None:
+            return count
+    return record.deposited_residues
+
+
+def quality_drop(record: CandidateRecord, cfg: Config) -> tuple[str, float] | None:
+    """First validation-report cap this record violates, else ``None``.
+
+    A cap fires only when both the cap and the metric are present and the metric
+    exceeds the cap; a missing metric never drops the entry. With
+    ``require_validation_report`` an entry that has no report at all is dropped.
+    Returns ``(reason, value)`` so the drop log records the offending number.
+    """
+    q = record.quality
+    if cfg.require_validation_report and not q.has_report:
+        return (DROP_NO_VALIDATION, 0.0)
+    checks = (
+        (cfg.max_clashscore, q.clashscore, DROP_CLASHSCORE),
+        (cfg.max_rfree, q.rfree, DROP_RFREE),
+        (cfg.max_ramachandran_outlier_pct, q.ramachandran_outlier_pct, DROP_RAMACHANDRAN),
+        (cfg.max_rotamer_outlier_pct, q.rotamer_outlier_pct, DROP_ROTAMER),
+        (cfg.max_rsrz_outlier_pct, q.rsrz_outlier_pct, DROP_RSRZ),
+    )
+    for cap, value, reason in checks:
+        if cap is not None and value is not None and value > cap:
+            return (reason, value)
+    return None
+
+
+def filter_candidates(
+    records: list[CandidateRecord], cfg: Config
+) -> tuple[list[CandidateRecord], list[dict]]:
+    """Return ``(kept, drops)`` where drops is a list of ``{entry_id, reason, ...}``."""
+    kept: list[CandidateRecord] = []
+    drops: list[dict] = []
+    for r in records:
+        proteins = [e for e in r.polymer_entities if e.is_protein]
+        if not proteins:
+            drops.append({"entry_id": r.entry_id, "reason": DROP_NO_PROTEIN})
+            continue
+        if not any(e.seq for e in proteins):
+            drops.append({"entry_id": r.entry_id, "reason": DROP_NO_SEQUENCE})
+            continue
+        tr = total_residues(r, cfg)
+        if tr is not None and tr >= cfg.max_total_residues:
+            drops.append({"entry_id": r.entry_id, "reason": DROP_TOO_LARGE, "residues": tr})
+            continue
+        qd = quality_drop(r, cfg)
+        if qd is not None:
+            reason, value = qd
+            drops.append({"entry_id": r.entry_id, "reason": reason, "value": value})
+            continue
+        kept.append(r)
+    kept.sort(key=lambda r: r.entry_id)
+    drops.sort(key=lambda d: d["entry_id"])
+    return kept, drops
+
+
+def drop_summary(drops: list[dict]) -> dict[str, int]:
+    """Count drops by reason (deterministic order via sorted keys downstream)."""
+    out: dict[str, int] = {}
+    for d in drops:
+        out[d["reason"]] = out.get(d["reason"], 0) + 1
+    return out

ifsplit/rcsb.py

@@ -0,0 +1,251 @@
+"""Thin, polite RCSB API client (Search v2 + Data GraphQL).
+
+No coordinates are ever fetched here — only metadata and sequences (see
+PLAN.md §1.5). Endpoints are centralized so field paths live in one place, and
+all requests retry with backoff on transient failures (429 / 5xx / network).
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Iterator
+
+import httpx
+
+from .config import Config
+
+SEARCH_URL = "https://search.rcsb.org/rcsbsearch/v2/query"
+DATA_GRAPHQL_URL = "https://data.rcsb.org/graphql"
+
+# Conservative page/batch sizes: large enough to keep request counts sane,
+# small enough to stay friendly and within response-size limits. The Data API
+# caps batch endpoints at 1000 ids.
+SEARCH_PAGE_ROWS = 5000
+DATA_BATCH_SIZE = 200
+
+_RETRY_STATUS = {429, 500, 502, 503, 504}
+
+# GraphQL: everything Stages 3-6 need, no coordinates. Validated live against
+# 4HHB / 1A1F / 1IEP. Notable curation signals:
+#   - rcsb_cluster_membership: precomputed cluster id per identity level
+#     (30/50/70/90/95/100) for protein entities -> no cluster-file download.
+#   - rcsb_entry_info.nonpolymer_bound_components: comp ids that actually contact
+#     the protein (the cheap buffer-vs-ligand gate; e.g. 4HHB -> ["HEM"], its
+#     PO4/Cl buffer is absent).
+#   - rcsb_binding_affinity.comp_id: comps with a measured affinity (sparse but a
+#     strong positive "this is a real ligand" signal).
+#   - pdbx_vrpt_summary_{geometry,diffraction,em}: wwPDB validation-report metrics
+#     (clashscore, Ramachandran/rotamer outliers, R-free, RSRZ). Geometry is
+#     reported for X-ray AND EM; diffraction is X-ray-only; each comes back as a
+#     1-element list. Metadata, not coordinates — keeps the no-download invariant.
+#   - rcsb_assembly_info.num_prot_na_interface_entities: RCSB-computed count of
+#     protein<->nucleic-acid interface entities in the assembly; > 0 verifies a real
+#     protein/NA contact (the holo gate for the nucleotide class). A single integer
+#     in the assembly-info block — far cheaper than listing every interface object
+#     (a ribosome has hundreds). Present for X-ray AND EM.
+_ENTRY_QUERY = """
+query($ids: [String!]!) {
+  entries(entry_ids: $ids) {
+    rcsb_id
+    exptl { method }
+    rcsb_entry_info {
+      resolution_combined
+      deposited_polymer_monomer_count
+      nonpolymer_bound_components
+    }
+    rcsb_accession_info { initial_release_date }
+    rcsb_binding_affinity { comp_id }
+    pdbx_vrpt_summary_geometry {
+      clashscore
+      percent_ramachandran_outliers
+      percent_rotamer_outliers
+    }
+    pdbx_vrpt_summary_diffraction {
+      DCC_Rfree
+      percent_RSRZ_outliers
+    }
+    pdbx_vrpt_summary_em {
+      atom_inclusion_backbone
+    }
+    polymer_entities {
+      rcsb_id
+      entity_poly {
+        rcsb_entity_polymer_type
+        pdbx_seq_one_letter_code_can
+      }
+      rcsb_cluster_membership { cluster_id identity }
+    }
+    nonpolymer_entities {
+      nonpolymer_comp { chem_comp { id name formula type } }
+    }
+    assemblies {
+      rcsb_id
+      rcsb_assembly_info {
+        polymer_monomer_count
+        num_prot_na_interface_entities
+      }
+    }
+  }
+}
+"""
+
+
+class RcsbError(RuntimeError):
+    """Raised when an RCSB request fails after exhausting retries."""
+
+
+class RcsbClient:
+    """Minimal client for the two RCSB services IF-Split needs."""
+
+    def __init__(
+        self,
+        *,
+        timeout: float = 60.0,
+        max_retries: int = 4,
+        backoff_base: float = 1.5,
+        sleep=time.sleep,
+    ) -> None:
+        self._client = httpx.Client(
+            timeout=timeout,
+            headers={"User-Agent": "IF-Split/0.1 (reproducible PDB splitter)"},
+        )
+        self._max_retries = max_retries
+        self._backoff_base = backoff_base
+        self._sleep = sleep
+
+    def __enter__(self) -> RcsbClient:
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self._client.close()
+
+    # --- low-level POST with retry/backoff ---
+    def _post(self, url: str, json_body: dict) -> httpx.Response:
+        last_exc: Exception | None = None
+        for attempt in range(self._max_retries + 1):
+            try:
+                resp = self._client.post(url, json=json_body)
+            except httpx.HTTPError as exc:  # network/timeout
+                last_exc = exc
+            else:
+                if resp.status_code not in _RETRY_STATUS:
+                    return resp
+                last_exc = RcsbError(f"{url} -> HTTP {resp.status_code}")
+            if attempt < self._max_retries:
+                self._sleep(self._backoff_base**attempt)
+        raise RcsbError(f"request to {url} failed after retries: {last_exc}")
+
+    # --- Search API: entry IDs matching the snapshot filters ---
+    def _search_query_body(self, cfg: Config) -> dict:
+        cutoff = f"{cfg.snapshot_date.isoformat()}T23:59:59Z"
+        return {
+            "query": {
+                "type": "group",
+                "logical_operator": "and",
+                "nodes": [
+                    {
+                        "type": "terminal",
+                        "service": "text",
+                        "parameters": {
+                            "attribute": "exptl.method",
+                            "operator": "in",
+                            "value": list(cfg.experimental_methods),
+                        },
+                    },
+                    {
+                        "type": "terminal",
+                        "service": "text",
+                        "parameters": {
+                            "attribute": "rcsb_entry_info.resolution_combined",
+                            "operator": "less_or_equal",
+                            "value": cfg.resolution_max_A,
+                        },
+                    },
+                    {
+                        "type": "terminal",
+                        "service": "text",
+                        "parameters": {
+                            "attribute": "rcsb_accession_info.initial_release_date",
+                            "operator": "less_or_equal",
+                            "value": cutoff,
+                        },
+                    },
+                ],
+            },
+            "return_type": "entry",
+        }
+
+    def count_entries(self, cfg: Config) -> int:
+        """Total entries matching the snapshot filters (no paging)."""
+        body = self._search_query_body(cfg)
+        body["request_options"] = {"return_counts": True}
+        resp = self._post(SEARCH_URL, body)
+        resp.raise_for_status()
+        return int(resp.json()["total_count"])
+
+    def search_entry_ids(
+        self, cfg: Config, limit: int | None = None, *, progress=None
+    ) -> list[str]:
+        """All matching entry IDs, sorted ascending for determinism.
+
+        ``limit`` (dev convenience) takes the first N in sorted order, so a
+        limited run is itself reproducible. ``progress`` (optional) is called with
+        a status string after each page, so a full-PDB enumeration isn't silent.
+        """
+        ids: list[str] = []
+        start = 0
+        while True:
+            rows = SEARCH_PAGE_ROWS
+            if limit is not None:
+                remaining = limit - len(ids)
+                if remaining <= 0:
+                    break
+                rows = min(rows, remaining)
+            body = self._search_query_body(cfg)
+            body["request_options"] = {
+                "paginate": {"start": start, "rows": rows},
+                "sort": [
+                    {
+                        "sort_by": "rcsb_entry_container_identifiers.entry_id",
+                        "direction": "asc",
+                    }
+                ],
+            }
+            resp = self._post(SEARCH_URL, body)
+            if resp.status_code == 204:  # no (more) results
+                break
+            resp.raise_for_status()
+            page = resp.json().get("result_set", [])
+            if not page:
+                break
+            ids.extend(hit["identifier"] for hit in page)
+            start += len(page)
+            if progress:
+                progress(f"search: {len(ids)} entry ids found...")
+            if len(page) < rows:
+                break
+        return ids
+
+    # --- Data API: batched metadata enrichment ---
+    def fetch_entries(self, ids: list[str]) -> Iterator[dict]:
+        """Yield raw Data-API entry objects for ``ids``, batched."""
+        for batch in _chunks(ids, DATA_BATCH_SIZE):
+            resp = self._post(
+                DATA_GRAPHQL_URL,
+                {"query": _ENTRY_QUERY, "variables": {"ids": batch}},
+            )
+            resp.raise_for_status()
+            payload = resp.json()
+            if payload.get("errors"):
+                raise RcsbError(f"GraphQL errors: {payload['errors'][:1]}")
+            for entry in payload["data"]["entries"]:
+                if entry is not None:
+                    yield entry
+
+
+def _chunks(seq: list[str], size: int) -> Iterator[list[str]]:
+    for i in range(0, len(seq), size):
+        yield seq[i : i + size]

ifsplit/schema.py

@@ -0,0 +1,241 @@
+"""Candidate-record schema + canonical serialization.
+
+A ``CandidateRecord`` is one entry's snapshot metadata — everything Stages 3-6
+need, with no coordinates. ``candidates.jsonl`` is the canonical, byte-stable
+serialization of these records (sorted entries, sorted keys), which is what the
+snapshot lock hashes.
+
+PDB-ID compatibility: identifiers (entry_id, entity_id) are stored *verbatim* as
+returned by the RCSB Data API in ``rcsb_id`` — never sliced, length-validated, or
+case-folded. This makes the schema agnostic to legacy 4-character IDs (``4HHB``,
+entity ``4HHB_1``) and the extended ``pdb_xxxxxxxx`` form alike.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from collections.abc import Iterable
+
+from pydantic import BaseModel, ConfigDict
+
+
+class PolymerEntity(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    entity_id: str  # RCSB rcsb_id, verbatim (e.g. "4HHB_1")
+    polymer_type: str  # rcsb_entity_polymer_type: Protein / DNA / RNA / NA-hybrid / Other
+    seq_len: int
+    seq: str
+    # RCSB precomputed cluster ids by identity level, e.g. {30: 48, 95: 1239}.
+    # Empty for non-protein entities (RCSB clusters proteins only).
+    cluster_ids: dict[int, int] = {}
+
+    @property
+    def is_protein(self) -> bool:
+        return "PROTEIN" in self.polymer_type.upper()
+
+    @property
+    def is_nucleic(self) -> bool:
+        t = self.polymer_type.upper()
+        return "DNA" in t or "RNA" in t
+
+
+class NonpolymerComp(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    comp_id: str  # chem_comp id, e.g. "HEM", "ZN" (CCD codes are uppercase)
+    name: str | None = None
+    formula: str | None = None
+    comp_type: str | None = None
+
+
+class QualityMetrics(BaseModel):
+    """wwPDB validation-report summary metrics (no coordinates).
+
+    Geometry metrics (clashscore, Ramachandran, rotamer) are reported for both
+    X-ray and cryo-EM; diffraction metrics (R-free, RSRZ) are X-ray only; EM
+    map-fit (backbone atom inclusion) is cryo-EM only. Any field may be ``None``
+    when the validation report does not provide it — a missing metric never
+    penalizes an entry (see :func:`ifsplit.parse.quality_drop`).
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    clashscore: float | None = None  # all-atom clashes / 1000 atoms (lower better)
+    ramachandran_outlier_pct: float | None = None  # % backbone Ramachandran outliers
+    rotamer_outlier_pct: float | None = None  # % sidechain rotamer outliers
+    rfree: float | None = None  # diffraction DCC_Rfree (X-ray only)
+    rsrz_outlier_pct: float | None = None  # diffraction % RSRZ outliers (X-ray only)
+    em_backbone_inclusion: float | None = None  # EM backbone atom inclusion (higher better)
+
+    @property
+    def has_report(self) -> bool:
+        """True if the validation report supplied at least one metric."""
+        return any(
+            v is not None
+            for v in (
+                self.clashscore,
+                self.ramachandran_outlier_pct,
+                self.rotamer_outlier_pct,
+                self.rfree,
+                self.rsrz_outlier_pct,
+                self.em_backbone_inclusion,
+            )
+        )
+
+
+class CandidateRecord(BaseModel):
+    """One PDB entry's snapshot metadata."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    entry_id: str
+    methods: list[str]
+    resolution_A: float | None
+    release_date: str  # YYYY-MM-DD
+    deposited_residues: int | None
+    assemblies: dict[str, int]  # assembly_id -> polymer_monomer_count
+    polymer_entities: list[PolymerEntity]
+    nonpolymer_comps: list[NonpolymerComp]
+    # Curation signals (Stage 4). comp ids that actually contact the protein
+    # (rcsb_entry_info.nonpolymer_bound_components) and comp ids with a measured
+    # binding affinity (rcsb_binding_affinity). Both are the buffer-vs-ligand gate.
+    bound_components: list[str] = []
+    affinity_comp_ids: list[str] = []
+    # RCSB-computed count of protein<->nucleic-acid interface entities across the
+    # entry's assemblies (rcsb_assembly_info.num_prot_na_interface_entities). > 0
+    # verifies a real protein/NA contact (Stage 4 holo gate for the nucleotide
+    # class) — distinguishing a true complex from a co-deposited oligo.
+    protein_na_interface_count: int = 0
+    # wwPDB validation-report summary (Stage 3 quality filters). All fields
+    # optional; defaults to an empty report when RCSB has none.
+    quality: QualityMetrics = QualityMetrics()
+
+    @classmethod
+    def from_data_api(cls, entry: dict) -> CandidateRecord:
+        """Build a record from a raw Data-API entry object (deterministic)."""
+        # Verbatim canonical id from RCSB (legacy or extended) — never reformat.
+        entry_id = entry["rcsb_id"]
+
+        methods = sorted(m["method"] for m in (entry.get("exptl") or []) if m.get("method"))
+
+        info = entry.get("rcsb_entry_info") or {}
+        res_list = info.get("resolution_combined") or []
+        resolution = min(res_list) if res_list else None
+        deposited = info.get("deposited_polymer_monomer_count")
+        bound = sorted({c.upper() for c in (info.get("nonpolymer_bound_components") or [])})
+
+        acc = entry.get("rcsb_accession_info") or {}
+        rel = acc.get("initial_release_date")
+        release_date = rel[:10] if rel else ""
+
+        affinity = sorted(
+            {
+                a["comp_id"].upper()
+                for a in (entry.get("rcsb_binding_affinity") or [])
+                if a.get("comp_id")
+            }
+        )
+
+        assemblies: dict[str, int] = {}
+        prot_na_interfaces = 0
+        for asm in entry.get("assemblies") or []:
+            aid = asm.get("rcsb_id")
+            info = asm.get("rcsb_assembly_info") or {}
+            count = info.get("polymer_monomer_count")
+            if aid is not None and count is not None:
+                assemblies[aid] = count
+            prot_na_interfaces += info.get("num_prot_na_interface_entities") or 0
+
+        polymers: list[PolymerEntity] = []
+        for p in entry.get("polymer_entities") or []:
+            poly = p.get("entity_poly") or {}
+            seq = poly.get("pdbx_seq_one_letter_code_can") or ""
+            seq = "".join(seq.split())  # strip newlines/whitespace the API may insert
+            cluster_ids: dict[int, int] = {}
+            for cm in p.get("rcsb_cluster_membership") or []:
+                ident = cm.get("identity")
+                cid = cm.get("cluster_id")
+                if ident is not None and cid is not None:
+                    cluster_ids[int(ident)] = int(cid)
+            polymers.append(
+                PolymerEntity(
+                    entity_id=p["rcsb_id"],  # verbatim
+                    polymer_type=poly.get("rcsb_entity_polymer_type") or "Other",
+                    seq_len=len(seq),
+                    seq=seq,
+                    cluster_ids=cluster_ids,
+                )
+            )
+        polymers.sort(key=lambda e: e.entity_id)
+
+        comps: dict[str, NonpolymerComp] = {}
+        for n in entry.get("nonpolymer_entities") or []:
+            cc = (n.get("nonpolymer_comp") or {}).get("chem_comp") or {}
+            cid = cc.get("id")
+            if cid and cid.upper() not in comps:
+                comps[cid.upper()] = NonpolymerComp(
+                    comp_id=cid.upper(),
+                    name=cc.get("name"),
+                    formula=cc.get("formula"),
+                    comp_type=cc.get("type"),
+                )
+        nonpolymers = [comps[k] for k in sorted(comps)]
+
+        # Validation-report summaries arrive as 1-element lists (or null).
+        def _first(items) -> dict:
+            return (items or [None])[0] or {}
+
+        geo = _first(entry.get("pdbx_vrpt_summary_geometry"))
+        dif = _first(entry.get("pdbx_vrpt_summary_diffraction"))
+        em = _first(entry.get("pdbx_vrpt_summary_em"))
+        quality = QualityMetrics(
+            clashscore=geo.get("clashscore"),
+            ramachandran_outlier_pct=geo.get("percent_ramachandran_outliers"),
+            rotamer_outlier_pct=geo.get("percent_rotamer_outliers"),
+            rfree=dif.get("DCC_Rfree"),
+            rsrz_outlier_pct=dif.get("percent_RSRZ_outliers"),
+            em_backbone_inclusion=em.get("atom_inclusion_backbone"),
+        )
+
+        return cls(
+            entry_id=entry_id,
+            methods=methods,
+            resolution_A=resolution,
+            release_date=release_date,
+            deposited_residues=deposited,
+            assemblies=dict(sorted(assemblies.items())),
+            polymer_entities=polymers,
+            nonpolymer_comps=nonpolymers,
+            bound_components=bound,
+            affinity_comp_ids=affinity,
+            protein_na_interface_count=prot_na_interfaces,
+            quality=quality,
+        )
+
+    def to_canonical_json(self) -> str:
+        """Single-line, sorted-key JSON for byte-stable serialization."""
+        return json.dumps(self.model_dump(mode="json"), sort_keys=True, separators=(",", ":"))
+
+
+def canonical_jsonl_bytes(records: Iterable[CandidateRecord]) -> bytes:
+    """Byte-stable candidates.jsonl content: records sorted by entry_id."""
+    ordered = sorted(records, key=lambda r: r.entry_id)
+    text = "".join(r.to_canonical_json() + "\n" for r in ordered)
+    return text.encode("utf-8")
+
+
+def sha256_hex(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def read_candidates_jsonl(path) -> list[CandidateRecord]:
+    """Parse a candidates.jsonl file back into validated records."""
+    records: list[CandidateRecord] = []
+    with open(path, encoding="utf-8") as fh:
+        for line in fh:
+            line = line.strip()
+            if line:
+                records.append(CandidateRecord.model_validate_json(line))
+    return records