"""Score ConStellaration leaderboard submissions again with the installed evaluator.
Each boundary goes through the forward model once, with the settings the
official evaluator uses for its problem. The script records every metric, the
objective, the normalised constraint violations and the time and memory each
evaluation took, one JSON object per line. Scores of a whole submission are
computed from these lines afterwards by `leaderboard/build_data.py` and the
leaderboard notebook.
python rescore.py --entries baselines --out baselines.jsonl
python rescore.py --entries top --problem mhd_stable --top 4 --jobs 8 --out p3.jsonl
The leaderboard results are read from Hugging Face at run time and are not
copied anywhere.
"""
from __future__ import annotations
import argparse
import concurrent.futures
import importlib.metadata
import io
import json
import platform
import resource
import sys
import time
import traceback
import urllib.request
RESULTS_URL = (
"https://huggingface.co/datasets/proxima-fusion/constellaration-bench-results"
"/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet"
)
REFERENCE_BOUNDARY_URL = (
"https://raw.githubusercontent.com/proximafusion/constellaration/v0.2.6"
"/hugging_face_competition/inputs/boundary.json"
)
PROBLEMS = ("geometrical", "simple_to_build", "mhd_stable")
def problem_object(name: str):
from constellaration import problems
return {
"geometrical": problems.GeometricalProblem,
"simple_to_build": problems.SimpleToBuildQIStellarator,
"mhd_stable": problems.MHDStableQIStellarator,
}[name]()
def settings_for(name: str):
from constellaration import forward_model
settings = forward_model.ConstellarationSettings
return settings.default_high_fidelity_skip_qi() if name == "geometrical" else settings.default_high_fidelity()
def flux_compression_both_weights(equilibrium, settings) -> tuple[float, float]:
"""Flux compression in bad curvature averaged with the VMEC and the PEST Jacobian.
The first value is what constellaration 0.2.6 computes. The second changes
only the averaging weight, as constellaration 0.3.0 does (PR #109), and
leaves everything else, including the equilibrium, as it is.
"""
import numpy as np
from constellaration.mhd import turbulent_transport, vmec_utils
from simsopt import mhd
vmec = vmec_utils.as_simsopt_vmec(equilibrium)
alpha, theta = turbulent_transport._set_up_surface_grid(
settings=settings, n_field_periods=equilibrium.nfp, is_stellarator_symmetric=not equilibrium.lasym
)
data = mhd.vmec_fieldlines(vs=vmec, s=settings.normalized_toroidal_flux, alpha=alpha, theta1d=theta)
grad_r_dot_grad_r = equilibrium.Aminor_p**2 / (4 * settings.normalized_toroidal_flux[:, None, None]) * data.grad_s_dot_grad_s
curvature_drift = data.B_cross_kappa_dot_grad_alpha * np.sign(equilibrium.phipf[-1])
compression = np.heaviside(curvature_drift, 0) * grad_r_dot_grad_r
def averaged(jacobian):
weights = np.abs(jacobian) / np.abs(jacobian).mean(axis=(1, 2))[:, None, None]
return float(np.mean(compression * weights, axis=(1, 2))[0])
return averaged(data.sqrt_g_vmec), averaged(data.edge_toroidal_flux_over_2pi / data.B_sup_theta_pest)
def evaluate_boundary(task: dict) -> dict:
"""One forward-model run in a fresh worker process."""
from constellaration import forward_model
from constellaration.geometry import surface_rz_fourier
started = time.perf_counter()
record = {key: task[key] for key in ("entry", "problem", "index")}
try:
boundary = surface_rz_fourier.SurfaceRZFourier.model_validate_json(task["boundary"])
record["n_field_periods"] = boundary.n_field_periods
record["modes"] = list(boundary.r_cos.shape)
problem = problem_object(task["problem"])
if task.get("flux_compression_only"):
# Without the Boozer transform and QI residual, which the
# correction does not touch and which need about 3 GiB.
settings = settings_for(task["problem"]).model_copy(update={"boozer_preset_settings": None, "qi_settings": None})
metrics, equilibrium = forward_model.forward_model(boundary, settings=settings)
record["metrics"] = metrics.model_dump(mode="json")
vmec_weight, pest_weight = flux_compression_both_weights(equilibrium, settings.turbulent_settings)
record["flux_compression_vmec_jacobian"] = vmec_weight
record["flux_compression_pest_jacobian"] = pest_weight
record["objectives"] = problem.get_objectives(metrics)
else:
metrics, _ = forward_model.forward_model(boundary, settings=settings_for(task["problem"]))
record["metrics"] = metrics.model_dump(mode="json")
record["violations"] = problem._normalized_constraint_violations(metrics).tolist()
record["feasibility"] = problem.compute_feasibility(metrics)
record["feasible"] = problem.is_feasible(metrics)
if task["problem"] == "mhd_stable":
record["objectives"] = problem.get_objectives(metrics)
else:
record["objective"] = problem.get_objective(metrics)[0]
record["score"] = problem._score(metrics) if record["feasible"] else 0.0
except Exception as error: # VMEC++ non-convergence and QI failures are results too.
record["error"] = f"{type(error).__name__}: {error}"
record["traceback"] = traceback.format_exc(limit=3)
record["seconds"] = round(time.perf_counter() - started, 2)
record["max_rss_mib"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss // (1024 if sys.platform != "darwin" else 1024 * 1024)
return record
def load_results():
import pandas as pd
with urllib.request.urlopen(RESULTS_URL, timeout=120) as response:
frame = pd.read_parquet(io.BytesIO(response.read()))
return frame[frame["user"] != "test"]
def select_entries(frame, which: str, problem: str | None, top: int):
if which == "baselines":
# The first submission per problem by the paper's first author holds the
# ALM-NGOpt baselines of Cadena et al. (2025), Tables 3 and 5.
chosen = frame[frame["user"] == "scadena-pf"].sort_values("submission_time").groupby("problem_type").head(1)
else:
candidates = frame[frame["problem_type"].isin([problem] if problem else PROBLEMS)]
candidates = candidates.drop_duplicates(subset=["problem_type", "boundary_json"])
chosen = candidates.sort_values("score", ascending=False).groupby("problem_type").head(top)
return chosen
def tasks_for(entries) -> list[dict]:
tasks = []
for _, row in entries.iterrows():
entry = f"{row['user']}@{row['submission_time']}"
if row["problem_type"] == "mhd_stable":
boundaries = json.loads(row["boundary_json"])
else:
boundaries = [row["boundary_json"]]
for index, boundary in enumerate(boundaries):
tasks.append({"entry": entry, "problem": row["problem_type"], "index": index, "boundary": boundary})
return tasks
def environment() -> dict:
packages = ("constellaration", "vmecpp", "booz-xform", "simsopt", "numpy", "scipy", "jax", "pymoo")
versions = {}
for name in packages:
try:
versions[name] = importlib.metadata.version(name)
except importlib.metadata.PackageNotFoundError:
versions[name] = None
return {"python": platform.python_version(), "machine": platform.machine(), "system": platform.system(), "packages": versions}
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--entries", choices=("baselines", "top", "reference"), default="baselines")
parser.add_argument("--problem", choices=PROBLEMS)
parser.add_argument("--top", type=int, default=1)
parser.add_argument("--rank", type=int, help="with --entries top: only the entry at this rank (1 = best)")
parser.add_argument("--jobs", type=int, default=1)
parser.add_argument("--flux-compression-only", action="store_true", help="MHD-stable only: skip QI and compute the flux compression with both Jacobian weights")
parser.add_argument("--limit", type=int, help="evaluate only the first N boundaries")
parser.add_argument("--sample", type=int, help="evaluate N boundaries drawn at random (seed 0)")
parser.add_argument("--shard", help="I/N: evaluate every N-th boundary starting at I (0-based)")
parser.add_argument("--out", required=True)
args = parser.parse_args()
if args.entries == "reference":
with urllib.request.urlopen(REFERENCE_BOUNDARY_URL, timeout=60) as response:
boundary = response.read().decode()
tasks = [{"entry": "reference", "problem": p, "index": 0, "boundary": boundary} for p in ("geometrical", "simple_to_build")]
recorded = {}
else:
entries = select_entries(load_results(), args.entries, args.problem, max(args.top, args.rank or 0))
if args.rank:
entries = entries.sort_values("score", ascending=False).iloc[[args.rank - 1]]
tasks = tasks_for(entries)
recorded = {
f"{row['user']}@{row['submission_time']}": {
"problem": row["problem_type"],
"score": row["score"],
"objectives": row["objectives"],
"feasibilities": row["feasibilities"],
"objective": row["objective"],
"feasibility": row["feasibility"],
}
for _, row in entries.iterrows()
}
for task in tasks:
task["flux_compression_only"] = args.flux_compression_only
if args.sample and args.sample < len(tasks):
import random
tasks = sorted(random.Random(0).sample(tasks, args.sample), key=lambda t: t["index"])
if args.shard:
start, step = (int(x) for x in args.shard.split("/"))
tasks = tasks[start::step]
if args.limit:
tasks = tasks[: args.limit]
with open(args.out, "w") as out:
out.write(json.dumps({"kind": "environment", **environment(), "recorded": recorded}) + "\n")
out.flush()
# max_tasks_per_child=1 (Python 3.11+) gives each evaluation its own
# process, so the memory figure belongs to one evaluation.
options = {"max_tasks_per_child": 1} if sys.version_info >= (3, 11) else {}
with concurrent.futures.ProcessPoolExecutor(max_workers=args.jobs, **options) as pool:
for record in pool.map(evaluate_boundary, tasks):
out.write(json.dumps({"kind": "boundary", **record}) + "\n")
out.flush()
print(record["entry"], record["problem"], record["index"], record.get("feasibility"), record.get("error", ""), record["seconds"], "s", flush=True)
if __name__ == "__main__":
main()
SHA-256 84f443bcd25ec00b0a9609de9636ecbbf5e5e2e15e1bd7b9bc0be903b3b89045