protcross 0.1.1__py3-none-any.whl

evopoint_da/__init__.py

@@ -0,0 +1,4 @@
+"""ProtCross core package."""
+
+__version__ = "0.1.1"
+

evopoint_da/assets.py

@@ -0,0 +1,178 @@
+"""Download and validate ProtCross runtime assets."""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import shlex
+from dataclasses import dataclass
+from pathlib import Path
+
+import requests
+from tqdm import tqdm
+
+
+RELEASE_TAG = "v0.1.1"
+GITHUB_RELEASE_BASE = f"https://github.com/GeraltZeroZhong/ProtCross/releases/download/{RELEASE_TAG}"
+DEFAULT_ASSETS_DIR = Path.home() / ".cache" / "protcross" / "assets" / RELEASE_TAG
+
+DEFAULT_ESM_URL = (
+    "https://huggingface.co/EvolutionaryScale/esmc-600m-2024-12/"
+    "resolve/main/data/weights/esmc_600m_2024_12_v0.pth"
+)
+DEFAULT_CHECKPOINT_URL = f"{GITHUB_RELEASE_BASE}/best-epoch.59.ckpt"
+DEFAULT_PCA_URL = f"{GITHUB_RELEASE_BASE}/pca_esmc_128.pkl"
+
+
+@dataclass(frozen=True)
+class AssetSpec:
+    name: str
+    filename: str
+    url: str
+    sha256: str | None = None
+
+
+DEFAULT_ASSETS = (
+    AssetSpec(
+        name="ESM-C 600M weights",
+        filename="esmc_600m_2024_12_v0.pth",
+        url=DEFAULT_ESM_URL,
+        sha256="8ef856e1a237ee3f995442df997a962e70057faadecf38fc0c8561bd3c2f4324",
+    ),
+    AssetSpec(
+        name="ProtCross checkpoint",
+        filename="best-epoch=59.ckpt",
+        url=DEFAULT_CHECKPOINT_URL,
+        sha256="3eb6d8c9ef94541efc0444508e15d630c156a98e164a6caa08f2ae7a20371e45",
+    ),
+    AssetSpec(
+        name="ProtCross PCA reducer",
+        filename="pca_esmc_128.pkl",
+        url=DEFAULT_PCA_URL,
+        sha256="c4317684fb94c1337a44b844381d7e84472a6958b34a604b0f982984b629098b",
+    ),
+)
+
+
+def get_default_assets_dir() -> Path:
+    return Path(os.environ.get("PROTCROSS_ASSETS_DIR", DEFAULT_ASSETS_DIR)).expanduser()
+
+
+def setup_assets(
+    output_dir: str | Path | None = None,
+    *,
+    esm_url: str = DEFAULT_ESM_URL,
+    checkpoint_url: str = DEFAULT_CHECKPOINT_URL,
+    pca_url: str = DEFAULT_PCA_URL,
+    force: bool = False,
+    verify: bool = True,
+    skip_esm: bool = False,
+) -> Path:
+    """Download ProtCross assets and return the asset directory."""
+    output_dir = Path(output_dir).expanduser() if output_dir else get_default_assets_dir()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    specs = [
+        AssetSpec(
+            DEFAULT_ASSETS[0].name,
+            DEFAULT_ASSETS[0].filename,
+            esm_url,
+            DEFAULT_ASSETS[0].sha256,
+        ),
+        AssetSpec(
+            DEFAULT_ASSETS[1].name,
+            DEFAULT_ASSETS[1].filename,
+            checkpoint_url,
+            DEFAULT_ASSETS[1].sha256,
+        ),
+        AssetSpec(
+            DEFAULT_ASSETS[2].name,
+            DEFAULT_ASSETS[2].filename,
+            pca_url,
+            DEFAULT_ASSETS[2].sha256,
+        ),
+    ]
+    if skip_esm:
+        specs = specs[1:]
+
+    print(f"Installing ProtCross assets into {output_dir}")
+    print("Note: ESM-C weights are distributed by EvolutionaryScale under their Hugging Face model terms.")
+    for spec in specs:
+        download_asset(spec, output_dir / spec.filename, force=force, verify=verify)
+
+    write_env_file(output_dir, include_esm=(output_dir / DEFAULT_ASSETS[0].filename).exists())
+    print("\nAsset setup complete.")
+    print("Use with: protcross predict input.pdb --output output.pdb")
+    print(f"Environment file written to: {output_dir / 'protcross.env'}")
+    return output_dir
+
+
+def download_asset(spec: AssetSpec, output_path: Path, *, force: bool = False, verify: bool = True) -> None:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    if output_path.exists() and not force:
+        if not verify or not spec.sha256 or sha256_file(output_path) == spec.sha256:
+            print(f"[skip] {spec.name}: {output_path}")
+            return
+        print(f"[warn] Existing file failed SHA256 verification and will be replaced: {output_path}")
+
+    tmp_path = output_path.with_suffix(output_path.suffix + ".part")
+    if tmp_path.exists():
+        tmp_path.unlink()
+
+    print(f"[download] {spec.name}")
+    print(f"           {spec.url}")
+    with requests.get(spec.url, stream=True, timeout=30) as response:
+        response.raise_for_status()
+        total = int(response.headers.get("content-length", 0))
+        with tmp_path.open("wb") as file, tqdm(
+            total=total if total > 0 else None,
+            unit="B",
+            unit_scale=True,
+            unit_divisor=1024,
+            desc=spec.filename,
+        ) as progress:
+            for chunk in response.iter_content(chunk_size=1024 * 1024):
+                if not chunk:
+                    continue
+                file.write(chunk)
+                progress.update(len(chunk))
+
+    if verify and spec.sha256:
+        actual = sha256_file(tmp_path)
+        if actual != spec.sha256:
+            tmp_path.unlink(missing_ok=True)
+            raise RuntimeError(
+                f"SHA256 mismatch for {spec.filename}: expected {spec.sha256}, got {actual}"
+            )
+
+    tmp_path.replace(output_path)
+    print(f"[ok] {output_path}")
+
+
+def sha256_file(path: str | Path) -> str:
+    digest = hashlib.sha256()
+    with Path(path).open("rb") as file:
+        for chunk in iter(lambda: file.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def write_env_file(output_dir: Path, *, include_esm: bool = True) -> None:
+    env_path = output_dir / "protcross.env"
+    lines = [
+        _export_line("PROTCROSS_ASSETS_DIR", output_dir),
+        _export_line("PROTCROSS_CHECKPOINT", output_dir / "best-epoch=59.ckpt"),
+    ]
+    if include_esm:
+        lines.append(_export_line("PROTCROSS_ESM_WEIGHTS", output_dir / "esmc_600m_2024_12_v0.pth"))
+    lines.extend(
+        [
+            _export_line("PROTCROSS_PCA", output_dir / "pca_esmc_128.pkl"),
+            "",
+        ]
+    )
+    env_path.write_text("\n".join(lines), encoding="utf-8")
+
+
+def _export_line(name: str, value: str | Path) -> str:
+    return f"export {name}={shlex.quote(str(value))}"

evopoint_da/cli/__init__.py

@@ -0,0 +1,2 @@
+"""Command line entry points."""
+

evopoint_da/cli/download_af2.py

@@ -0,0 +1,37 @@
+"""CLI for downloading matching AlphaFold structures."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from evopoint_da.data.af2 import AF2DownloadConfig, download_af2_structures
+
+
+def build_parser(prog: str | None = None) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog=prog, description="Download AlphaFold structures for PDB files.")
+    parser.add_argument("--raw-pdb-dir", default="data/raw_pdb")
+    parser.add_argument("--output-dir", default="data/raw_af2")
+    parser.add_argument("--mapping-file", default="pdb_uniprot_mapping.json")
+    parser.add_argument("--max-workers", type=int, default=8)
+    parser.add_argument("--uniprot-candidates", type=int, default=3)
+    parser.add_argument("--timeout-seconds", type=int, default=30)
+    return parser
+
+
+def main(argv: list[str] | None = None, *, prog: str | None = None) -> int:
+    args = build_parser(prog=prog).parse_args(argv)
+    config = AF2DownloadConfig(
+        raw_pdb_dir=Path(args.raw_pdb_dir),
+        output_dir=Path(args.output_dir),
+        mapping_file=Path(args.mapping_file),
+        max_workers=args.max_workers,
+        uniprot_candidates=args.uniprot_candidates,
+        timeout_seconds=args.timeout_seconds,
+    )
+    download_af2_structures(config)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

evopoint_da/cli/main.py

@@ -0,0 +1,42 @@
+"""Unified ProtCross command line interface."""
+
+from __future__ import annotations
+
+import argparse
+from importlib import import_module
+
+from evopoint_da import __version__
+
+
+COMMANDS = {
+    "predict": "evopoint_da.cli.predict:main",
+    "setup-assets": "evopoint_da.cli.setup_assets:main",
+    "preprocess": "evopoint_da.cli.preprocess:main",
+    "download-af2": "evopoint_da.cli.download_af2:main",
+    "map-labels": "evopoint_da.cli.map_labels:main",
+}
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="protcross", description="ProtCross command line tools.")
+    parser.add_argument("--version", action="version", version=f"ProtCross {__version__}")
+    parser.add_argument("command", choices=sorted(COMMANDS))
+    parser.add_argument("args", nargs=argparse.REMAINDER)
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+    command_args = args.args
+    if command_args and command_args[0] == "--":
+        command_args = command_args[1:]
+    return _load_command(COMMANDS[args.command])(command_args, prog=f"protcross {args.command}")
+
+
+def _load_command(target: str):
+    module_name, function_name = target.split(":", 1)
+    return getattr(import_module(module_name), function_name)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

evopoint_da/cli/map_labels.py

@@ -0,0 +1,43 @@
+"""CLI for mapping PDB binding labels onto AF2 structures."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from evopoint_da.data.label_mapping import LabelMappingConfig, map_labels
+
+
+def build_parser(prog: str | None = None) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog=prog, description="Map PDB-derived labels onto processed AF2 samples.")
+    parser.add_argument("--processed-pdb-dir", default="data/processed_pdb")
+    parser.add_argument("--processed-af2-dir", default="data/processed_af2")
+    parser.add_argument("--raw-pdb-dir", default="data/raw_pdb")
+    parser.add_argument("--raw-af2-dir", default="data/raw_af2")
+    parser.add_argument("--mapping-file", default="pdb_uniprot_mapping.json")
+    parser.add_argument("--output-csv", default="mapping_report_final.csv")
+    parser.add_argument("--debug-limit", type=int, default=5)
+    parser.add_argument("--min-chain-score", type=float, default=0.15)
+    parser.add_argument("--max-rmsd", type=float, default=30.0)
+    return parser
+
+
+def main(argv: list[str] | None = None, *, prog: str | None = None) -> int:
+    args = build_parser(prog=prog).parse_args(argv)
+    config = LabelMappingConfig(
+        processed_pdb_dir=Path(args.processed_pdb_dir),
+        processed_af2_dir=Path(args.processed_af2_dir),
+        raw_pdb_dir=Path(args.raw_pdb_dir),
+        raw_af2_dir=Path(args.raw_af2_dir),
+        mapping_file=Path(args.mapping_file),
+        output_csv=Path(args.output_csv),
+        debug_limit=args.debug_limit,
+        min_chain_score=args.min_chain_score,
+        max_rmsd=args.max_rmsd,
+    )
+    map_labels(config)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

evopoint_da/cli/predict.py

@@ -0,0 +1,156 @@
+"""Command line interface for lightweight ProtCross prediction."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+
+LOCAL_CHECKPOINT = Path("checkpoint/best-epoch=59.ckpt")
+LOCAL_PCA = Path("data/pca_esmc_128.pkl")
+
+
+def build_parser(prog: str | None = None) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog=prog,
+        description="Predict binding-site probabilities for one PDB/mmCIF structure and write them to B-factors.",
+    )
+    parser.add_argument("input_pdb", nargs="?", help="Input PDB/mmCIF structure.")
+    parser.add_argument("--pdb_file", "--pdb-file", dest="pdb_file", help="Legacy input structure argument.")
+    parser.add_argument(
+        "--assets-dir",
+        help=(
+            "Directory containing best-epoch=59.ckpt, esmc_600m_2024_12_v0.pth, "
+            "and pca_esmc_128.pkl. Explicit file arguments override this."
+        ),
+    )
+    parser.add_argument(
+        "--ckpt_path",
+        "--ckpt-path",
+        "--checkpoint",
+        dest="ckpt_path",
+        default=os.environ.get("PROTCROSS_CHECKPOINT"),
+        help=(
+            "ProtCross Lightning checkpoint. Defaults to PROTCROSS_CHECKPOINT, "
+            "installed assets, or checkpoint/best-epoch=59.ckpt when present."
+        ),
+    )
+    parser.add_argument(
+        "--esm_weights",
+        "--esm-weights",
+        dest="esm_weights",
+        default=os.environ.get("PROTCROSS_ESM_WEIGHTS"),
+        help="Local ESM-C 600M weights. Can also be set with PROTCROSS_ESM_WEIGHTS.",
+    )
+    parser.add_argument(
+        "--pca_path",
+        "--pca-path",
+        "--pca",
+        dest="pca_path",
+        default=os.environ.get("PROTCROSS_PCA"),
+        help=(
+            "Fitted PCA pickle for ESM-C embeddings. Defaults to PROTCROSS_PCA, "
+            "installed assets, or data/pca_esmc_128.pkl when present."
+        ),
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        "--output_pdb",
+        "--output-pdb",
+        dest="output_pdb",
+        help="Output PDB path. Predicted probabilities overwrite B-factors.",
+    )
+    parser.add_argument("--scores-tsv", help="Optional residue-level score table.")
+    parser.add_argument("--threshold", type=float, default=0.5, help="Threshold used in the text summary.")
+    parser.add_argument("--chain", "--chain-id", dest="chain_id", help="Restrict prediction to one chain.")
+    parser.add_argument("--device", default="auto", help="auto, cpu, cuda or cuda:N.")
+    parser.add_argument("--pca_dim", "--pca-dim", dest="pca_dim", type=int, default=128)
+    parser.add_argument("--max-len", type=int, default=1022, help="Maximum residues passed to ESM-C.")
+    parser.add_argument("--fail-on-truncation", action="store_true")
+    parser.add_argument("--quiet", action="store_true")
+    return parser
+
+
+def main(argv: list[str] | None = None, *, prog: str | None = None) -> int:
+    parser = build_parser(prog=prog)
+    args = parser.parse_args(argv)
+
+    input_pdb = args.input_pdb or args.pdb_file
+    if not input_pdb:
+        parser.error("an input PDB/mmCIF path is required")
+
+    assets = _resolve_asset_directory(args.assets_dir)
+    if assets:
+        args.ckpt_path = args.ckpt_path or str(assets.checkpoint)
+        args.esm_weights = args.esm_weights or str(assets.esm_weights)
+        args.pca_path = args.pca_path or str(assets.pca)
+
+    if not args.ckpt_path and LOCAL_CHECKPOINT.exists():
+        args.ckpt_path = str(LOCAL_CHECKPOINT)
+    if not args.pca_path and LOCAL_PCA.exists():
+        args.pca_path = str(LOCAL_PCA)
+
+    if not args.esm_weights:
+        parser.error(
+            "--esm-weights is required unless --assets-dir, PROTCROSS_ESM_WEIGHTS, "
+            "or installed assets from `protcross setup-assets` are available"
+        )
+    if not args.ckpt_path:
+        parser.error("--checkpoint is required unless PROTCROSS_CHECKPOINT or installed assets are available")
+    if not args.pca_path:
+        parser.error("--pca is required unless PROTCROSS_PCA or installed assets are available")
+
+    try:
+        from evopoint_da.inference import ProtCrossPredictor
+
+        predictor = ProtCrossPredictor.from_files(
+            ckpt_path=args.ckpt_path,
+            esm_weights=args.esm_weights,
+            pca_path=args.pca_path,
+            device=args.device,
+            pca_dim=args.pca_dim,
+            max_len=args.max_len,
+        )
+        result = predictor.predict(
+            input_pdb,
+            chain_id=args.chain_id,
+            threshold=args.threshold,
+            output_pdb=args.output_pdb,
+            scores_tsv=args.scores_tsv,
+        )
+        if args.fail_on_truncation and result.truncated:
+            raise RuntimeError(
+                f"Input was truncated from {result.original_length} to {len(result.residue_ids)} residues."
+            )
+    except Exception as exc:
+        print(f"ProtCross prediction failed: {exc}", file=sys.stderr)
+        return 1
+
+    if not args.quiet:
+        print(result.format_summary())
+        if args.output_pdb:
+            print(f"Wrote B-factor prediction PDB: {Path(args.output_pdb)}")
+        if args.scores_tsv:
+            print(f"Wrote score table: {Path(args.scores_tsv)}")
+
+    return 0
+
+
+def _resolve_asset_directory(assets_dir: str | None) -> PredictorAssets | None:
+    from evopoint_da.inference import PredictorAssets
+
+    if assets_dir:
+        return PredictorAssets.from_dir(assets_dir)
+
+    default_assets = PredictorAssets.from_default_dir()
+    if default_assets.is_complete():
+        return default_assets
+
+    return None
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

evopoint_da/cli/preprocess.py

@@ -0,0 +1,45 @@
+"""CLI for ESM-C + PCA preprocessing."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from evopoint_da.data.preprocess import PreprocessConfig, preprocess_directory
+
+
+def build_parser(prog: str | None = None) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog=prog, description="Preprocess PDB/mmCIF structures into ProtCross .pt files.")
+    parser.add_argument("--data_dir", "--data-dir", dest="data_dir", required=True)
+    parser.add_argument("--output_dir", "--output-dir", dest="output_dir", required=True)
+    parser.add_argument("--model_name", "--model-name", dest="model_name", required=True)
+    parser.add_argument("--pca_model_path", "--pca-model-path", dest="pca_model_path", default="pca_esmc_128.pkl")
+    parser.add_argument("--fit_pca", "--fit-pca", dest="fit_pca", action="store_true")
+    parser.add_argument("--pca_dim", "--pca-dim", dest="pca_dim", type=int, default=128)
+    parser.add_argument("--is_af2", "--is-af2", dest="is_af2", action="store_true")
+    parser.add_argument("--sample_ratio", "--sample-ratio", dest="sample_ratio", type=float, default=0.1)
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--max-len", type=int, default=1022)
+    return parser
+
+
+def main(argv: list[str] | None = None, *, prog: str | None = None) -> int:
+    args = build_parser(prog=prog).parse_args(argv)
+    config = PreprocessConfig(
+        data_dir=Path(args.data_dir),
+        output_dir=Path(args.output_dir),
+        model_name=Path(args.model_name),
+        pca_model_path=Path(args.pca_model_path),
+        fit_pca=args.fit_pca,
+        pca_dim=args.pca_dim,
+        is_af2=args.is_af2,
+        sample_ratio=args.sample_ratio,
+        device=args.device,
+        max_len=args.max_len,
+    )
+    preprocess_directory(config)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

evopoint_da/cli/setup_assets.py

@@ -0,0 +1,51 @@
+"""CLI for downloading ProtCross runtime assets."""
+
+from __future__ import annotations
+
+import argparse
+
+from evopoint_da.assets import (
+    DEFAULT_CHECKPOINT_URL,
+    DEFAULT_ESM_URL,
+    DEFAULT_PCA_URL,
+    setup_assets,
+)
+
+
+def build_parser(prog: str | None = None) -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog=prog,
+        description="Download ProtCross checkpoint, PCA reducer, and ESM-C weights.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        help=(
+            "Directory where assets will be installed. Defaults to PROTCROSS_ASSETS_DIR "
+            "or ~/.cache/protcross/assets/v0.1.1."
+        ),
+    )
+    parser.add_argument("--esm-url", default=DEFAULT_ESM_URL)
+    parser.add_argument("--checkpoint-url", default=DEFAULT_CHECKPOINT_URL)
+    parser.add_argument("--pca-url", default=DEFAULT_PCA_URL)
+    parser.add_argument("--force", action="store_true", help="Re-download files even if they already exist.")
+    parser.add_argument("--no-verify", action="store_true", help="Skip SHA256 verification.")
+    parser.add_argument("--skip-esm", action="store_true", help="Only download ProtCross checkpoint and PCA reducer.")
+    return parser
+
+
+def main(argv: list[str] | None = None, *, prog: str | None = None) -> int:
+    args = build_parser(prog=prog).parse_args(argv)
+    setup_assets(
+        args.output_dir,
+        esm_url=args.esm_url,
+        checkpoint_url=args.checkpoint_url,
+        pca_url=args.pca_url,
+        force=args.force,
+        verify=not args.no_verify,
+        skip_esm=args.skip_esm,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

evopoint_da/cli/train.py

@@ -0,0 +1,19 @@
+"""Hydra CLI entry point for standard ProtCross training."""
+
+from __future__ import annotations
+
+import hydra
+from omegaconf import DictConfig
+from pathlib import Path
+
+from evopoint_da.training import run_training
+
+CONFIG_PATH = str(Path(__file__).resolve().parents[3] / "configs")
+
+@hydra.main(version_base="1.3", config_path=CONFIG_PATH, config_name="train")
+def main(cfg: DictConfig) -> None:
+    run_training(cfg)
+
+
+if __name__ == "__main__":
+    main()

evopoint_da/data/__init__.py

@@ -0,0 +1,19 @@
+"""Data parsing and loading utilities."""
+
+from .af2 import AF2DownloadConfig, AF2Downloader, download_af2_structures
+from .label_mapping import LabelMappingConfig, map_labels
+from .pca import PCAReducer
+from .structure import MAX_ESM_RESIDUES, STANDARD_AA, StructureParser, truncate_parsed_structure
+
+__all__ = [
+    "AF2DownloadConfig",
+    "AF2Downloader",
+    "LabelMappingConfig",
+    "MAX_ESM_RESIDUES",
+    "PCAReducer",
+    "STANDARD_AA",
+    "StructureParser",
+    "download_af2_structures",
+    "map_labels",
+    "truncate_parsed_structure",
+]