"""Build the leaderboard notebook's two data files.
python build_data.py results.parquet out/*.jsonl
`results.parquet` is the leaderboard's results dataset
(proxima-fusion/constellaration-bench-results, parquet conversion). From it,
data/entries.parquet keeps one row per submission and one per MHD-stable
point: who, when, the recorded objectives, feasibility and score. From the
JSON lines of scripts/rescore.py, data/rescored.parquet keeps one row per
boundary we evaluated, with the recorded values beside ours; a boundary
evaluated again keeps one row, and the build stops when the flux-compression
files cover only part of the first- or second-ranked front. For the
evaluate notebook, ../evaluate/data/leading-front.json keeps the feasible
(L∇B, A) points of the leading MHD-stable submission. No file holds
boundary coefficients.
"""
import json
import sys
from pathlib import Path
import pandas as pd
TOLERANCE = 0.01
RUN_KINDS = {"kernel": "kernel image", "top": "leaderboard environment", "space": "leaderboard environment", "flux": "flux compression only"}
def rows_from(path: Path):
environment, recorded = None, {}
for line in path.read_text().splitlines():
record = json.loads(line)
if record["kind"] == "environment":
environment, recorded = record, record.get("recorded", {})
continue
entry = record["entry"]
user, _, submitted = entry.rpartition("@")
known = recorded.get(entry, {})
row = {
"source_file": path.name,
"run": next((label for key, label in RUN_KINDS.items() if path.name.startswith(key)), path.stem),
"python": environment["python"],
"booz_xform": environment["packages"]["booz-xform"],
"user": user,
"submitted": submitted,
"problem": record["problem"],
"index": record["index"],
"n_field_periods": record.get("n_field_periods"),
"poloidal_modes": (record.get("modes") or [None])[0],
"seconds": record["seconds"],
"max_rss_mib": record["max_rss_mib"],
"error": record.get("error"),
"feasibility": record.get("feasibility"),
"score": record.get("score"),
"objective": record.get("objective"),
"flux_compression_pest": record.get("flux_compression_pest_jacobian"),
"flux_compression_vmec": record.get("flux_compression_vmec_jacobian"),
}
for name, value in (record.get("metrics") or {}).items():
row[name] = value
if known:
if record["problem"] == "mhd_stable":
feasibilities = json.loads(known["feasibilities"]) if isinstance(known["feasibilities"], str) else known["feasibilities"]
objectives = json.loads(known["objectives"]) if isinstance(known["objectives"], str) else known["objectives"]
row["recorded_feasibility"] = feasibilities[record["index"]]
row["recorded_gradient_scale_length"] = objectives[record["index"]][0][0]
row["recorded_aspect_ratio"] = objectives[record["index"]][1][0]
row["recorded_entry_score"] = known["score"]
else:
row["recorded_feasibility"] = known["feasibility"]
row["recorded_objective"] = known["objective"]
row["recorded_entry_score"] = known["score"]
yield row
def entries_from(results: Path) -> pd.DataFrame:
frame = pd.read_parquet(results)
frame = frame[(frame["user"] != "test") & frame["evaluated"]]
rows = []
for _, entry in frame.iterrows():
base = {"user": entry["user"], "submitted": entry["submission_time"], "problem": entry["problem_type"], "entry_score": entry["score"]}
if entry["problem_type"] == "mhd_stable":
objectives = json.loads(entry["objectives"]) if entry["objectives"] else []
feasibilities = json.loads(entry["feasibilities"]) if entry["feasibilities"] else []
for index, (objective, feasibility) in enumerate(zip(objectives, feasibilities)):
rows.append({**base, "index": index, "gradient_scale_length": objective[0][0], "aspect_ratio": objective[1][0], "feasibility": feasibility})
else:
rows.append({**base, "index": 0, "objective": entry["objective"], "feasibility": entry["feasibility"]})
return pd.DataFrame(rows)
def leading_front(entries: pd.DataFrame) -> dict:
points = entries[entries["problem"] == "mhd_stable"]
leader = points.loc[points["entry_score"].idxmax()]
chosen = points[(points["user"] == leader["user"]) & (points["submitted"] == leader["submitted"])]
feasible = chosen[chosen["feasibility"] <= TOLERANCE].sort_values("index")
return {
"user": leader["user"],
"submitted": leader["submitted"],
"recorded_score": float(leader["entry_score"]),
"columns": ["minimum_normalized_magnetic_gradient_scale_length", "aspect_ratio"],
"points": feasible[["gradient_scale_length", "aspect_ratio"]].astype(float).to_numpy().tolist(),
}
def without_reruns(rescored: pd.DataFrame) -> pd.DataFrame:
"""One row per boundary and run kind; a successful row wins over a failed one, then the later file."""
if rescored.empty:
return rescored
order = rescored.assign(_failed=rescored["error"].notna(), _position=range(len(rescored)))
order = order.sort_values(["_failed", "_position"], ascending=[False, True])
kept = order.drop_duplicates(["run", "user", "submitted", "problem", "index"], keep="last")
return kept.sort_values("_position").drop(columns=["_failed", "_position"])
def check_flux_fronts(entries: pd.DataFrame, rescored: pd.DataFrame) -> None:
"""Fail when the flux-compression run covers only part of the first- or second-ranked front.
Ranks 1 and 2 are evaluated point by point in shards; ranks 3 and 4 are
samples and are not checked. scripts/rescore.py ranks after dropping
resubmissions of identical boundaries, which have identical scores.
"""
flux = rescored[rescored["run"] == RUN_KINDS["flux"]]
if flux.empty:
return
points = entries[entries["problem"] == "mhd_stable"]
ranked = points.groupby(["user", "submitted"], as_index=False).agg(score=("entry_score", "first"), points=("index", "size"))
ranked = ranked.sort_values(["score", "submitted"], ascending=[False, True]).drop_duplicates("score")
for rank, leader in enumerate(ranked.head(2).itertuples(), start=1):
found = flux[(flux["user"] == leader.user) & (flux["submitted"] == leader.submitted)]["index"]
missing = sorted(set(range(leader.points)) - set(found))
if missing:
sys.exit(f"flux compression of rank {rank} ({leader.user}) lacks {len(missing)} of {leader.points} points, from index {missing[0]}")
def main() -> None:
data = Path(__file__).parent / "data"
data.mkdir(exist_ok=True)
entries = entries_from(Path(sys.argv[1]))
entries.to_parquet(data / "entries.parquet", compression="zstd", index=False)
print(f"{len(entries)} rows -> {data / 'entries.parquet'}")
front = Path(__file__).parent.parent / "evaluate" / "data" / "leading-front.json"
front.write_text(json.dumps(leading_front(entries)) + "\n")
print(f"leading MHD-stable front -> {front}")
rescored = without_reruns(pd.DataFrame([row for name in sys.argv[2:] for row in rows_from(Path(name))]))
check_flux_fronts(entries, rescored)
rescored.to_parquet(data / "rescored.parquet", compression="zstd", index=False)
print(f"{len(rescored)} rows -> {data / 'rescored.parquet'}")
if __name__ == "__main__":
main()
SHA-256 f29892510f0ecdb32dc79327d146084a38d5cd89b2bce1f3636d303092cc0fbb