#!/usr/bin/env python3 """Deterministically score PlantGeneBench predictions.""" from __future__ import annotations import argparse import json import math import sys from pathlib import Path from typing import Any, Dict, Iterable, List, Tuple DEFAULT_WEIGHTS = { "stage_f1": 0.35, "forbidden_ok": 0.20, "params": 0.25, "dependencies": 0.10, "analysis_type": 0.10, } STAGE_DEPENDENCIES = { "phylogeny": ["identify"], "motif": ["identify", "phylogeny"], "domain": ["identify", "phylogeny"], "structure": ["identify", "phylogeny"], "promoter": ["identify", "phylogeny"], "expression": ["identify", "phylogeny"], "synteny": ["identify"], "karyotype": ["identify", "synteny"], "interspecies_synteny": ["identify"], "ppi": ["identify", "interspecies_synteny"], "kaks": ["identify", "interspecies_synteny"], "combo": ["identify", "phylogeny", "motif", "domain", "structure"], "pep": ["identify"], "ortholog": ["identify", "phylogeny"], "reference_classification": ["phylogeny"], "rnaseq_align": ["rnaseq_qc"], "rnaseq_quant": ["rnaseq_align"], "rnaseq_sample_qc": ["rnaseq_quant"], "rnaseq_de": ["rnaseq_sample_qc"], "rnaseq_enrich": ["rnaseq_de"], "rnaseq_candidate_family": ["rnaseq_de"], } AUXILIARY_SETUP_STAGES = {"genome_prepare"} def load_jsonl(path: Path) -> List[Dict[str, Any]]: rows: List[Dict[str, Any]] = [] with path.open("r", encoding="utf-8") as handle: for lineno, line in enumerate(handle, start=1): line = line.strip() if not line: continue try: rows.append(json.loads(line)) except json.JSONDecodeError as exc: raise SystemExit(f"{path}:{lineno}: invalid JSON: {exc}") from exc return rows def dump_jsonl(path: Path, rows: Iterable[Dict[str, Any]]) -> None: path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") def as_list(value: Any) -> List[str]: if value is None: return [] if isinstance(value, list): return [str(item) for item in value] if isinstance(value, tuple): return [str(item) for item in value] return [str(value)] def normalize_scalar(value: Any) -> Any: if isinstance(value, str): return value.strip() return value def values_equal(predicted: Any, expected: Any) -> bool: if isinstance(expected, dict) and "any_of" in expected: return any(values_equal(predicted, option) for option in expected["any_of"]) if isinstance(expected, list): return sorted(as_list(predicted)) == sorted(as_list(expected)) return normalize_scalar(predicted) == normalize_scalar(expected) def dependency_closure(stages: List[str]) -> set[str]: closure: set[str] = set() stack = list(stages) while stack: stage = stack.pop() for dep in STAGE_DEPENDENCIES.get(stage, []): if dep not in closure: closure.add(dep) stack.append(dep) if stages and any(not stage.startswith("rnaseq_") for stage in stages): closure.update(AUXILIARY_SETUP_STAGES) return closure def stage_f1_predicted_set(predicted: List[str], expected: List[str]) -> List[str]: allowed_dependency_extras = dependency_closure(expected) - set(expected) return [stage for stage in predicted if stage not in allowed_dependency_extras] def f1(predicted: List[str], expected: List[str]) -> Tuple[float, float, float]: pred_set = set(predicted) exp_set = set(expected) if not pred_set and not exp_set: return 1.0, 1.0, 1.0 if not pred_set or not exp_set: return 0.0, 0.0, 0.0 tp = len(pred_set & exp_set) precision = tp / len(pred_set) recall = tp / len(exp_set) if precision + recall == 0: return precision, recall, 0.0 return precision, recall, 2 * precision * recall / (precision + recall) def choose_expected_stages(gold: Dict[str, Any], predicted: List[str]) -> Tuple[List[str], Dict[str, float]]: candidates = gold.get("expected_stage_sets") or [gold.get("expected_stages") or []] best_expected = list(candidates[0]) best_scores = {"stage_precision": 0.0, "stage_recall": 0.0, "stage_f1": -math.inf} for candidate in candidates: effective_predicted = stage_f1_predicted_set(predicted, list(candidate)) precision, recall, stage_f1 = f1(effective_predicted, list(candidate)) if stage_f1 > best_scores["stage_f1"]: best_expected = list(candidate) best_scores = {"stage_precision": precision, "stage_recall": recall, "stage_f1": stage_f1} return best_expected, best_scores def score_params(pred_params: Dict[str, Any], required: Dict[str, Any]) -> Tuple[float, List[str]]: if not required: return 1.0, [] hits = 0 errors: List[str] = [] for key, expected in required.items(): predicted = pred_params.get(key) if values_equal(predicted, expected): hits += 1 else: errors.append(f"param_mismatch:{key}") return hits / len(required), errors def score_dependencies(pred_deps: Dict[str, Any], required: Dict[str, Any]) -> Tuple[float, List[str]]: if not required: return 1.0, [] hits = 0 errors: List[str] = [] for stage, expected_deps in required.items(): predicted = pred_deps.get(stage, []) if set(as_list(expected_deps)).issubset(set(as_list(predicted))): hits += 1 else: errors.append(f"dependency_mismatch:{stage}") return hits / len(required), errors def normalize_weights(weights: Dict[str, float]) -> Dict[str, float]: merged = dict(DEFAULT_WEIGHTS) merged.update({key: float(value) for key, value in (weights or {}).items()}) total = sum(value for value in merged.values() if value > 0) if total <= 0: return DEFAULT_WEIGHTS return {key: max(value, 0.0) / total for key, value in merged.items()} def score_one(pred_row: Dict[str, Any], gold_row: Dict[str, Any]) -> Dict[str, Any]: prediction = pred_row.get("prediction") or {} gold = gold_row.get("gold") or {} predicted_stages = as_list(prediction.get("stages")) expected_stages, stage_scores = choose_expected_stages(gold, predicted_stages) forbidden = set(as_list(gold.get("forbidden_stages"))) forbidden_hits = sorted(set(predicted_stages) & forbidden) forbidden_ok = 0.0 if forbidden_hits else 1.0 params_score, param_errors = score_params(prediction.get("params") or {}, gold.get("required_params") or {}) dep_score, dep_errors = score_dependencies(prediction.get("dependencies") or {}, gold.get("required_dependencies") or {}) expected_analysis_type = str(gold.get("analysis_type") or "") predicted_analysis_type = str(prediction.get("analysis_type") or "") analysis_type_score = 1.0 if not expected_analysis_type or predicted_analysis_type == expected_analysis_type else 0.0 weights = normalize_weights(gold_row.get("weights") or {}) components = { **stage_scores, "forbidden_ok": forbidden_ok, "params": params_score, "dependencies": dep_score, "analysis_type": analysis_type_score, } total = ( weights["stage_f1"] * components["stage_f1"] + weights["forbidden_ok"] * components["forbidden_ok"] + weights["params"] * components["params"] + weights["dependencies"] * components["dependencies"] + weights["analysis_type"] * components["analysis_type"] ) errors: List[str] = [] if components["stage_f1"] < 1.0: errors.append("stage_set_mismatch") if forbidden_hits: errors.extend([f"forbidden_stage:{stage}" for stage in forbidden_hits]) if analysis_type_score < 1.0: errors.append("analysis_type_mismatch") errors.extend(param_errors) errors.extend(dep_errors) if not prediction.get("success", False): errors.append("provider_failure") return { "task_id": pred_row["task_id"], "provider": pred_row.get("provider", ""), "run_id": pred_row.get("run_id", ""), "score": round(total, 6), "components": {key: round(value, 6) for key, value in components.items()}, "weights": weights, "expected_stages": expected_stages, "predicted_stages": predicted_stages, "forbidden_hits": forbidden_hits, "errors": errors, "prediction": prediction, "gold": gold, } def run(args: argparse.Namespace) -> int: predictions = load_jsonl(args.predictions) gold_by_id = {row["task_id"]: row for row in load_jsonl(args.gold)} scored: List[Dict[str, Any]] = [] missing: List[str] = [] for row in predictions: task_id = row.get("task_id") gold_row = gold_by_id.get(task_id) if not gold_row: missing.append(str(task_id)) continue scored.append(score_one(row, gold_row)) if missing: print(f"warning: missing gold rows for {len(missing)} tasks: {', '.join(missing)}", file=sys.stderr) dump_jsonl(args.out, scored) print(f"Wrote {len(scored)} scored rows to {args.out}", file=sys.stderr) return 0 def parse_args(argv: List[str]) -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--predictions", type=Path, required=True) parser.add_argument("--gold", type=Path, required=True) parser.add_argument("--out", type=Path, required=True) return parser.parse_args(argv) if __name__ == "__main__": raise SystemExit(run(parse_args(sys.argv[1:])))