| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import re |
| import statistics |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| SLUG_LABELS = { |
| "reference": "Reference", |
| "prithvi_wxc": "Prithvi-WxC", |
| "stormcast": "StormCast", |
| "aurora": "Aurora", |
| "climax": "ClimaX", |
| "alphaearth": "AlphaEarth", |
| } |
|
|
|
|
| def load(path: Path) -> dict[str, Any]: |
| return json.loads(path.read_text(encoding="utf-8")) |
|
|
|
|
| def stats(values: list[float]) -> dict[str, float | int]: |
| values = [float(v) for v in values if math.isfinite(float(v))] |
| if not values: |
| return {"n": 0, "mean": math.nan, "std": math.nan} |
| return { |
| "n": len(values), |
| "mean": float(statistics.fmean(values)), |
| "std": float(statistics.stdev(values)) if len(values) > 1 else 0.0, |
| } |
|
|
|
|
| def seed_from_path(path: Path) -> int | None: |
| match = re.search(r"_seed_(\d+)", str(path)) |
| return int(match.group(1)) if match else None |
|
|
|
|
| def label_from_seed_dir(path: Path, prefix: str) -> str: |
| for part in path.parts: |
| if part.startswith(prefix) and "_seed_" in part: |
| slug = part[len(prefix) :].split("_seed_", 1)[0] |
| return SLUG_LABELS.get(slug, slug) |
| return "unknown" |
|
|
|
|
| def dedupe_rows(rows: list[dict[str, Any]], keys: tuple[str, ...]) -> list[dict[str, Any]]: |
| selected: dict[tuple[Any, ...], dict[str, Any]] = {} |
| for row in rows: |
| key = tuple(row.get(name) for name in keys) |
| old = selected.get(key) |
| if old is None: |
| selected[key] = row |
| continue |
| old_mtime = Path(str(old["path"])).stat().st_mtime |
| new_mtime = Path(str(row["path"])).stat().st_mtime |
| if new_mtime >= old_mtime: |
| selected[key] = row |
| return list(selected.values()) |
|
|
|
|
| def best_val_threshold(data: dict[str, Any]) -> str: |
| entries = data["splits"]["val"]["threshold_metrics"] |
| return max(entries, key=lambda key: (float(entries[key]["f1"]), -float(entries[key]["threshold"]))) |
|
|
|
|
| def collect_occupancy(run_root: Path) -> dict[str, Any]: |
| rows: list[dict[str, Any]] = [] |
| for path in sorted(run_root.glob("table3_occupancy_*_seed_*/run_*/summary.json")): |
| data = load(path) |
| threshold_key = best_val_threshold(data) |
| test = data["splits"]["test"] |
| rows.append( |
| { |
| "label": data.get("fm_family") or label_from_seed_dir(path, "table3_occupancy_"), |
| "seed": seed_from_path(path), |
| "strict_f1": float(test["threshold_metrics"][threshold_key]["f1"]), |
| "tolerant_f1": float(test["tolerant_threshold_metrics"]["t0_s3"][threshold_key]["f1"]), |
| "union_f1": float(test["tolerant_threshold_metrics"]["t3_s3"][threshold_key]["f1"]), |
| "path": str(path), |
| } |
| ) |
| return group(rows, ["strict_f1", "tolerant_f1", "union_f1"]) |
|
|
|
|
| def collect_headcontrol(run_root: Path) -> dict[str, Any]: |
| rows: list[dict[str, Any]] = [] |
| for path in sorted(run_root.glob("table2_prithvi_wxc_headcontrol_seed_*/run_*/summary.json")): |
| data = load(path) |
| seed = seed_from_path(path) |
| for row in data.get("selection_summary", {}).get("rows", []): |
| rows.append( |
| { |
| "label": "Prithvi-WxC", |
| "scope": row["scope"], |
| "seed": seed, |
| "ranking_selected_union_f1": float(row["ranking_selected_union_f1"]), |
| "decision_selected_union_f1": float(row["decision_selected_union_f1"]), |
| "decision_regret_union_f1": float(row["decision_regret_union_f1"]), |
| "selection_failure": bool(row.get("selection_failure", False)), |
| "path": str(path), |
| } |
| ) |
| grouped: dict[str, Any] = {} |
| rows = dedupe_rows(rows, ("label", "scope", "seed")) |
| for scope in sorted({str(row["scope"]) for row in rows}): |
| selected = [row for row in rows if row["scope"] == scope] |
| grouped[scope] = { |
| "n": len(selected), |
| "failure_count": int(sum(1 for row in selected if row["selection_failure"])), |
| "ranking_selected_union_f1": stats([row["ranking_selected_union_f1"] for row in selected]), |
| "decision_selected_union_f1": stats([row["decision_selected_union_f1"] for row in selected]), |
| "decision_regret_union_f1": stats([row["decision_regret_union_f1"] for row in selected]), |
| } |
| return {"rows": rows, "summary": grouped} |
|
|
|
|
| def collect_spread(run_root: Path) -> dict[str, Any]: |
| rows: list[dict[str, Any]] = [] |
| for pattern, prefix in [ |
| ("table3_spread_*_seed_*/run_*/summary.json", "table3_spread_"), |
| ("table3_reference_spread_seed_*/run_*/summary.json", "table3_reference_spread_"), |
| ]: |
| for path in sorted(run_root.glob(pattern)): |
| data = load(path) |
| headline = data["headline_metrics"] |
| label = data.get("fm_family") or ("Reference" if "reference_spread" in str(path) else label_from_seed_dir(path, prefix)) |
| rows.append( |
| { |
| "label": label, |
| "seed": seed_from_path(path), |
| "strict_f1": float(headline["strict_f1"]), |
| "spatial_f1": float(headline["same_sample_spatial_tolerance_f1"]["s4"]), |
| "ap": float(headline["strict_AP"]), |
| "path": str(path), |
| } |
| ) |
| return group(rows, ["strict_f1", "spatial_f1", "ap"]) |
|
|
|
|
| def collect_task(run_root: Path, glob_pattern: str, prefix: str, metrics_path: list[str], metric_keys: list[str]) -> dict[str, Any]: |
| rows: list[dict[str, Any]] = [] |
| for path in sorted(run_root.glob(glob_pattern)): |
| data = load(path) |
| label = data.get("fm_family") or label_from_seed_dir(path, prefix) |
| node: Any = data |
| for key in metrics_path: |
| node = node[key] |
| row = {"label": label, "seed": seed_from_path(path), "path": str(path)} |
| for key in metric_keys: |
| row[key] = float(node[key]) |
| rows.append(row) |
| return group(rows, metric_keys) |
|
|
|
|
| def group(rows: list[dict[str, Any]], metric_keys: list[str]) -> dict[str, Any]: |
| if rows and "seed" in rows[0]: |
| rows = dedupe_rows(rows, ("label", "seed")) |
| summary: dict[str, Any] = {} |
| for label in sorted({str(row["label"]) for row in rows}): |
| selected = [row for row in rows if row["label"] == label] |
| summary[label] = {"n": len(selected)} |
| for key in metric_keys: |
| summary[label][key] = stats([row[key] for row in selected]) |
| return {"rows": rows, "summary": summary} |
|
|
|
|
| def fmt(value: dict[str, Any], scale: float = 1.0, digits: int = 2) -> str: |
| if int(value["n"]) == 0: |
| return "missing" |
| return f"{float(value['mean']) * scale:.{digits}f} +/- {float(value['std']) * scale:.{digits}f} (n={int(value['n'])})" |
|
|
|
|
| def write_markdown(out: Path, summary: dict[str, Any]) -> None: |
| lines = ["# Forced Mean/Std Gap-Fill Summary", ""] |
| for section in [ |
| "table2_headcontrol", |
| "table3_occupancy", |
| "table3_spread", |
| "table4_final_area", |
| "table4_analog", |
| "table4_smoke", |
| "table4_heat", |
| ]: |
| lines += [f"## {section}", ""] |
| sec = summary.get(section, {}).get("summary", {}) |
| if section == "table2_headcontrol": |
| for scope, row in sec.items(): |
| lines.append( |
| f"- {scope}: regret {fmt(row['decision_regret_union_f1'], 100.0)}; " |
| f"ranking union {fmt(row['ranking_selected_union_f1'], 100.0)}; " |
| f"decision union {fmt(row['decision_selected_union_f1'], 100.0)}; " |
| f"failures {row['failure_count']}/{row['n']}" |
| ) |
| else: |
| for label, row in sec.items(): |
| pieces = [f"{key} {fmt(val, 100.0 if key.endswith('_f1') or key == 'ap' else 1.0)}" for key, val in row.items() if isinstance(val, dict)] |
| lines.append(f"- {label}: " + "; ".join(pieces)) |
| lines.append("") |
| out.write_text("\n".join(lines), encoding="utf-8") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--run-root", type=Path, default=Path("${RUN_ROOT}")) |
| parser.add_argument("--out-json", type=Path, default=Path("${OUT_JSON}")) |
| parser.add_argument("--out-md", type=Path, default=Path("${OUT_MD}")) |
| args = parser.parse_args() |
|
|
| summary = { |
| "run_root": str(args.run_root), |
| "table2_headcontrol": collect_headcontrol(args.run_root), |
| "table3_occupancy": collect_occupancy(args.run_root), |
| "table3_spread": collect_spread(args.run_root), |
| "table4_final_area": collect_task(args.run_root, "table4_final_area_*_seed_*/run_*/summary.json", "table4_final_area_", ["headline_metrics"], ["log_rmse", "log_mae", "log_spearman"]), |
| "table4_analog": collect_task(args.run_root, "table4_analog_*_seed_*/run_*/summary.json", "table4_analog_", ["test_metrics"], ["ndcg_at_10", "log_rmse", "log_mae"]), |
| "table4_smoke": collect_task(args.run_root, "table4_smoke_*_seed_*/run_*/summary.json", "table4_smoke_", ["test_metrics"], ["rmse", "mae", "pearson_r"]), |
| "table4_heat": collect_task(args.run_root, "table4_heat_*_seed_*/run_*/summary.json", "table4_heat_", ["test_metrics"], ["rmse_c", "mae_c", "pearson_r"]), |
| } |
| args.out_json.parent.mkdir(parents=True, exist_ok=True) |
| args.out_json.write_text(json.dumps(summary, indent=2), encoding="utf-8") |
| write_markdown(args.out_md, summary) |
| print(f"wrote={args.out_json}") |
| print(f"wrote={args.out_md}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
