#!/usr/bin/env python3 """Run Project Chimera on development-only protein-cluster folds.""" from __future__ import annotations import argparse import importlib.metadata import json import time from pathlib import Path from typing import Any import numpy as np from mitointeract_recovery.chimera import ( LightGBMConfig, ResidualVariant, evaluate_chimera, ligand_feature_matrix, read_manifest, read_mmseqs_clusters, sha256_file, validate_and_align, ) def read_jsonl(path: Path) -> list[dict[str, Any]]: with path.open() as handle: return [json.loads(line) for line in handle if line.strip()] def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--sample", type=Path, required=True) parser.add_argument("--manifest", type=Path, required=True) parser.add_argument("--embeddings", type=Path, required=True) parser.add_argument("--clusters", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--target-key", default="pkd") parser.add_argument("--outer-splits", type=int, default=5) parser.add_argument("--control-crossfit-splits", type=int, default=4) parser.add_argument("--residual-selection-splits", type=int, default=3) parser.add_argument("--bootstrap-iterations", type=int, default=2000) parser.add_argument("--max-estimators", type=int, default=1000) parser.add_argument("--early-stopping-rounds", type=int, default=50) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--minimum-rmse-improvement", type=float, default=0.05) parser.add_argument("--catastrophic-fold-tolerance", type=float, default=0.10) parser.add_argument("--overwrite", action="store_true") return parser.parse_args() def main() -> None: args = parse_args() for path in (args.sample, args.manifest, args.embeddings, args.clusters): if not path.is_file(): raise FileNotFoundError(path) if args.output_dir.exists() and not args.overwrite: raise FileExistsError(args.output_dir) if args.output_dir.exists() and not args.output_dir.is_dir(): raise NotADirectoryError(args.output_dir) args.output_dir.mkdir(parents=True, exist_ok=True) started = time.monotonic() rows = read_jsonl(args.sample) manifest = read_manifest(args.manifest) clusters = read_mmseqs_clusters(args.clusters) with np.load(args.embeddings) as embedding_arrays: aligned = validate_and_align( rows, embedding_arrays, manifest, clusters, target_key=args.target_key, ) ligand_control = ligand_feature_matrix(rows) variants = [ ResidualVariant("protein_ridge_residual", aligned["protein"]), ResidualVariant( "protein_ligand_ridge_residual", np.concatenate([aligned["protein"], aligned["ligand"]], axis=1), ), ] result = evaluate_chimera( targets=aligned["targets"], ligand_control_features=ligand_control, variants=variants, development_indices=aligned["development_indices"], cluster_ids=aligned["cluster_ids"], protein_ids=aligned["protein_ids"], outer_splits=args.outer_splits, control_crossfit_splits=args.control_crossfit_splits, residual_selection_splits=args.residual_selection_splits, bootstrap_iterations=args.bootstrap_iterations, seed=args.seed, minimum_rmse_improvement=args.minimum_rmse_improvement, catastrophic_fold_tolerance=args.catastrophic_fold_tolerance, lightgbm_config=LightGBMConfig( max_estimators=args.max_estimators, early_stopping_rounds=args.early_stopping_rounds, ), ) prediction_path = args.output_dir / "predictions.jsonl" development_indices = aligned["development_indices"] with prediction_path.open("w") as handle: for index in development_indices: row = { "observation_id": str(aligned["observation_ids"][index]), "pair_id": str(aligned["pair_ids"][index]), "protein_id": str(aligned["protein_ids"][index]), "protein_cluster_id": str(aligned["cluster_ids"][index]), "original_split": str(aligned["row_splits"][index]), "development_fold": int(result["fold_assignments"][index]), "target_pkd": float(aligned["targets"][index]), "ligand_control_prediction": float(result["control_predictions"][index]), "variant_predictions": { name: float(predictions[index]) for name, predictions in result["variant_predictions"].items() }, } handle.write(json.dumps(row, sort_keys=True) + "\n") variant_reports = result["variants"] best_variant = min( variant_reports, key=lambda name: variant_reports[name]["metrics"]["rmse"], ) passed_variants = [name for name, report in variant_reports.items() if report["gate"]["passed"]] report = { "schema_version": "1.0", "experiment": "project_chimera_development_gate", "status": "development_only_no_test_evaluation", "decision": "advance_to_locked_confirmation" if passed_variants else "reject", "target": "pKd", "seed": args.seed, "test_evaluations": 0, "population": { "sample_rows": len(rows), "development_rows": int(len(development_indices)), "test_rows_excluded": int(len(aligned["test_indices"])), "development_pairs": len(set(aligned["pair_ids"][development_indices].tolist())), "development_proteins": len(set(aligned["protein_ids"][development_indices].tolist())), "development_clusters": result["development_clusters"], "source_partitions": ["train", "validation"], "excluded_partitions": ["test"], }, "inputs": { "sample_sha256": sha256_file(args.sample), "manifest_sha256": sha256_file(args.manifest), "embeddings_sha256": sha256_file(args.embeddings), "clusters_sha256": sha256_file(args.clusters), }, "encoders": { "protein": "facebook/esm2_t12_35M_UR50D", "protein_revision": "6fbf070e65b0b7291e7bbcd451118c216cff79d8", "ligand": "DeepChem/ChemBERTa-77M-MLM", "ligand_revision": "ed8a5374f2024ec8da53760af91a33fb8f6a15ff", "frozen": True, }, "control": { "name": "ligand_morgan_descriptors_lightgbm", **result["control"], }, "variants": variant_reports, "best_variant": best_variant, "passed_variants": passed_variants, "folds": result["folds"], "protocol": { "outer_splits": args.outer_splits, "outer_group": "MMseqs2 50% identity / 80% bidirectional coverage cluster", "control_crossfit_splits": args.control_crossfit_splits, "residual_selection_splits": args.residual_selection_splits, "bootstrap_iterations": args.bootstrap_iterations, "bootstrap_unit": "MMseqs2 protein cluster", "minimum_rmse_improvement": args.minimum_rmse_improvement, "catastrophic_fold_tolerance": args.catastrophic_fold_tolerance, "test_partition_policy": ( "masked and excluded from all fitting, selection, and evaluation" ), }, "packages": { package: importlib.metadata.version(package) for package in ("lightgbm", "numpy", "rdkit", "scikit-learn") }, "artifacts": { "predictions": prediction_path.name, "predictions_sha256": sha256_file(prediction_path), }, "elapsed_seconds": time.monotonic() - started, } report_path = args.output_dir / "report.json" report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") print( json.dumps( { "decision": report["decision"], "best_variant": best_variant, "control_rmse": report["control"]["metrics"]["rmse"], "best_variant_rmse": report["variants"][best_variant]["metrics"]["rmse"], "best_variant_gate": report["variants"][best_variant]["gate"], "test_evaluations": 0, "report": str(report_path), }, indent=2, ) ) if __name__ == "__main__": main()