#!/usr/bin/env python from __future__ import annotations import argparse import csv import json import os import tempfile from collections import Counter, defaultdict from pathlib import Path from typing import Any from assaylens.data.access import iter_limited_rows from assaylens.data.hf_upload import auto_upload_folder from assaylens.data.schemas import EVE_HF_ID, assay_frame_coarse_key, compound_target_key from assaylens.eval.ablations import ABLATION_ROWS from assaylens.eval.metrics import binary_metrics from assaylens.models.baselines import FeatureClassifier from assaylens.reports.markdown import write_markdown def _float(value: str) -> float: try: return float(value) except (TypeError, ValueError): return 0.0 def maybe_download_baseline_dir(baseline_dir: str) -> Path: path = Path(baseline_dir) if (path / "ablation_table.csv").exists(): return path repo_id = os.environ.get("ARTIFACT_REPO") if not repo_id: return path from huggingface_hub import hf_hub_download download_dir = Path(tempfile.mkdtemp(prefix="assaylens_models_")) for file_name in ("models/ablation_table.csv", "models/baseline_metrics.json"): try: downloaded = Path(hf_hub_download(repo_id, file_name, repo_type="dataset")) except Exception: continue (download_dir / Path(file_name).name).write_bytes(downloaded.read_bytes()) return download_dir def maybe_download_n1_metrics() -> dict[str, Any] | None: repo_id = os.environ.get("ARTIFACT_REPO") if not repo_id: return None try: from huggingface_hub import hf_hub_download path = Path(hf_hub_download(repo_id, "models/n1_metrics.json", repo_type="dataset")) return json.loads(path.read_text(encoding="utf-8")) except Exception: # noqa: BLE001 return None def append_n1_rows(rows: list[dict[str, str]], columns: list[str]) -> None: if any(row.get("model") == "N1_assaylens_multitask" for row in rows): return metrics = maybe_download_n1_metrics() if not metrics: return split_metrics = metrics.get("splits") if isinstance(metrics.get("splits"), dict) else None if split_metrics: iterator = ( (split_name, split_payload.get("n_train", ""), split_payload.get("n_test", ""), target_name, target_metrics) for split_name, split_payload in split_metrics.items() if isinstance(split_payload, dict) for target_name, target_metrics in split_payload.items() if isinstance(target_metrics, dict) ) else: iterator = ( ("n1_smoke_holdout", "1600", "400", target_name, target_metrics) for target_name, target_metrics in metrics.items() if isinstance(target_metrics, dict) ) for split_name, n_train, n_test, target_name, target_metrics in iterator: row = {column: "" for column in columns} row.update( { "model": "N1_assaylens_multitask", "split_name": str(split_name), "target_name": target_name, "n_train": str(n_train), "n_test": str(n_test), } ) if "AUPRC" in target_metrics: for key in ("positive_rate", "AUPRC", "AUROC", "Brier", "ECE", "balanced_accuracy"): row[key] = str(target_metrics.get(key, "")) if target_name == "outcome_potency_pxc50": row["MAE_pxc50"] = str(target_metrics.get("MAE", "")) row["RMSE_pxc50"] = str(target_metrics.get("RMSE", "")) row["Spearman_pxc50"] = str(target_metrics.get("Spearman", "")) if target_name == "outcome_max_activity": row["MAE_max_activity"] = str(target_metrics.get("MAE", "")) row["Spearman_max_activity"] = str(target_metrics.get("Spearman", "")) rows.append(row) def evaluate_transfer_policy( rows: list[dict[str, Any]], eligible_rows: list[dict[str, Any]], eligible_pairs: set[tuple[Any, Any]], policy_name: str, heldout_predicate: Any, fingerprint_bits: int, ) -> dict[str, Any]: train_rows = [row for row in eligible_rows if not heldout_predicate(row)] train_pairs = {compound_target_key(row) for row in train_rows} test_rows = [ row for row in eligible_rows if heldout_predicate(row) and compound_target_key(row) in train_pairs ] result: dict[str, Any] = { "policy_name": policy_name, "n_rows": len(rows), "n_eligible_pairs": len(eligible_pairs), "n_train": len(train_rows), "n_test": len(test_rows), } if not train_rows or not test_rows: result["status"] = "not_enough_rows" return result y_train = [int(bool(row.get("outcome_is_active"))) for row in train_rows] y_true = [int(bool(row.get("outcome_is_active"))) for row in test_rows] if len(set(y_train)) < 2 or len(set(y_true)) < 2: result["status"] = "insufficient_positive_labels" return result b2 = FeatureClassifier("ligand_target", ("compound_id", "target_id"), fingerprint_bits=fingerprint_bits) b3 = FeatureClassifier( "ligand_target_assay", ("compound_id", "target_id", "mode", "assay__mechanism"), fingerprint_bits=fingerprint_bits, ) b2.fit(train_rows, "outcome_is_active") b3.fit(train_rows, "outcome_is_active") b2_scores = b2.predict_proba(test_rows) b3_scores = b3.predict_proba(test_rows) b2_metrics = binary_metrics(y_true, b2_scores) b3_metrics = binary_metrics(y_true, b3_scores) target_class_breakdown = {} for target_class in sorted({str(row.get("target__class")) for row in test_rows}): idx = [i for i, row in enumerate(test_rows) if str(row.get("target__class")) == target_class] class_y = [y_true[i] for i in idx] if len(idx) < 2 or len(set(class_y)) < 2: continue class_b2 = binary_metrics(class_y, [b2_scores[i] for i in idx]) class_b3 = binary_metrics(class_y, [b3_scores[i] for i in idx]) target_class_breakdown[target_class] = { "n_test": len(idx), "cross_frame_active_AUPRC": class_b3["AUPRC"], "cross_frame_Brier": class_b3["Brier"], "delta_AUPRC": class_b3["AUPRC"] - class_b2["AUPRC"], "delta_Brier": class_b2["Brier"] - class_b3["Brier"], } rank_correlation = None try: from scipy.stats import spearmanr rank_correlation = float(spearmanr(b2_scores, b3_scores).statistic) except Exception: # noqa: BLE001 rank_correlation = None result.update({ "status": "ok", "cross_frame_active_AUPRC": b3_metrics["AUPRC"], "cross_frame_Brier": b3_metrics["Brier"], "counterfactual_rank_correlation": rank_correlation, "B2_ligand_target": b2_metrics, "B3_ligand_target_assay": b3_metrics, "delta_AUPRC": b3_metrics["AUPRC"] - b2_metrics["AUPRC"], "delta_Brier": b2_metrics["Brier"] - b3_metrics["Brier"], "target_class_breakdown": target_class_breakdown, }) return result def assay_transfer_metrics(limit: int | None, fingerprint_bits: int) -> dict[str, Any]: rows = list(iter_limited_rows(EVE_HF_ID, "train", limit)) frames_by_pair: dict[tuple[Any, Any], set[tuple[Any, Any, Any]]] = defaultdict(set) for row in rows: frames_by_pair[compound_target_key(row)].add(assay_frame_coarse_key(row)) eligible_pairs = {key for key, frames in frames_by_pair.items() if len(frames) >= 2} eligible_rows = [row for row in rows if compound_target_key(row) in eligible_pairs] frame_counts = Counter(assay_frame_coarse_key(row) for row in eligible_rows) heldout_frames = { frame for idx, (frame, _count) in enumerate(sorted(frame_counts.items(), key=lambda item: str(item[0]))) if idx % 5 == 0 } modes = sorted({row.get("mode") for row in eligible_rows}, key=lambda value: str(value)) mechanisms = sorted({row.get("assay__mechanism") for row in eligible_rows}, key=lambda value: str(value)) mode_holdout = modes[0] if modes else None mechanism_holdout = mechanisms[0] if mechanisms else None policies = [ ("coarse_frame_every_5", lambda row: assay_frame_coarse_key(row) in heldout_frames), (f"mode_holdout::{mode_holdout}", lambda row: row.get("mode") == mode_holdout), ( f"mechanism_holdout::{mechanism_holdout}", lambda row: row.get("assay__mechanism") == mechanism_holdout, ), ( "live_cell_like_technology_holdout", lambda row: str(row.get("assay__technology")).lower() in {"luminescence", "live_cell"}, ), ] policy_results = [ evaluate_transfer_policy(rows, eligible_rows, eligible_pairs, name, predicate, fingerprint_bits) for name, predicate in policies ] ok_results = [row for row in policy_results if row.get("status") == "ok"] return { "status": "ok" if ok_results else "coverage_too_sparse", "n_rows": len(rows), "n_eligible_pairs": len(eligible_pairs), "n_policies": len(policy_results), "n_ok_policies": len(ok_results), "policies": policy_results, } def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--baseline-dir", default="/tmp/assaylens_outputs/models") parser.add_argument("--skip-assay-transfer", action="store_true") parser.add_argument("--assay-transfer-limit", type=int, default=None) parser.add_argument("--assay-transfer-fingerprint-bits", type=int, default=128) parser.add_argument("--out-dir", default="/tmp/assaylens_outputs/eval") parser.add_argument("--reports-dir", default="reports") args = parser.parse_args() out = Path(args.out_dir) out.mkdir(parents=True, exist_ok=True) columns = [ "model", "split_name", "target_name", "n_train", "n_test", "positive_rate", "AUPRC", "AUROC", "Brier", "ECE", "balanced_accuracy", "threshold", "tp", "fp", "tn", "fn", "MAE_pxc50", "RMSE_pxc50", "Spearman_pxc50", "MAE_max_activity", "Spearman_max_activity", ] baseline_table = maybe_download_baseline_dir(args.baseline_dir) / "ablation_table.csv" if baseline_table.exists(): rows = list(csv.DictReader(baseline_table.open(encoding="utf-8"))) else: rows = [ {"model": model, "split_name": "pending", "target_name": "outcome_is_active"} for model in ABLATION_ROWS ] append_n1_rows(rows, columns) path = out / "ablation_table.csv" with path.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=columns) writer.writeheader() for row in rows: writer.writerow({column: row.get(column, "") for column in columns}) by_key = {(row.get("split_name"), row.get("target_name"), row.get("model")): row for row in rows} deltas = [] hard_splits = {"scaffold_split", "target_split", "release_time_split", "assay_frame_holdout"} for split_name in sorted({row.get("split_name") for row in rows}): for target_name in sorted({row.get("target_name") for row in rows}): b2 = by_key.get((split_name, target_name, "B2_ligand_target")) b3 = by_key.get((split_name, target_name, "B3_ligand_target_assay")) if not b2 or not b3: continue deltas.append( { "split_name": split_name, "target_name": target_name, "delta_AUPRC": _float(b3.get("AUPRC", "")) - _float(b2.get("AUPRC", "")), "delta_Brier": _float(b2.get("Brier", "")) - _float(b3.get("Brier", "")), "hard_split": split_name in hard_splits, } ) hard_active_improvements = [ row for row in deltas if row["hard_split"] and row["target_name"] == "outcome_is_active" and row["delta_AUPRC"] > 0 ] summary = { "n_ablation_rows": len(rows), "n_delta_rows": len(deltas), "hard_active_splits_with_b3_auprc_improvement": len(hard_active_improvements), "primary_success_signal": len(hard_active_improvements) >= 2, } (out / "metrics_by_split.json").write_text( json.dumps({"summary": summary, "deltas": deltas}, indent=2), encoding="utf-8" ) assay_transfer = None if not args.skip_assay_transfer: assay_transfer = assay_transfer_metrics( args.assay_transfer_limit, args.assay_transfer_fingerprint_bits, ) (out / "counterfactual_assay_transfer.json").write_text( json.dumps(assay_transfer, indent=2), encoding="utf-8" ) write_markdown( Path(args.reports_dir) / "assaylens_evaluation.md", "AssayLens Evaluation", [ ("Ablation Rows", f"{len(rows)} rows copied to `{path}`."), ( "Primary Scientific Test", json.dumps(summary, indent=2) + "\n\nPositive success should not be claimed unless hard-split improvements are supported by full HF Jobs results.", ), ( "Counterfactual Assay Transfer", json.dumps(assay_transfer, indent=2) if assay_transfer else "Assay-transfer metrics were skipped for this run.", ), ], ) uploaded_eval = auto_upload_folder(out, "eval") uploaded_reports = auto_upload_folder(args.reports_dir, "reports") if uploaded_eval: print(f"uploaded evaluation artifacts to {uploaded_eval}/eval") if uploaded_reports: print(f"uploaded evaluation reports to {uploaded_reports}/reports") print(f"wrote {path}") if __name__ == "__main__": main()