reproduction-bundle/src/summarize_molecular_progress.py

#!/usr/bin/env python3
"""Create a publishable, path-sanitized molecular campaign checkpoint."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import pandas as pd

from summarize_molecular_campaign import summarize_frame


SAFE_RESULT_KEYS = (
    "sample_index",
    "sample_seed",
    "attempt",
    "wall_seconds",
    "sampling_seconds",
    "sha256",
    "bytes",
    "valid_reconstruction",
    "smiles",
    "posterior_min",
    "posterior_max",
)


def build_report(
    shard_root: Path,
    merged_root: Path,
    results_csv: Path | None = None,
) -> dict:
    progress = json.loads((shard_root / "PROGRESS.json").read_text())
    merge_manifest = json.loads(
        (merged_root / "29" / "371" / "merge_manifest.json").read_text()
    )
    completed = int(progress["completed"])
    if completed != len(progress["results"]):
        raise ValueError("progress result count mismatch")
    if int(merge_manifest["samples"]) != completed:
        raise ValueError("partial merge does not match current progress count")
    safe_results = [
        {key: row[key] for key in SAFE_RESULT_KEYS if key in row}
        for row in progress["results"]
    ]
    if not all(row.get("valid_reconstruction") for row in safe_results):
        raise ValueError("not every merged shard has a valid reconstruction")
    smiles = [row.get("smiles") for row in safe_results]
    if len(smiles) != len(set(smiles)):
        raise ValueError("partial campaign contains duplicate SMILES")

    report = {
        "paper": "OAM1jJsMGp",
        "claim": "C4: GRM5/RRM1 molecular composition",
        "status": "partial_campaign_checkpoint",
        "scope": (
            f"{completed} completed ligands from an active "
            f"{progress['samples']}-ligand AND campaign; not a replacement "
            "for the complete paper protocol"
        ),
        "protocol": {
            "logic": progress["logic"],
            "target_indices": progress["targets"]["indices"],
            "target_names": progress["targets"]["names"],
            "num_steps": progress["num_steps"],
            "inverse_temperature_beta": progress["inverse_temperature_beta"],
            "posterior_mc_samples": progress["posterior_mc_samples"],
            "num_atoms": progress["num_atoms"],
            "compile_forward": progress["compile_forward"],
            "workers": progress["workers"],
            "max_attempts_per_sample": progress["max_attempts_per_sample"],
        },
        "campaign_target_samples": int(progress["samples"]),
        "completed_samples": completed,
        "valid_reconstructions": sum(
            bool(row["valid_reconstruction"]) for row in safe_results
        ),
        "unique_smiles": len(set(smiles)),
        "campaign_elapsed_seconds_at_checkpoint": progress["elapsed_seconds"],
        "merged_sample_sha256": merge_manifest["output_sha256"],
        "results": safe_results,
    }
    if results_csv is not None:
        if not results_csv.is_file():
            raise FileNotFoundError(results_csv)
        docking = summarize_frame(
            pd.read_csv(results_csv), samples_per_run=completed
        )
        if progress["logic"] == "logdiff_and":
            docking["diagnostic_mean_rrm1_minus_grm5"] = docking.pop(
                "mean_and_not_separation"
            )
            docking["and_not_metric_status"] = (
                "not_applicable_to_and_samples; requires a separate "
                "logdiff_and_not campaign"
            )
        report["docking"] = {
            "mode": "vina_full",
            "exhaustiveness": 32,
            **docking,
        }
    return report


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("shard_root", type=Path)
    parser.add_argument("merged_root", type=Path)
    parser.add_argument("--results-csv", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    report = build_report(args.shard_root, args.merged_root, args.results_csv)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()