#!/usr/bin/env python from __future__ import annotations import argparse import csv import json import tempfile from collections import Counter, defaultdict from pathlib import Path from statistics import mean, median from typing import Any from assaylens.data.access import iter_limited_rows from assaylens.data.hf_upload import auto_upload_enabled, auto_upload_folder, upload_folder_to_path from assaylens.data.schemas import ( EVE_HF_ID, REQUIRED_EVE_COLUMNS, assay_frame_coarse_key, assay_frame_key, compound_target_key, ) from assaylens.data.validate import null_rates, require_columns def write_json(path: Path, payload: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8") def write_csv(path: Path, rows: list[dict[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) if not rows: path.write_text("", encoding="utf-8") return with path.open("w", newline="", encoding="utf-8") as handle: writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys())) writer.writeheader() writer.writerows(rows) def count_rows(counter: Counter[Any], key_name: str, value_name: str = "n_rows") -> list[dict[str, Any]]: return [{key_name: str(key), value_name: count} for key, count in counter.most_common()] def summary(values: list[int]) -> dict[str, float | int]: if not values: return {"min": 0, "mean": 0, "median": 0, "p75": 0, "p90": 0, "p95": 0, "max": 0} ordered = sorted(values) def pct(q: float) -> int: return ordered[min(int(q * (len(ordered) - 1)), len(ordered) - 1)] return { "min": min(values), "mean": mean(values), "median": median(values), "p75": pct(0.75), "p90": pct(0.90), "p95": pct(0.95), "max": max(values), } def profile(rows: list[dict[str, Any]], out_dir: Path) -> str: if not rows: raise RuntimeError("No EvE rows were available for profiling") require_columns(rows[0].keys(), REQUIRED_EVE_COLUMNS, "EvE") out_dir.mkdir(parents=True, exist_ok=True) compound_target_frames: dict[tuple[Any, Any], set[tuple[Any, Any, Any]]] = defaultdict(set) fine_frames: Counter[Any] = Counter() coarse_frames: Counter[Any] = Counter() target_class_counts: Counter[Any] = Counter() release_counts: Counter[Any] = Counter() assay_ids: set[Any] = set() compounds: set[Any] = set() targets: set[Any] = set() target_class_mode_frequency: Counter[Any] = Counter() target_class_mode_viability: Counter[Any] = Counter() target_class_mode_rows: Counter[Any] = Counter() for row in rows: compounds.add(row.get("compound_id")) targets.add(row.get("target_id")) assay_ids.add(row.get("assay_id")) fine_frames[assay_frame_key(row)] += 1 coarse = assay_frame_coarse_key(row) coarse_frames[coarse] += 1 compound_target_frames[compound_target_key(row)].add(coarse) target_class_counts[row.get("target__class")] += 1 release_counts[row.get("release")] += 1 class_mode = (row.get("target__class"), row.get("mode")) target_class_mode_rows[class_mode] += 1 if row.get("frequency_flag"): target_class_mode_frequency[class_mode] += 1 if row.get("viability_flag"): target_class_mode_viability[class_mode] += 1 repeated_counts = [len(frames) for frames in compound_target_frames.values()] repeated_pairs = sum(value >= 2 for value in repeated_counts) repeated_rows = sum( 1 for row in rows if len(compound_target_frames[compound_target_key(row)]) >= 2 ) repeated_pair_rate = repeated_pairs / max(len(compound_target_frames), 1) assaylens_mode = ( "COUNTERFACTUAL_READY" if repeated_pair_rate >= 0.10 or repeated_rows >= 5000 else "ASSAY_CONDITIONED_ONLY" ) row_counts = { "n_rows": len(rows), "n_compounds": len(compounds), "n_targets": len(targets), "n_assays": len(assay_ids), "n_compound_target_pairs": len(compound_target_frames), "n_assay_frames_fine": len(fine_frames), "n_assay_frames_coarse": len(coarse_frames), "assaylens_mode": assaylens_mode, } density_rows = [ { "compound_id": compound, "target_id": target, "n_coarse_frames": len(frames), } for (compound, target), frames in compound_target_frames.items() ] flag_rows = [] for key, total in target_class_mode_rows.items(): flag_rows.append( { "target_class": key[0], "mode": key[1], "frequency_flag_rate": target_class_mode_frequency[key] / total, "viability_flag_rate": target_class_mode_viability[key] / total, "n_rows": total, } ) write_json(out_dir / "row_counts.json", row_counts) write_json(out_dir / "schema.json", {"columns": sorted(rows[0].keys())}) write_csv(out_dir / "null_rates.csv", null_rates(rows, rows[0].keys())) write_csv(out_dir / "target_class_counts.csv", count_rows(target_class_counts, "target_class")) write_csv(out_dir / "frame_counts.csv", count_rows(coarse_frames, "assay_frame_coarse")) write_csv(out_dir / "activity_by_frame.csv", activity_by_frame(rows)) write_csv(out_dir / "compound_target_frame_density.csv", density_rows) write_csv(out_dir / "release_counts.csv", count_rows(release_counts, "release")) write_csv(out_dir / "flag_prevalence.csv", flag_rows) title = ( "AssayLens: Counterfactual assay-frame modeling for functional pharmacology" if assaylens_mode == "COUNTERFACTUAL_READY" else "AssayLens: Assay-conditioned activity and artifact prediction for functional pharmacology" ) report = "\n".join( [ f"# {title}", "", f"Go/no-go mode: `{assaylens_mode}`", "", "## Row Counts", "", json.dumps(row_counts, indent=2), "", "## Repeated Frame Coverage", "", json.dumps( { "frames_per_compound_target": summary(repeated_counts), "count_pairs_with_at_least_2_frames": repeated_pairs, "count_pairs_with_at_least_3_frames": sum( value >= 3 for value in repeated_counts ), "rows_in_pairs_with_at_least_2_frames": repeated_rows, "repeated_pair_rate": repeated_pair_rate, }, indent=2, ), "", "The report models observed functional pharmacology under assay frames. " "It does not infer causal mechanism or replace orthogonal confirmation.", "", ] ) (out_dir / "coverage_report.md").write_text(report, encoding="utf-8") return assaylens_mode def activity_by_frame(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: totals: Counter[Any] = Counter() positives: Counter[Any] = Counter() quantified: Counter[Any] = Counter() pxc50_values: dict[Any, list[float]] = defaultdict(list) for row in rows: key = assay_frame_coarse_key(row) totals[key] += 1 if row.get("outcome_is_active"): positives[key] += 1 if row.get("is_quantified"): quantified[key] += 1 value = row.get("outcome_potency_pxc50") if value is not None: pxc50_values[key].append(float(value)) return [ { "assay_frame_coarse": str(key), "n_rows": total, "active_rate": positives[key] / total, "quantified_rate": quantified[key] / total, "mean_pxc50": mean(pxc50_values[key]) if pxc50_values[key] else "", } for key, total in totals.most_common() ] def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--limit", type=int, default=None) parser.add_argument("--out-dir", default="coverage") args = parser.parse_args() rows = list(iter_limited_rows(EVE_HF_ID, "train", args.limit)) out_dir = Path(args.out_dir) mode = profile(rows, out_dir) uploaded = auto_upload_folder(out_dir, "data_profiles") if auto_upload_enabled(): with tempfile.TemporaryDirectory(prefix="assaylens_reports_") as tmp_dir: reports_dir = Path(tmp_dir) (reports_dir / "coverage_report.md").write_text( (out_dir / "coverage_report.md").read_text(encoding="utf-8"), encoding="utf-8", ) upload_folder_to_path(reports_dir, "reports") if uploaded: print(f"uploaded coverage artifacts to {uploaded}/data_profiles") print(mode) if __name__ == "__main__": main()