269 lines
9.6 KiB
Python
269 lines
9.6 KiB
Python
|
#!/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:])))
|
||
|
|
|