sciencehub-team/constellarationPublic

# /// script
# requires-python = ">=3.10,<3.14"
# dependencies = [
#     "marimo>=0.25.1",
#     "constellaration==0.2.6",
#     "pandas",
#     "pyarrow",
#     "matplotlib",
# ]
#
# [tool.uv]
# # booz-xform 0.0.9, which constellaration 0.2.6 pins, ships no wheels and
# # needs a compiler and libnetcdf to build. 0.1.0 has wheels. Both give
# # identical metrics: see README.md of sciencehub-team/constellaration/evaluate.
# override-dependencies = ["booz-xform==0.1.0"]
# ///

import marimo

__generated_with = "0.25.1"
app = marimo.App(width="medium", app_title="Evaluate a stellarator boundary")

with app.setup:
    import json
    import subprocess
    import sys
    import tempfile
    import time
    from pathlib import Path

    import numpy as np
    import pandas as pd
    from constellaration import forward_model
    from constellaration import initial_guess
    from constellaration import problems
    from constellaration.geometry import surface_rz_fourier

    EVALUATOR_VERSION = "0.2.6"

    # The constraints in the order constellaration 0.2.6 returns their
    # normalised violations (problems.py), with the metric each one reads.
    PROBLEMS = {
        "geometrical": {
            "title": "Geometrical",
            "evaluator": problems.GeometricalProblem,
            "objective": ("max_elongation", "minimise"),
            "constraints": [
                ("aspect_ratio", "≤", 4.0),
                ("average_triangularity", "≤", -0.5),
                ("|edge_rotational_transform_over_n_field_periods|", "≥", 0.3),
            ],
        },
        "simple_to_build": {
            "title": "Simple-to-build QI",
            "evaluator": problems.SimpleToBuildQIStellarator,
            "objective": ("minimum_normalized_magnetic_gradient_scale_length", "maximise"),
            "constraints": [
                ("aspect_ratio", "≤", 10.0),
                ("|edge_rotational_transform_over_n_field_periods|", "≥", 0.25),
                ("log10(qi)", "≤", -4.0),
                ("edge_magnetic_mirror_ratio", "≤", 0.2),
                ("max_elongation", "≤", 5.0),
            ],
        },
        "mhd_stable": {
            "title": "MHD-stable QI",
            "evaluator": problems.MHDStableQIStellarator,
            "objective": ("hypervolume of (minimum_normalized_magnetic_gradient_scale_length, aspect_ratio)", "maximise"),
            "constraints": [
                ("|edge_rotational_transform_over_n_field_periods|", "≥", 0.25),
                ("log10(qi)", "≤", -3.5),
                ("edge_magnetic_mirror_ratio", "≤", 0.25),
                ("flux_compression_in_regions_of_bad_curvature", "≤", 0.9),
                ("vacuum_well", "≥", 0.0),
            ],
        },
    }
    TOLERANCE = 0.01
    HYPERVOLUME_REFERENCE = (1.0, 20.0)
    # The largest peak resident memory of one QI evaluation we measured on
    # x86-64, 3,136 MiB (README.md of sciencehub-team/constellaration/evaluate,
    # "Cost of one evaluation").
    QI_MEMORY_GIB = 3.1


@app.function
def load_boundary(source) -> surface_rz_fourier.SurfaceRZFourier:
    """A plasma boundary from a JSON string, a dict, or a path to a JSON file.

    The format is the one the ConStellaration dataset and leaderboard use:
    `r_cos` and `z_sin` as (m_max + 1) × (2 n_max + 1) arrays, `n_field_periods`,
    and `is_stellarator_symmetric: true`.
    """
    if isinstance(source, surface_rz_fourier.SurfaceRZFourier):
        return source
    if isinstance(source, dict):
        return surface_rz_fourier.SurfaceRZFourier.model_validate(source)
    if isinstance(source, Path) or (isinstance(source, str) and not source.lstrip().startswith("{")):
        source = Path(source).read_text()
    return surface_rz_fourier.SurfaceRZFourier.model_validate_json(source)


@app.function
def rotating_ellipse(aspect_ratio: float, elongation: float, rotational_transform: float, n_field_periods: int) -> surface_rz_fourier.SurfaceRZFourier:
    """The rotating-ellipse boundary the paper's baselines start from, with major radius 1."""
    return initial_guess.generate_rotating_ellipse(
        aspect_ratio=aspect_ratio,
        elongation=elongation,
        rotational_transform=rotational_transform,
        n_field_periods=n_field_periods,
    )


@app.function
def compute_metrics(boundary, problem: str = "simple_to_build", fidelity: str = "high_fidelity", corrected_flux_compression: bool = False) -> dict:
    """Run the forward model with the settings the official evaluator uses for `problem`.

    Returns the metrics as a dict. The geometrical problem skips the Boozer
    transform and the QI residual, as the evaluator does; the other two
    compute everything. VMEC++ raises RuntimeError when the equilibrium does
    not converge.

    `fidelity` selects the VMEC++ preset: "high_fidelity" is what the
    leaderboard uses; "low_fidelity" and "very_low_fidelity" are cheaper
    presets for the inner loop of an optimiser, whose results should be
    checked at high fidelity. With `corrected_flux_compression` the dict also
    holds `flux_compression_pest`, the proxy averaged with the field-aligned
    Jacobian as constellaration 0.3.0 does, from the same equilibrium.
    """
    settings = forward_model.ConstellarationSettings
    chosen = settings.default_high_fidelity_skip_qi() if problem == "geometrical" else settings.default_high_fidelity()
    if fidelity != "high_fidelity":
        chosen = chosen.model_copy(update={"vmec_preset_settings": chosen.vmec_preset_settings.model_copy(update={"fidelity": fidelity})})
    metrics, equilibrium = forward_model.forward_model(load_boundary(boundary), settings=chosen)
    out = metrics.model_dump(mode="json")
    if corrected_flux_compression:
        out["flux_compression_pest"] = flux_compression_both_weights(equilibrium, chosen.turbulent_settings)[1]
    return out


@app.function
def flux_compression_both_weights(equilibrium, settings) -> tuple[float, float]:
    """Flux compression in bad curvature averaged with the VMEC and with the PEST Jacobian.

    The first value is what constellaration 0.2.6 reports. The second changes
    only the averaging weight, as constellaration 0.3.0 does
    (proximafusion/constellaration#109), on the same equilibrium.
    """
    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_squared = equilibrium.Aminor_p**2 / (4 * settings.normalized_toroidal_flux[:, None, None]) * data.grad_s_dot_grad_s
    bad_curvature = np.heaviside(data.B_cross_kappa_dot_grad_alpha * np.sign(equilibrium.phipf[-1]), 0)

    def averaged(jacobian):
        weights = np.abs(jacobian) / np.abs(jacobian).mean(axis=(1, 2))[:, None, None]
        return float(np.mean(bad_curvature * grad_r_squared * weights, axis=(1, 2))[0])

    return averaged(data.sqrt_g_vmec), averaged(data.edge_toroidal_flux_over_2pi / data.B_sup_theta_pest)


@app.function
def constraint_report(problem: str, metrics: dict) -> pd.DataFrame:
    """Each constraint of `problem` with its value, bound and normalised violation.

    The violation is (value − bound) / |bound|, signed so that a positive
    number means the constraint is broken; the evaluator accepts violations up
    to 0.01. For the vacuum well, whose bound is 0, it divides by 0.1. The
    numbers come from constellaration's own code.
    """
    spec = PROBLEMS[problem]
    evaluator = spec["evaluator"]()
    typed = forward_model.ConstellarationMetrics.model_validate({k: v for k, v in metrics.items() if k in forward_model.ConstellarationMetrics.model_fields})
    violations = evaluator._normalized_constraint_violations(typed)
    rows = []
    for (name, relation, bound), violation in zip(spec["constraints"], violations):
        if name.startswith("|"):
            value = abs(metrics[name.strip("|")])
        elif name == "log10(qi)":
            value = float(np.log10(metrics["qi"]))
        else:
            value = metrics[name]
        rows.append({
            "constraint": f"{name} {relation} {bound:g}",
            "value": value,
            "normalised violation": float(violation),
            "satisfied": bool(violation <= TOLERANCE),
        })
    return pd.DataFrame(rows)


@app.function
def hypervolume(points) -> float:
    """The MHD-stable score of feasible (gradient scale length, aspect ratio) points.

    The area in the plane of −L∇B and A that the points dominate, up to the
    reference point (−L∇B, A) = (1, 20). Pass only feasible points.
    """
    if len(points) == 0:
        return 0.0
    from pymoo.indicators import hv

    X = np.array([(-float(L), float(A)) for L, A in points])
    return float(hv.Hypervolume(ref_point=np.array(HYPERVOLUME_REFERENCE))(X))


@app.function
def added_hypervolume(candidate, front) -> float:
    """How much a feasible (L∇B, A) point would add to the hypervolume of `front`.

    `front` is a list of feasible (L∇B, A) points, for example the leading
    submission's. Zero when the candidate is dominated.
    """
    return hypervolume([*front, candidate]) - hypervolume(front)


@app.function
def evaluate(problem: str, boundary, force: bool = False) -> dict:
    """Score a boundary (or, for "mhd_stable", a set of boundaries) as the leaderboard does.

    For "geometrical" and "simple_to_build" the result holds objective,
    feasibility, feasible, score, metrics, constraints (a list of dicts as in
    `constraint_report`) and seconds. For "mhd_stable", `boundary` is a list
    of boundaries or the leaderboard's submission format (a JSON array of
    boundary JSON strings, as text or as a path), and the result holds one
    such record per boundary under "points" and the set's score, the
    hypervolume of its feasible points. A boundary that does not parse, or
    whose evaluation raises, usually because VMEC++ did not converge, gets a
    record with "error" and feasible False, where the leaderboard would reject the whole submission;
    for the single-objective problems the record has score 0. The calls are
    those of constellaration 0.2.6's `evaluate`, with one forward-model run
    per boundary.

    The two QI problems raise MemoryError before any work when the machine
    has less memory than one QI evaluation needs (`memory_shortfall`), because
    running out of memory would end the Python process. `force=True` runs
    them anyway.
    """
    if problem not in PROBLEMS:
        raise ValueError(f"problem must be one of {', '.join(PROBLEMS)}")
    refuse_without_memory(problem, force)
    if problem == "mhd_stable":
        points = [evaluate_single(problem, b) for b in boundary_list(boundary)]
        return {"problem": problem, "evaluator": EVALUATOR_VERSION, "points": points, "score": hypervolume(feasible_points(points))}
    return {"problem": problem, "evaluator": EVALUATOR_VERSION, **evaluate_single(problem, boundary)}


@app.function
def boundary_list(boundaries) -> list:
    """The boundaries of a list, or of the leaderboard's submission format.

    The submission format is a JSON array of boundary JSON strings, given as
    text or as a path to a file. A single boundary becomes a list of one.
    """
    if isinstance(boundaries, Path) or (isinstance(boundaries, str) and not boundaries.lstrip().startswith("{")):
        text = boundaries if isinstance(boundaries, str) and boundaries.lstrip().startswith("[") else Path(boundaries).read_text()
        boundaries = json.loads(text)
    if isinstance(boundaries, (str, dict, surface_rz_fourier.SurfaceRZFourier)):
        return [boundaries]
    return list(boundaries)


@app.function
def feasible_points(records) -> list[tuple[float, float]]:
    """The (L∇B, A) points of the feasible records among `records`, as `hypervolume` takes them."""
    return [
        (r["metrics"]["minimum_normalized_magnetic_gradient_scale_length"], r["metrics"]["aspect_ratio"])
        for r in records
        if r["feasible"]
    ]


@app.function
def error_record(problem: str, message: str, seconds: float) -> dict:
    """The record of a boundary whose evaluation failed: not feasible, and score 0 for the single-objective problems."""
    record = {"error": message, "feasible": False, "feasibility": None, "metrics": None, "constraints": None}
    if problem != "mhd_stable":
        record.update(objective=None, score=0.0)
    record["seconds"] = round(seconds, 1)
    return record


@app.function
def evaluate_single(problem: str, boundary) -> dict:
    """One boundary's record for `evaluate`: one forward-model run, then the evaluator's own feasibility and score."""
    started = time.perf_counter()
    try:
        metrics = compute_metrics(load_boundary(boundary), problem)
    except Exception as error:
        return error_record(problem, f"{type(error).__name__}: {error}", time.perf_counter() - started)
    evaluator = PROBLEMS[problem]["evaluator"]()
    typed = forward_model.ConstellarationMetrics.model_validate(metrics)
    feasible = evaluator.is_feasible(typed)
    record = {
        "feasibility": evaluator.compute_feasibility(typed),
        "feasible": feasible,
        "metrics": metrics,
        "constraints": constraint_report(problem, metrics).to_dict("records"),
    }
    if problem != "mhd_stable":
        record["objective"] = evaluator.get_objective(typed)[0]
        record["score"] = evaluator._score(typed) if feasible else 0.0
    record["seconds"] = round(time.perf_counter() - started, 1)
    return record


@app.function
def evaluate_many(problem: str, boundaries, jobs: int = 1, timeout: float = 600.0, force: bool = False) -> list[dict]:
    """Score each boundary in its own Python process, `jobs` at a time.

    `boundaries` is a list or the leaderboard's submission format, as
    `boundary_list` reads it. Returns one record per boundary, in order, as
    `evaluate_single` makes them, except that "seconds" is the wall time of
    the whole process, imports included. A boundary that does not parse, a
    process that runs longer than `timeout` seconds, and a process that
    crashes each give an error record with feasible False; the other
    boundaries still run. The timeout counts wall time, so keep `jobs` at or
    below the number of CPUs. For "mhd_stable",
    `hypervolume(feasible_points(records))` is the set's score.

    The process boundary does not protect the caller from running out of
    memory where one limit covers the caller and its children, as on
    ScienceHub compute: a child that exceeds the kernel's memory ends the
    whole kernel. Use `jobs=1` there. For the same reason the two QI problems
    raise MemoryError, as in `evaluate`, unless `force=True`. Each process
    imports this file again, so the file must exist on disk.
    """
    if problem not in PROBLEMS:
        raise ValueError(f"problem must be one of {', '.join(PROBLEMS)}")
    refuse_without_memory(problem, force)
    source = globals().get("__file__")
    if not source or not Path(source).is_file():
        raise RuntimeError("evaluate_many runs this notebook's file in new processes; import it from evaluate.py on disk.")
    worker = (
        "import importlib.util, json, sys\n"
        "spec = importlib.util.spec_from_file_location('constellaration_evaluate', sys.argv[1])\n"
        "module = importlib.util.module_from_spec(spec)\n"
        "spec.loader.exec_module(module)\n"
        "task = json.loads(open(sys.argv[2]).read())\n"
        "record = module.evaluate_single(task['problem'], task['boundary'])\n"
        "open(sys.argv[3], 'w').write(json.dumps(record))\n"
    )
    texts, records = [], []
    for b in boundary_list(boundaries):
        try:
            texts.append(load_boundary(b).model_dump_json())
            records.append(None)
        except Exception as error:
            texts.append(None)
            records.append(error_record(problem, f"{type(error).__name__}: {error}", 0.0))
    waiting = [i for i in reversed(range(len(texts))) if records[i] is None]
    running = {}
    with tempfile.TemporaryDirectory() as scratch:
        scratch = Path(scratch)

        def finished(i, child, seconds):
            out = scratch / f"{i}.out"
            if child.returncode == 0 and out.is_file():
                return {**json.loads(out.read_text()), "seconds": round(seconds, 1)}
            if child.returncode < 0:
                reason = f"killed by signal {-child.returncode}"
                if child.returncode == -9:
                    reason += ", usually because it ran out of memory"
            else:
                lines = (scratch / f"{i}.err").read_text().strip().splitlines()
                reason = lines[-1] if lines else f"exit status {child.returncode}"
            return error_record(problem, f"ProcessError: {reason}", seconds)

        try:
            while waiting or running:
                while waiting and len(running) < max(1, jobs):
                    i = waiting.pop()
                    (scratch / f"{i}.json").write_text(json.dumps({"problem": problem, "boundary": texts[i]}))
                    with open(scratch / f"{i}.err", "w") as err:
                        command = [sys.executable, "-c", worker, str(source), str(scratch / f"{i}.json"), str(scratch / f"{i}.out")]
                        running[i] = (subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=err), time.perf_counter())
                time.sleep(0.2)
                for i, (child, started) in list(running.items()):
                    seconds = time.perf_counter() - started
                    if child.poll() is not None:
                        records[i] = finished(i, child, seconds)
                    elif seconds > timeout:
                        child.kill()
                        child.wait()
                        records[i] = error_record(problem, f"TimeoutError: no result after {timeout:g} s", seconds)
                    else:
                        continue
                    del running[i]
        finally:
            for child, _ in running.values():
                child.kill()
                child.wait()
    return records


@app.function
def leader_front(path=None) -> list[tuple[float, float]]:
    """The feasible (L∇B, A) points of the leading MHD-stable submission.

    Read from `data/leading-front.json` beside this notebook, or from `path`.
    The leaderboard notebook's `build_data.py` writes that file from the
    leaderboard as recorded on 25 October 2026; the points' hypervolume is
    the leader's score. Pass them to `added_hypervolume` as the front.
    """
    if path is None:
        path = Path(globals()["__file__"]).parent / "data" / "leading-front.json"
    return [tuple(point) for point in json.loads(Path(path).read_text())["points"]]


@app.function
def memory_limit_gib() -> float | None:
    """The memory available to this process in GiB, or None when it cannot tell.

    Reads the container's limit from cgroup v2, then cgroup v1, which reports
    no limit as a number near 2**63. Without a limit there, or without
    cgroup files, it falls back to MemTotal in /proc/meminfo. macOS has
    neither and gives None.
    """
    for name in ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"):
        try:
            raw = Path(name).read_text().strip()
        except OSError:
            continue
        if raw != "max" and int(raw) < 2**60:
            return int(raw) / 2**30
        break
    try:
        for line in Path("/proc/meminfo").read_text().splitlines():
            if line.startswith("MemTotal:"):
                return int(line.split()[1]) / 2**20
    except (OSError, ValueError, IndexError):
        pass
    return None


@app.function
def memory_shortfall(problem: str) -> str | None:
    """Why `problem` cannot run in this machine's memory, or None when it can.

    Only the two QI problems need more than a ScienceHub kernel's 1.5 GiB.
    """
    limit = memory_limit_gib()
    if problem == "geometrical" or limit is None or limit >= QI_MEMORY_GIB:
        return None
    return (
        f"This machine has {limit:.1f} GiB of memory. The QI residual of constellaration 0.2.6 needs "
        f"about {QI_MEMORY_GIB:g} GiB, so only the geometrical problem runs here. Run the other two on "
        "a computer with at least 4 GiB."
    )


@app.function
def refuse_without_memory(problem: str, force: bool = False) -> None:
    """Raise MemoryError with `memory_shortfall`'s reason, unless `force` is set."""
    reason = None if force else memory_shortfall(problem)
    if reason:
        raise MemoryError(f"{reason} Pass force=True to run it anyway.")


@app.function
def cross_sections(boundary, toroidal_angles=None, n_theta: int = 256):
    """R and Z of the boundary in planes of constant toroidal angle φ.

    Evaluates R = Σ r_cos[m, n] cos(mθ − n N φ) and Z = Σ z_sin[m, n] sin(mθ − n N φ)
    on the stored coefficients, without solving anything. Returns a list of
    (φ, R, Z) with φ in radians.
    """
    b = load_boundary(boundary)
    nfp = b.n_field_periods
    if toroidal_angles is None:
        toroidal_angles = [0.0, np.pi / (2 * nfp), np.pi / nfp]
    r_cos, z_sin = np.asarray(b.r_cos), np.asarray(b.z_sin)
    m = np.arange(r_cos.shape[0])[:, None]
    n = np.arange(r_cos.shape[1])[None, :] - (r_cos.shape[1] - 1) // 2
    theta = np.linspace(0, 2 * np.pi, n_theta)
    out = []
    for phi in toroidal_angles:
        angle = m[..., None] * theta - (n * nfp * phi)[..., None]
        out.append((phi, (r_cos[..., None] * np.cos(angle)).sum((0, 1)), (z_sin[..., None] * np.sin(angle)).sum((0, 1))))
    return out


@app.cell
def _():
    import marimo as mo
    import matplotlib.pyplot as plt

    return mo, plt


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    # Evaluate a stellarator boundary

    A stellarator confines a fusion plasma in a twisted magnetic field made by
    external coils alone. Its performance is set mostly by the shape of the
    outermost magnetic surface, the plasma boundary, and finding a good shape
    is an optimisation problem with expensive physics in the loop.
    ConStellaration, released by Proxima Fusion with a dataset of about 180,000
    boundaries, poses three such problems with fixed evaluation code, so that
    results from different groups can be compared
    ([Cadena et al. 2025](https://arxiv.org/abs/2506.19583)).

    This notebook holds that evaluation code, pinned to **constellaration
    0.2.6**, the version the public leaderboard uses, and exposes it as
    functions another notebook can call. The booz-xform override in this
    notebook's header changes no number; `README.md` beside the notebook
    gives the comparison and the time and memory one evaluation takes. The
    project's other notebooks look at the dataset and at the leaderboard,
    where the recorded scores reproduce in the leaderboard's environment.
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## The boundary

    All three problems take a stellarator-symmetric boundary in cylindrical
    coordinates (R, φ, Z), written as a Fourier series in the poloidal angle θ
    and the toroidal angle φ:

    $$R(\theta, \varphi) = \sum_{m,n} r_{mn} \cos(m\theta - n N_\mathrm{fp}\varphi), \qquad
    Z(\theta, \varphi) = \sum_{m,n} z_{mn} \sin(m\theta - n N_\mathrm{fp}\varphi),$$

    where $N_\mathrm{fp}$ is the number of field periods. The dataset uses
    modes up to m = n = 4 and fixes $r_{00} = 1$, which leaves 80 free
    coefficients; the evaluator accepts more modes, and the leading
    geometrical entries use poloidal modes up to m = 12. The boundary below is the
    example Proxima ships with the benchmark, with four field periods.
    """)
    return


@app.cell
def _(mo):
    reference = load_boundary(mo.notebook_dir() / "data" / "reference-boundary.json")
    reference_set = [load_boundary(s) for s in json.loads((mo.notebook_dir() / "data" / "reference-boundaries.json").read_text())]
    return reference, reference_set


@app.cell
def _(mo, plt, reference):
    _fig, _ax = plt.subplots(figsize=(5.5, 4))
    for (_phi, _R, _Z), _color in zip(cross_sections(reference), ("#1f2328", "#0969da", "#9a6700")):
        _ax.plot(_R, _Z, color=_color, lw=1.6, label=f"φ = {_phi * reference.n_field_periods / np.pi:.1f} π/N")
    _ax.set_aspect("equal")
    _ax.set_xlabel("R")
    _ax.set_ylabel("Z")
    _ax.legend(frameon=False, fontsize=9)
    for _side in ("top", "right"):
        _ax.spines[_side].set_visible(False)
    _fig.tight_layout()
    mo.vstack([
        _fig,
        mo.md(
            f"Cross-sections of the reference boundary at the start, quarter and half of a field period. "
            f"N = {reference.n_field_periods} field periods, {reference.r_cos.shape[0]} × {reference.r_cos.shape[1]} coefficients for R and for Z."
        ),
    ])
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## The three problems

    The forward model solves for the ideal-MHD equilibrium inside the boundary
    with VMEC++, with essentially no pressure and no net toroidal current, at
    a resolution of at least ten poloidal and toroidal modes that grows with
    the boundary's own, and computes the quantities the constraints use
    from it. A boundary is feasible when every normalised constraint violation,
    (value − bound) / |bound| with the sign chosen so that positive means
    broken, is at most 0.01.

    **Geometrical.** Minimise the maximum elongation of the cross-sections at
    aspect ratio at most 4, average triangularity at most −0.5 and edge
    rotational transform per field period at least 0.3. Score
    1 − (ε_max − 1)/9, clipped to [0, 1], when feasible. No quasi-isodynamicity is required; this
    is the warm-up problem.

    **Simple-to-build QI.** Maximise the smallest magnetic gradient scale
    length on the boundary, $L_{\nabla B} = B\sqrt{2}/\lVert\nabla \mathbf{B}\rVert_F$,
    normalised as $\tilde L_{\nabla B} = N_\mathrm{fp} L_{\nabla B}/a$ with $a$
    the minor radius. It is a proxy for how far the coils can stand from the
    plasma (Kappel et al. 2024). The field must be quasi-isodynamic (QI):
    omnigenous, so that trapped particles have no bounce-averaged radial
    drift, with contours of |B| that close poloidally, which keeps the
    bootstrap current small. The QI residual of Goodman et al. (2023)
    measures the distance from such a field. Score $\tilde L_{\nabla B}/20$,
    clipped to [0, 1], when feasible.

    **MHD-stable QI.** The trade-off between $\tilde L_{\nabla B}$ and compactness,
    for QI fields that also have a magnetic well (W ≥ 0, a proxy for ideal-MHD
    stability) and little flux compression in regions of bad curvature: the
    flux-surface average at ρ = 0.7 of $|\nabla r|^2$ where the curvature
    drift is bad and zero elsewhere, a geometric proxy for turbulent
    transport. A submission is a set of
    boundaries; its score is the area its feasible points dominate in the
    plane of $-\tilde L_{\nabla B}$ and aspect ratio, measured up to the reference
    point (1, 20). The score has no upper bound.
    """)
    return


@app.cell
def _(mo):
    _rows = []
    for _key, _spec in PROBLEMS.items():
        for _name, _relation, _bound in _spec["constraints"]:
            _rows.append({"problem": _spec["title"], "objective": f"{_spec['objective'][1]} {_spec['objective'][0]}", "constraint": f"{_name} {_relation} {_bound:g}"})
    mo.ui.table(pd.DataFrame(_rows), selection=None, page_size=15, label="Constraints of constellaration 0.2.6")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Run the evaluator

    One evaluation is a full equilibrium solve at the leaderboard's
    resolution. On one x86-64 core it takes about a minute for the example
    boundary and up to about four minutes for the largest boundaries on the
    leaderboard. The two QI problems also compute a Boozer
    transform and the QI residual, which need about 3 GiB of memory, so on a
    ScienceHub kernel (1.5 GiB) this notebook refuses them. The geometrical
    problem needs 1.0 GiB for boundaries of the dataset's size. The largest
    leaderboard boundaries, with 13 × 25 modes, may not fit: our runs
    measured at most 1.72 GiB for them, in worker processes that had already
    run other evaluations, so a single evaluation may need less.

    Choose a problem and start it. The MHD-stable problem evaluates the
    two-boundary example set and reports each point and the hypervolume of
    the feasible ones. On the geometrical problem the example boundary
    breaks its constraints by a normalised violation of 1.23 and scores 0;
    the constraint table shows by how much each one misses.
    """)
    return


@app.cell
def _(mo):
    problem_choice = mo.ui.dropdown(
        options={spec["title"]: key for key, spec in PROBLEMS.items()},
        value="Geometrical",
        label="Problem",
    )
    run = mo.ui.run_button(label="Evaluate")
    mo.hstack([problem_choice, run], justify="start", gap=1)
    return problem_choice, run


@app.cell
def _(mo, problem_choice, reference, reference_set, run):
    mo.stop(not run.value, mo.md("The evaluation starts when you press **Evaluate**."))
    _problem = problem_choice.value
    _shortfall = memory_shortfall(_problem)
    mo.stop(
        _shortfall is not None,
        mo.md(f"{_shortfall} The section on calling these functions shows how to install the evaluator there."),
    )
    with mo.status.spinner(title="Solving the equilibrium with VMEC++"):
        result = evaluate(_problem, reference_set if _problem == "mhd_stable" else reference)
    return (result,)


@app.cell
def _(mo, result):
    if result["problem"] == "mhd_stable":
        _points = pd.DataFrame([
            {
                "aspect ratio": p["metrics"]["aspect_ratio"] if p["metrics"] else None,
                "L∇B": p["metrics"]["minimum_normalized_magnetic_gradient_scale_length"] if p["metrics"] else None,
                "feasibility": p["feasibility"],
                "feasible": p["feasible"],
                "error": p.get("error"),
                "seconds": p["seconds"],
            }
            for p in result["points"]
        ])
        _view = mo.vstack([mo.md(f"Hypervolume score: **{result['score']:.3f}**"), mo.ui.table(_points, selection=None)])
    elif "error" in result:
        _view = mo.md(f"The evaluation failed after {result['seconds']:.0f} s and scores 0: `{result['error']}`")
    else:
        _view = mo.vstack([
            mo.md(
                f"Objective **{result['objective']:.4f}**, feasibility {result['feasibility']:.4f} "
                f"({'feasible' if result['feasible'] else 'not feasible'}), score **{result['score']:.4f}**. "
                f"The evaluation took {result['seconds']:.0f} s."
            ),
            mo.ui.table(pd.DataFrame(result["constraints"]), selection=None),
        ])
    _view
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Calling these functions from another notebook

    **Where each problem runs.** A ScienceHub kernel has one CPU and 1.5 GiB
    of memory. That is enough for the geometrical problem on boundaries of
    the dataset's size. The two QI problems need about 3 GiB, and ScienceHub
    compute has no larger kernel yet, so evaluate them on your own computer
    or another machine with at least 4 GiB. Every dependency installs there
    without a compiler on Python 3.10 to 3.13, on Linux x86-64 and on
    macOS 14 or later with Apple silicon.

    **On your own machine** the `sciencehub` package downloads this notebook
    with its `data/` folder and imports it. Install the evaluator with the
    same booz-xform override as this notebook's header, because pip alone
    would try to build booz-xform 0.0.9 from source:

    ```sh
    echo "booz-xform==0.1.0" > overrides.txt
    uv pip install sciencehub "constellaration==0.2.6" pandas pyarrow --override overrides.txt
    ```

    ```python
    import sciencehub as sh

    bench = sh.load("sciencehub-team/constellaration/evaluate")
    start = bench.rotating_ellipse(aspect_ratio=8, elongation=1.5, rotational_transform=0.4, n_field_periods=3)
    print(bench.evaluate("simple_to_build", start)["constraints"])
    ```

    `sh.load` runs each function with the downloaded notebook's folder as
    the working directory, so a relative path such as `"my-boundary.json"`
    would be looked up there. Pass absolute paths, or boundaries already
    loaded with `load_boundary`.

    **On ScienceHub compute** a kernel cannot reach sciencehub.sh, so
    `sh.load` fails there. Copy `evaluate.py` and its `data/` folder into
    the files of your own notebook, copy the dependencies and the
    `[tool.uv]` override of this notebook's header into yours, and import
    the functions with `from evaluate import evaluate_many, leader_front`.
    The notebook's folder is the kernel's working directory and is on
    `sys.path`.

    **In an optimisation loop**, `evaluate_many(problem, boundaries, jobs,
    timeout)` scores each boundary in its own Python process, at most `jobs`
    at a time. A boundary that does not parse, a run that crashes, and a run
    that exceeds `timeout` seconds of wall time (default 600; the slowest
    run we measured took 252 s on its own CPU) each get an error record
    with feasible False (and score 0 on the single-objective problems), so
    one bad candidate does not stop the loop. Keep `jobs` at or below the
    number of CPUs, or slow runs share a CPU and reach the timeout. Each record serialises to JSON and
    can be appended to a JSON-lines log.
    `compute_metrics(..., fidelity="low_fidelity")` is cheaper for an inner
    loop; check what it finds at the leaderboard's fidelity.

    Running out of memory is different on ScienceHub compute. The kernel's
    1.5 GiB limit covers the kernel and every process it starts, so a run
    that exceeds it ends the kernel, not only the run. The kernel already
    holds constellaration's imports, and each run imports them again and
    peaks at about 1.0 GiB for the geometrical problem; we have not measured
    the two together. Use `jobs=1` in a kernel. `evaluate` and
    `evaluate_many` raise MemoryError for the two QI problems when the
    memory is below what one QI run needs; `force=True` overrides that.
    Without a container memory
    limit, as on your own computer, the operating system usually kills only
    the run, which then gets an error record like any other crash.

    **Against the current leader**, `leader_front()` returns the 688
    feasible (L∇B, A) points of the best MHD-stable submission as recorded
    on 25 October 2026, read from `data/leading-front.json`, and
    `added_hypervolume(point, leader_front())` is the area a new feasible
    point would add to that front. For a set of records from
    `evaluate_many`, `feasible_points(records)` gives the points to pass.
    """)
    return


@app.cell
def _(mo):
    leading = leader_front(mo.notebook_dir() / "data" / "leading-front.json")
    _candidate = (15.0, 10.0)
    mo.md(
        f"The leading front holds {len(leading)} feasible points with a hypervolume of "
        f"**{hypervolume(leading):.2f}**, the leader's recorded score. A feasible point at "
        f"L∇B = {_candidate[0]:g}, A = {_candidate[1]:g} would add **{added_hypervolume(_candidate, leading):.2f}** to it."
    )
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## Limits and references

    The evaluator is constellaration 0.2.6 unchanged; only booz-xform, the
    Boozer transform it calls, is 0.1.0 instead of 0.0.9 because 0.0.9 cannot
    be installed without a compiler. Version 0.3.0 corrects how the flux
    compression is averaged
    ([constellaration#109](https://github.com/proximafusion/constellaration/pull/109))
    and changes which MHD-stable boundaries are feasible; the leaderboard
    stays on 0.2.6, and so does this notebook. `compute_metrics(...,
    corrected_flux_compression=True)` adds the corrected proxy alone. A full
    0.3.0 installation also moves VMEC++ from 0.4.11 to 0.5.4 and other
    numerical dependencies, so its scores can differ in more than that one
    quantity. The equilibrium is computed in near-vacuum, and the gradient
    scale length stands in for coil complexity; neither replaces a finite-β
    equilibrium or a coil design. The Hugging Face dataset includes VMEC++
    equilibria at volume-averaged β of 1 to 5 % for many of its boundaries.
    Optimisers that work at lower VMEC++ resolution should check their
    results here, because a boundary that converges at low resolution can
    fail at this one.

    Cadena, S. A., Merlo, A., Laude, E., Bauer, A., Agrawal, A., Pascu, M.,
    Savtchouk, M., Guiraud, E., Bonauer, L., Hudson, S. R. and Kaiser, M.
    (2025). ConStellaration: A dataset of QI-like stellarator plasma
    boundaries and optimization benchmarks. *NeurIPS 2025 Datasets and
    Benchmarks Track*. [OpenReview NQSbGKlCpx](https://openreview.net/forum?id=NQSbGKlCpx),
    [arXiv:2506.19583](https://arxiv.org/abs/2506.19583).

    Code: [proximafusion/constellaration](https://github.com/proximafusion/constellaration)
    at tag v0.2.6, MIT licence, Copyright (c) 2025-present Proxima Fusion GmbH.
    The two example boundaries in `data/` are `hugging_face_competition/inputs/boundary.json`
    and `boundaries.json` at that tag, renamed and otherwise unchanged.

    Kappel, J., Landreman, M. and Malhotra, D. (2024). The magnetic gradient
    scale length explains why certain plasmas require close external magnetic
    coils. *Plasma Phys. Control. Fusion* 66, 025018.

    Goodman, A. G. et al. (2023). Constructing precisely quasi-isodynamic
    magnetic fields. *J. Plasma Phys.* 89, 905890504.
    """)
    return


if __name__ == "__main__":
    app.run()