if-split 0.1.0__py3-none-any.whl

ifsplit/cli.py

@@ -0,0 +1,317 @@
+"""IF-Split command-line interface.
+
+`build` runs the full pipeline: Stage 1 enumerate (RCSB Search + Data API ->
+candidates.jsonl + dataset.lock), Stage 3 filter, Stage 4 ligand classification,
+Stage 5 cluster, Stage 6 deterministic split, Stage 7 manifest + registry. No
+structure coordinates are downloaded. `verify` re-enumerates from a lock and
+reports drift; `stats` summarizes a manifest.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from pydantic import ValidationError
+
+from . import __version__
+from .config import load_config
+from .rcsb import RcsbError
+
+DEFAULT_CONFIG = "config/default.yaml"
+SPLITS_CHOICES = ("train", "val", "test")
+
+
+def cmd_build(args: argparse.Namespace) -> int:
+    from .cluster import build_clusters
+    from .enumerate import enumerate_candidates, make_console_progress
+    from .ligands import classify_components
+    from .manifest import (
+        build_lock,
+        build_manifest,
+        build_tiers_doc,
+        read_registry,
+        write_classes,
+        write_clusters,
+        write_lock,
+        write_manifest,
+        write_registry,
+        write_split_files,
+        write_tiers,
+    )
+    from .parse import drop_summary, filter_candidates
+    from .split import assign_splits, check_no_leakage
+
+    cfg = load_config(args.config)
+    sf = cfg.split_fractions
+    assembly = "biological (assembly 1)" if cfg.use_biological_assembly else "asymmetric unit"
+    print(f"IF-Split {__version__}")
+    print(f"  config file:   {args.config}")
+    print(f"  config hash:   {cfg.config_hash()}")
+    print(f"  dataset:       {cfg.dataset_version}")
+    print(f"  snapshot_date: {cfg.snapshot_date}  (selects release_date <= this)")
+    print(f"  methods:       {', '.join(cfg.experimental_methods)}")
+    print(f"  resolution:    <= {cfg.resolution_max_A} A")
+    print(f"  max residues:  < {cfg.max_total_residues + 1}")
+    print(f"  assembly:      {assembly}")
+    print(f"  identity:      {cfg.identity_threshold}")
+    print(f"  clustering:    {cfg.clustering_backend}")
+    print(f"  splits:        train={sf.train} val={sf.val} test={sf.test}  salt={cfg.split_salt!r}")
+    if args.limit is not None:
+        print(f"  limit:         {args.limit} (dev: first N by sorted entry id)")
+    print()
+
+    print("Stage 1 - enumerate candidates (Search + Data API, no coordinates):")
+    say = make_console_progress()  # timestamped + line-flushed (survives redirect)
+    records, _candidates_path, sha = enumerate_candidates(
+        cfg, args.out, limit=args.limit, progress=say
+    )
+    lock_path = write_lock(
+        build_lock(
+            cfg,
+            entry_ids=[r.entry_id for r in records],
+            candidates_sha256=sha,
+            limit=args.limit,
+        ),
+        args.out,
+    )
+    print(f"  wrote {lock_path}")
+
+    print("Stage 3 - filter (metadata only):")
+    kept, drops = filter_candidates(records, cfg)
+    dcounts = drop_summary(drops)
+    print(f"  kept {len(kept)} / {len(records)}; dropped {len(drops)} {dcounts or ''}")
+
+    print("Stage 4 - ligand classification + curation:")
+    class_map = {r.entry_id: classify_components(r, cfg) for r in kept}
+    n_artifact = sum(1 for v in class_map.values() if v["purification_artifact"])
+    print(f"  classified {len(class_map)} entries; {n_artifact} purification artifact(s) flagged")
+
+    print(f"Stage 5 - cluster ({cfg.clustering_backend} @ {cfg.identity_level}%):")
+    clusters = build_clusters(kept, cfg)
+    print(
+        f"  {clusters.n_clusters} leakage-safe components "
+        f"from {clusters.n_raw_clusters} raw clusters "
+        f"({len(clusters.multichain_entries)} multi-chain merged)"
+    )
+
+    print("Stage 6 - assign splits (deterministic hash):")
+    registry = read_registry(args.registry) if args.registry else {}
+    entry_classes = {eid: info["classes"] for eid, info in class_map.items()}
+    splits = assign_splits(clusters, cfg, registry=registry, entry_classes=entry_classes)
+    check_no_leakage(splits, clusters)  # structural guarantee; raises on violation
+    c = splits.counts
+    print(f"  train={c['train']} val={c['val']} test={c['test']}  (no cross-split leakage)")
+    if cfg.test_min_per_class:
+        if splits.minimum_shortfalls:
+            short = ", ".join(f"{k}:{v}" for k, v in splits.minimum_shortfalls.items())
+            print(f"  test minimums: applied; SHORTFALL (not enough supply) -> {short}")
+        else:
+            print("  test minimums: applied; all per-class floors met")
+
+    print("Stage 7 - manifest + registry:")
+    manifest = build_manifest(
+        cfg,
+        candidates_sha256=sha,
+        n_candidates=len(records),
+        drops=drops,
+        drop_counts=dcounts,
+        clusters=clusters,
+        splits=splits,
+        class_map=class_map,
+    )
+    mpath = write_manifest(manifest, args.out)
+    split_paths = write_split_files(splits, class_map, args.out)
+    write_clusters(clusters.entry_to_cluster, args.out)
+    write_classes(class_map, args.out)
+    write_registry(splits.cluster_split, args.out)
+    write_tiers(build_tiers_doc(class_map), args.out)
+    print(f"  wrote {mpath} (provenance + counts)")
+    for s in ("train", "val", "test"):
+        print(f"  wrote {split_paths[s]}")
+    test_class_paths = [p for k, p in split_paths.items() if k.startswith("test:")]
+    for p in sorted(test_class_paths):
+        print(f"  wrote {p}")
+    print("  wrote clusters.json, ligands.classes.json, ligands.tiers.json, splits.registry.json")
+    print()
+    print(f"Build complete: {len(kept)} structures across train/val/test.")
+    print(f"Run `if-split stats {mpath}`.")
+    return 0
+
+
+def cmd_verify(args: argparse.Namespace) -> int:
+    from .manifest import verify_lock
+
+    return verify_lock(args.lock)
+
+
+def cmd_stats(args: argparse.Namespace) -> int:
+    from .manifest import summarize_manifest
+
+    return summarize_manifest(args.manifest)
+
+
+# Confirm before a pull larger than this many structures unless --yes is given.
+_FETCH_CONFIRM_THRESHOLD = 1000
+
+
+def cmd_fetch(args: argparse.Namespace) -> int:
+    from .download import SPLITS, StructureFetcher
+    from .hydrate import hydrate, select_targets
+    from .manifest import read_manifest
+
+    if args.all and args.split:
+        print("error: use either --all or --split, not both", file=sys.stderr)
+        return 2
+    if not args.all and not args.split:
+        print(
+            "error: choose a scope: --split test (repeatable) or --all.\n"
+            "       (explicit by design — `fetch` can pull a lot of data)",
+            file=sys.stderr,
+        )
+        return 2
+
+    splits = list(SPLITS) if args.all else args.split
+    unknown = [s for s in splits if s not in SPLITS]
+    if unknown:
+        print(f"error: unknown split(s): {', '.join(unknown)}", file=sys.stderr)
+        return 2
+
+    manifest = read_manifest(args.manifest)
+    targets = select_targets(manifest, splits)
+    if not targets:
+        print(f"nothing to fetch for split(s): {', '.join(splits)}")
+        return 0
+
+    assembly = not args.asymmetric_unit
+    print(f"fetch: {len(targets)} structures across {', '.join(splits)} -> {args.out}")
+
+    # Size estimate + large-pull guard (no accidental terabyte).
+    unit_label = "assembly 1" if assembly else "asymmetric unit"
+    with StructureFetcher(assembly=assembly, workers=args.workers) as fetcher:
+        est = fetcher.estimate_bytes([e for e, _ in targets])
+        if est is not None:
+            print(f"  estimated download: ~{est / 1e9:.2f} GB ({unit_label})")
+        if len(targets) > _FETCH_CONFIRM_THRESHOLD and not args.yes:
+            print(
+                f"  refusing to fetch {len(targets)} structures without --yes "
+                f"(threshold {_FETCH_CONFIRM_THRESHOLD}).",
+                file=sys.stderr,
+            )
+            return 5
+        summary = hydrate(
+            args.manifest,
+            args.out,
+            splits=splits,
+            assembly=assembly,
+            workers=args.workers,
+            fetcher=fetcher,
+            progress=lambda m: print(f"  {m}"),
+        )
+
+    print(
+        f"done: {summary['fetched']} fetched, {summary['skipped']} cached, "
+        f"{len(summary['failed'])} failed"
+    )
+    if summary["failed"]:
+        for eid, reason in summary["failed"][:10]:
+            print(f"  ! {eid}: {reason}", file=sys.stderr)
+    for kind, path in summary["index"].items():
+        print(f"  index ({kind}): {path}")
+    print(f"  dataset card: {args.out}/DATASET_CARD.md")
+    return 0 if not summary["failed"] else 6
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="if-split",
+        description="Reproducible, ligand-aware PDB train/val/test splitter (LigandMPNN-style).",
+    )
+    p.add_argument("--version", action="version", version=f"if-split {__version__}")
+    sub = p.add_subparsers(dest="command", required=True)
+
+    pb = sub.add_parser("build", help="Run the pipeline (Stages 1-7) and emit manifest + lock.")
+    pb.add_argument(
+        "--config",
+        default=DEFAULT_CONFIG,
+        help=f"Path to config YAML (default: {DEFAULT_CONFIG}).",
+    )
+    pb.add_argument(
+        "--out",
+        default="data/out",
+        help="Output dir for candidates.jsonl + dataset.lock (default: data/out).",
+    )
+    pb.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="Dev only: cap to the first N candidates by sorted entry id (reproducible).",
+    )
+    pb.add_argument(
+        "--registry",
+        default=None,
+        help="Optional prior splits.registry.json to pin existing cluster->split "
+        "assignments (growth stability).",
+    )
+    pb.set_defaults(func=cmd_build)
+
+    pv = sub.add_parser(
+        "verify", help="Re-enumerate from a lock file and report drift vs the live PDB."
+    )
+    pv.add_argument("lock", help="Path to dataset.lock")
+    pv.set_defaults(func=cmd_verify)
+
+    ps = sub.add_parser(
+        "stats", help="Report split sizes and per-class test counts from a manifest."
+    )
+    ps.add_argument("manifest", help="Path to manifest.json")
+    ps.set_defaults(func=cmd_stats)
+
+    pf = sub.add_parser(
+        "fetch",
+        help="OPTIONAL: download structures for a built manifest into an ML-ready tree.",
+    )
+    pf.add_argument("manifest", help="Path to manifest.json")
+    pf.add_argument(
+        "--out", default="data/structures", help="Output root (default: data/structures)."
+    )
+    pf.add_argument(
+        "--split",
+        action="append",
+        choices=list(SPLITS_CHOICES),
+        help="Split to fetch (repeatable, e.g. --split test --split val).",
+    )
+    pf.add_argument("--all", action="store_true", help="Fetch all splits.")
+    pf.add_argument(
+        "--asymmetric-unit",
+        action="store_true",
+        help="Fetch the asymmetric unit instead of biological assembly 1.",
+    )
+    pf.add_argument("--workers", type=int, default=8, help="Concurrent downloads (default: 8).")
+    pf.add_argument("--yes", action="store_true", help="Proceed without confirming large pulls.")
+    pf.set_defaults(func=cmd_fetch)
+
+    return p
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        return args.func(args)
+    except FileNotFoundError as e:
+        print(f"error: {e}", file=sys.stderr)
+        return 2
+    except ValidationError as e:
+        print(f"invalid config:\n{e}", file=sys.stderr)
+        return 2
+    except NotImplementedError as e:
+        print(f"not implemented yet: {e}", file=sys.stderr)
+        return 3
+    except RcsbError as e:
+        print(f"RCSB request failed: {e}", file=sys.stderr)
+        return 4
+
+
+if __name__ == "__main__":
+    sys.exit(main())

ifsplit/cluster.py

@@ -0,0 +1,130 @@
+"""Stage 5 - Cluster protein entities into leakage-safe groups.
+
+The ``precomputed`` backend reads each protein entity's RCSB cluster id at the
+configured identity level from ``PolymerEntity.cluster_ids`` (captured in Stage 1
+from the Data API ``rcsb_cluster_membership`` field) - no file download, no
+mmseqs2 binary.
+
+A *raw cluster* is the set of protein entities sharing an RCSB cluster id. But an
+entry with several protein chains can touch several raw clusters, so raw clusters
+alone are NOT a leakage-safe split unit: if entry X has chain a (raw cluster A)
+and chain b (raw cluster B), then A and B must land in the same split or X's b
+sequence leaks across splits. We therefore merge raw clusters joined by a shared
+entry into **components** (connected components, union-find). The component is the
+unit Stage 6 assigns to a split, which makes cross-split sequence overlap
+impossible by construction - no heuristic, no after-the-fact audit.
+
+A component's canonical key is the lexicographically smallest entity id across all
+its members. (Equivalently the smallest raw-cluster key, since each raw key is
+itself a min-entity-id - so min-of-mins = global min.) Keying the split hash on a
+stable member id, not RCSB's volatile integer cluster id, keeps assignments
+stable as the dataset grows (PLAN.md §6). Sub-10-aa peptides that RCSB does not
+cluster become their own singleton components.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from .config import Config
+from .schema import CandidateRecord
+
+SINGLETON_PREFIX = "singleton:"
+
+
+@dataclass
+class ClusterResult:
+    """Stage 5 output: raw sequence clusters merged into leakage-safe components."""
+
+    identity: int
+    entry_to_cluster: dict[str, str]  # entry_id -> component key (the split unit)
+    cluster_members: dict[str, list[str]]  # component key -> sorted entry_ids
+    entry_raw_clusters: dict[str, list[str]]  # entry_id -> raw cluster keys it touches
+    multichain_entries: list[str] = field(default_factory=list)
+    unclustered_entries: list[str] = field(default_factory=list)
+    n_raw_clusters: int = 0
+
+    @property
+    def n_clusters(self) -> int:
+        """Number of components (the split units)."""
+        return len(self.cluster_members)
+
+
+def build_clusters(records: list[CandidateRecord], cfg: Config) -> ClusterResult:
+    """Cluster filtered records at ``cfg.identity_level``, merged into components."""
+    if cfg.clustering_backend != "precomputed":
+        raise NotImplementedError(
+            f"clustering_backend {cfg.clustering_backend!r} not implemented "
+            "(only 'precomputed' is available)."
+        )
+    level = cfg.identity_level
+
+    # 1. Raw clusters: RCSB cluster id -> member entity ids -> canonical raw key
+    #    (the smallest member entity id).
+    raw_entities: dict[int, set[str]] = {}
+    for r in records:
+        for e in r.polymer_entities:
+            if e.is_protein and level in e.cluster_ids:
+                raw_entities.setdefault(e.cluster_ids[level], set()).add(e.entity_id)
+    raw_key = {cid: min(ents) for cid, ents in raw_entities.items()}
+
+    # 2. Each entry -> the raw cluster keys it touches (a singleton key if no
+    #    protein chain is clustered at this level).
+    entry_raw: dict[str, list[str]] = {}
+    multichain: list[str] = []
+    unclustered: list[str] = []
+    all_keys: set[str] = set(raw_key.values())
+    for r in records:
+        proteins = [e for e in r.polymer_entities if e.is_protein]
+        if not proteins:
+            continue  # defensive; Stage 3 already drops no-protein entries
+        keys = sorted({raw_key[e.cluster_ids[level]] for e in proteins if level in e.cluster_ids})
+        if not keys:
+            singleton = SINGLETON_PREFIX + min(e.entity_id for e in proteins)
+            keys = [singleton]
+            all_keys.add(singleton)
+            unclustered.append(r.entry_id)
+        elif len(keys) > 1:
+            multichain.append(r.entry_id)
+        entry_raw[r.entry_id] = keys
+
+    # 3. Union-find: merge raw clusters joined by a shared entry into components.
+    #    The smaller key is always made the root, so a component's root is its
+    #    global-minimum key (order-independent -> deterministic).
+    parent = {k: k for k in all_keys}
+
+    def find(x: str) -> str:
+        root = x
+        while parent[root] != root:
+            root = parent[root]
+        while parent[x] != root:  # path compression
+            parent[x], x = root, parent[x]
+        return root
+
+    def union(a: str, b: str) -> None:
+        ra, rb = find(a), find(b)
+        if ra != rb:
+            lo, hi = (ra, rb) if ra < rb else (rb, ra)
+            parent[hi] = lo
+
+    for keys in entry_raw.values():
+        for k in keys[1:]:
+            union(keys[0], k)
+
+    # 4. Materialize components: component key -> entries; entry -> component.
+    entry_to_cluster: dict[str, str] = {}
+    members: dict[str, set[str]] = {}
+    for entry, keys in entry_raw.items():
+        comp = find(keys[0])
+        entry_to_cluster[entry] = comp
+        members.setdefault(comp, set()).add(entry)
+
+    return ClusterResult(
+        identity=level,
+        entry_to_cluster=dict(sorted(entry_to_cluster.items())),
+        cluster_members={k: sorted(v) for k, v in sorted(members.items())},
+        entry_raw_clusters=dict(sorted(entry_raw.items())),
+        multichain_entries=sorted(multichain),
+        unclustered_entries=sorted(unclustered),
+        n_raw_clusters=len(raw_key),
+    )

ifsplit/config.py

@@ -0,0 +1,146 @@
+"""Load, validate, and hash an IF-Split run configuration.
+
+The config is the single source of truth for a build. Its canonical hash is
+embedded in the manifest so that two manifests sharing a config hash are
+guaranteed to have used identical, output-affecting settings.
+
+The hash is computed over the *validated, normalized* settings (not the raw YAML
+text), so comments and formatting do not change it — only values do.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from datetime import date
+from pathlib import Path
+from typing import Any, Literal
+
+import yaml
+from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+
+
+class SplitFractions(BaseModel):
+    """Train/val/test partition fractions; must sum to 1.0."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    train: float = Field(gt=0, lt=1)
+    val: float = Field(gt=0, lt=1)
+    test: float = Field(gt=0, lt=1)
+
+    @model_validator(mode="after")
+    def _sum_to_one(self) -> SplitFractions:
+        total = self.train + self.val + self.test
+        if abs(total - 1.0) > 1e-9:
+            raise ValueError(f"split_fractions must sum to 1.0, got {total}")
+        return self
+
+
+class Config(BaseModel):
+    """A fully-validated IF-Split build configuration."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    # --- snapshot definition (reproducibility anchor) ---
+    snapshot_date: date
+    experimental_methods: list[str] = Field(min_length=1)
+    resolution_max_A: float = Field(gt=0)
+    max_total_residues: int = Field(gt=0)
+    excluded_het: list[str] = Field(default_factory=list)
+    use_biological_assembly: bool = True
+
+    # --- curation: purification-artifact detection (Stage 4) ---
+    # A poly-His tag coordinating Ni/Co is a purification artifact, not a
+    # biological metal site (a known blemish in the LigandMPNN metal set). An
+    # entry whose *only* metal is a purification metal AND that carries a His-tag
+    # is flagged so it can be excluded from the metal class. Empty
+    # purification_metals disables the heuristic.
+    purification_metals: list[str] = Field(default_factory=lambda: ["NI", "CO"])
+    histag_min_run: int = Field(default=6, gt=0)
+    exclude_purification_artifacts: bool = True
+
+    # --- clustering + split ---
+    identity_threshold: float = Field(gt=0, le=1)
+    # "precomputed": reuse RCSB's entity clusters (default, no external binary).
+    # "mmseqs2": run our own over the snapshot's sequences.
+    clustering_backend: Literal["precomputed", "mmseqs2"] = "precomputed"
+    split_fractions: SplitFractions
+    split_salt: str = Field(min_length=1)
+    seed: int = Field(ge=0)
+
+    # --- test-set minimums (opt-in stratification top-up) ---
+    # Floor on the number of test entries carrying each functional ligand class,
+    # e.g. {"metal": 500, "nucleotide": 200}. Empty (default) = pure deterministic
+    # hash, no top-up. When set, after the base assignment any class below its floor
+    # recruits *whole components* (never individual entries, so no leakage) into
+    # test in deterministic hash order, skipping components already pinned by a
+    # registry (so growth stays stable). A floor larger than the available supply
+    # is satisfied as far as possible and the shortfall is reported, not forced.
+    test_min_per_class: dict[str, int] = Field(default_factory=dict)
+
+    # --- quality filters (Stage 3): wwPDB validation-report metrics ---
+    # Metadata only (no coordinates). Each cap is optional (None disables it). An
+    # entry is dropped only when the metric is present AND violates the cap;
+    # entries whose report lacks a metric are kept (never penalized for an absent
+    # metric). Geometry caps (clashscore/Ramachandran/rotamer) apply to X-ray and
+    # cryo-EM; diffraction caps (R-free/RSRZ) naturally no-op on EM entries.
+    max_clashscore: float | None = Field(default=None, gt=0)
+    max_rfree: float | None = Field(default=None, gt=0)
+    max_ramachandran_outlier_pct: float | None = Field(default=None, ge=0)
+    max_rotamer_outlier_pct: float | None = Field(default=None, ge=0)
+    max_rsrz_outlier_pct: float | None = Field(default=None, ge=0)
+    require_validation_report: bool = False
+
+    # --- featurization (downstream-optional; not part of the split definition) ---
+    ligand_context_radius_A: float = Field(gt=0)
+    max_ligand_atoms: int = Field(gt=0)
+
+    @field_validator("experimental_methods")
+    @classmethod
+    def _normalize_methods(cls, v: list[str]) -> list[str]:
+        return [m.strip().upper() for m in v]
+
+    @field_validator("excluded_het", "purification_metals")
+    @classmethod
+    def _normalize_codes(cls, v: list[str]) -> list[str]:
+        return [h.strip().upper() for h in v]
+
+    @field_validator("test_min_per_class")
+    @classmethod
+    def _check_minimums(cls, v: dict[str, int]) -> dict[str, int]:
+        for k, n in v.items():
+            if n < 0:
+                raise ValueError(f"test_min_per_class[{k!r}] must be >= 0, got {n}")
+        return v
+
+    @property
+    def dataset_version(self) -> str:
+        """Versioned dataset name, e.g. 'IF-Split-2026.05.30'."""
+        return f"IF-Split-{self.snapshot_date:%Y.%m.%d}"
+
+    @property
+    def identity_level(self) -> int:
+        """``identity_threshold`` as an integer percent (e.g. 0.30 -> 30)."""
+        return round(self.identity_threshold * 100)
+
+    def canonical_dict(self) -> dict[str, Any]:
+        """JSON-mode dump (dates -> ISO strings) used for hashing and manifests."""
+        return self.model_dump(mode="json")
+
+    def config_hash(self) -> str:
+        """Deterministic, formatting-independent hash of the settings."""
+        canonical = json.dumps(self.canonical_dict(), sort_keys=True, separators=(",", ":"))
+        return hashlib.blake2b(canonical.encode("utf-8"), digest_size=16).hexdigest()
+
+
+def load_config(path: str | Path) -> Config:
+    """Load and validate a YAML config file into a :class:`Config`."""
+    path = Path(path)
+    if not path.exists():
+        raise FileNotFoundError(f"Config not found: {path}")
+    with path.open("r", encoding="utf-8") as fh:
+        raw = yaml.safe_load(fh)
+    if not isinstance(raw, dict):
+        raise ValueError(f"Config must be a YAML mapping, got {type(raw).__name__}: {path}")
+    return Config.model_validate(raw)