| |
| """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() |
|
|
