sciencehub-team/constellarationPublic

# /// script
# requires-python = ">=3.11"
# dependencies = ["marimo>=0.25.1", "numpy", "pandas", "pyarrow", "matplotlib"]
# ///

import marimo

__generated_with = "0.25.1"
app = marimo.App(width="medium", app_title="The ConStellaration dataset against its benchmarks")

with app.setup:
    import numpy as np
    import pandas as pd

    DATA = "data/metrics.parquet"
    TOLERANCE = 0.01

    # Bounds of constellaration 0.2.6 (problems.py), as (column, upper or
    # lower, bound). "log10_qi" and "abs_edge_iota" are derived in `violations`.
    CONSTRAINTS = {
        "geometrical": [
            ("aspect_ratio", "upper", 4.0),
            ("average_triangularity", "upper", -0.5),
            ("abs_edge_iota", "lower", 0.3),
        ],
        "simple_to_build": [
            ("aspect_ratio", "upper", 10.0),
            ("abs_edge_iota", "lower", 0.25),
            ("log10_qi", "upper", -4.0),
            ("edge_mirror_ratio", "upper", 0.2),
            ("max_elongation", "upper", 5.0),
        ],
        "mhd_stable": [
            ("abs_edge_iota", "lower", 0.25),
            ("log10_qi", "upper", -3.5),
            ("edge_mirror_ratio", "upper", 0.25),
            ("flux_compression", "upper", 0.9),
            ("vacuum_well", "lower", 0.0),
        ],
    }
    TITLES = {"geometrical": "Geometrical", "simple_to_build": "Simple-to-build QI", "mhd_stable": "MHD-stable QI"}


@app.function
def load_metrics(path: str = DATA) -> pd.DataFrame:
    """The metrics of the 158,685 ConStellaration boundaries that passed the forward model.

    One row per boundary: `id` (the dataset's plasma_config_id),
    `n_field_periods`, `optimiser` (DESC or VMEC++), and eleven of the twelve metrics of
    constellaration's forward model under shorter names (all but the aspect
    ratio over edge rotational transform). Values are stored as
    32-bit floats.
    """
    return pd.read_parquet(path)


@app.function
def violations(problem: str, metrics: pd.DataFrame) -> pd.DataFrame:
    """Normalised constraint violations of every row for one benchmark problem.

    (value − bound)/|bound| for upper bounds and (bound − value)/|bound| for
    lower bounds, as constellaration computes them; positive means broken.
    The vacuum-well bound is 0, and the evaluator divides by 0.1 instead.
    One column per constraint plus `worst`, the largest violation.
    """
    derived = metrics.assign(
        abs_edge_iota=metrics["edge_iota_per_period"].abs(),
        log10_qi=np.log10(metrics["qi"]),
    )
    out = {}
    for column, side, bound in CONSTRAINTS[problem]:
        scale = abs(bound) if bound != 0 else 0.1
        sign = 1 if side == "upper" else -1
        out[f"{column} {'≤' if side == 'upper' else '≥'} {bound:g}"] = sign * (derived[column] - bound) / scale
    table = pd.DataFrame(out, index=metrics.index)
    table["worst"] = table.max(axis=1)
    return table


@app.function
def closest(problem: str, n: int = 10, metrics: pd.DataFrame | None = None) -> pd.DataFrame:
    """The n rows nearest to feasibility for a problem, with their worst violation and the constraint behind it.

    Their `id` selects the full boundary in the Hugging Face dataset, which
    makes these rows natural starting points for an optimiser.
    """
    metrics = load_metrics() if metrics is None else metrics
    v = violations(problem, metrics)
    best = v["worst"].nsmallest(n).index
    return pd.DataFrame({
        "id": metrics.loc[best, "id"],
        "n_field_periods": metrics.loc[best, "n_field_periods"],
        "worst violation": v.loc[best, "worst"],
        "binding constraint": v.loc[best].drop(columns="worst").idxmax(axis=1),
    }).reset_index(drop=True)


@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"""
    # The ConStellaration dataset against its benchmarks

    Proxima Fusion's ConStellaration dataset holds 182,222 stellarator plasma
    boundaries, generated to be close to quasi-isodynamic (QI), together with
    the metrics of a fixed forward model for the 158,685 of them on which that
    model succeeded ([Cadena et al. 2025](https://arxiv.org/abs/2506.19583)).
    About 149,000 come from DESC optimisations towards sampled omnigenous
    target fields and about 10,000 from optimisations with VMEC++ in the loop;
    the field periods run from 1 to 5. The same paper defines three
    optimisation benchmarks on these metrics.

    How many of these boundaries already solve a benchmark, and how far are
    the rest? This notebook applies each problem's constraints, with the
    evaluator's 1 % tolerance, to the stored metrics. The answer matters to
    anyone who wants to start an optimiser from the data or train a
    generative model on it.
    """)
    return


@app.cell
def _():
    metrics = load_metrics()
    table = {problem: violations(problem, metrics) for problem in CONSTRAINTS}
    return metrics, table


@app.cell
def _(metrics, mo, table):
    _summary, _each = [], []
    for _problem, _v in table.items():
        _closest = _v["worst"].idxmin()
        _summary.append({
            "problem": TITLES[_problem],
            "feasible rows": int((_v["worst"] <= TOLERANCE).sum()),
            "smallest worst violation": round(float(_v["worst"].min()), 3),
            "binding at the closest row": _v.drop(columns="worst").loc[_closest].idxmax(),
        })
        for _c in _v.columns.drop("worst"):
            _each.append({"problem": TITLES[_problem], "constraint": _c, "rows meeting it": int((_v[_c] <= TOLERANCE).sum())})
    mo.vstack([
        mo.md(f"Rows evaluated: **{len(metrics):,}**."),
        mo.ui.table(pd.DataFrame(_summary), selection=None),
        mo.ui.table(pd.DataFrame(_each), selection=None, page_size=13),
    ])
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    **No boundary in the dataset is feasible for any of the three problems.**
    The nearest misses break a constraint by 20 to 27 % of its bound. Each
    constraint on its own is met by many rows; it is their combination that
    the data does not reach. For the geometrical problem the limit is the
    aspect ratio: only 552 rows meet A ≤ 4 within the tolerance, and none of
    them combines it with strongly negative triangularity and enough
    rotational transform. For the two QI problems it is the QI residual
    itself: within the tolerance, 9 rows meet log₁₀ QI ≤ −4 and 169 meet
    log₁₀ QI ≤ −3.5.

    The paper reports the same situation indirectly. Its generative-model
    experiment had to relax problems 1 and 2 to find 41 and 52 feasible
    training points (§5, Fig. 8), and Ran et al. (2025, v5, §6.1.2) note
    that sweeping the released designs gives no non-zero score. Every
    non-zero leaderboard score therefore comes from boundaries outside the
    dataset.
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## How much relaxation the data needs

    Raise the tolerance to see how quickly the dataset starts to contain
    "feasible" boundaries. At a tolerance of 0.25 an aspect ratio of 5 passes
    the bound of 4, a QI residual ten times the bound passes (log₁₀ QI of −3
    against −4), and the vacuum well may reach −0.025.
    """)
    return


@app.cell
def _(mo):
    tolerance = mo.ui.slider(0.01, 1.0, step=0.01, value=0.3, label="Tolerance", show_value=True, full_width=True)
    tolerance
    return (tolerance,)


@app.cell
def _(mo, plt, table, tolerance):
    _fig, _ax = plt.subplots(figsize=(7, 3.6))
    _grid = np.linspace(0, 1, 201)
    for (_problem, _v), _color in zip(table.items(), ("#1f2328", "#0969da", "#9a6700")):
        _sorted = np.sort(_v["worst"].to_numpy())
        _ax.plot(_grid, np.searchsorted(_sorted, _grid, side="right"), color=_color, lw=1.8, label=TITLES[_problem])
    _ax.axvline(TOLERANCE, color="#cf222e", lw=0.8, ls="--")
    _ax.text(TOLERANCE + 0.01, 1.5, "evaluator's tolerance", color="#cf222e", fontsize=8)
    _ax.axvline(tolerance.value, color="#57606a", lw=0.8)
    _ax.set_yscale("symlog", linthresh=1)
    _ax.set_xlabel("tolerance on the normalised constraint violation")
    _ax.set_ylabel("rows meeting every constraint")
    _ax.legend(frameon=False, loc="upper left")
    for _side in ("top", "right"):
        _ax.spines[_side].set_visible(False)
    _fig.tight_layout()
    _counts = ", ".join(f"{TITLES[p]} {int((v['worst'] <= tolerance.value).sum()):,}" for p, v in table.items())
    mo.vstack([_fig, mo.md(f"At tolerance {tolerance.value:.2f}: {_counts}.")])
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    ## QI quality against coil simplicity

    The two QI problems reward a large normalised magnetic gradient scale
    length $\tilde L_{\nabla B}$, a stand-in for coils that can sit far from
    the plasma, while asking for a small QI residual. Boundaries with a
    smaller QI residual tend to have a longer $\tilde L_{\nabla B}$
    (Spearman ρ = −0.34), but the longest $\tilde L_{\nabla B}$ falls as the
    residual bound tightens: 13.5 at log₁₀ QI ≤ −2.5, 7.1 at −3.5 and 4.6 at
    −4. The dashed lines are the
    two QI bounds.
    """)
    return


@app.cell
def _(metrics, mo, plt):
    _fig, _ax = plt.subplots(figsize=(7, 4))
    _x = np.log10(metrics["qi"].to_numpy())
    _y = metrics["gradient_scale_length"].to_numpy()
    _ok = np.isfinite(_x) & np.isfinite(_y)
    _h = _ax.hexbin(_x[_ok], _y[_ok], gridsize=70, bins="log", cmap="Greys", mincnt=1)
    for _bound in (-4.0, -3.5):
        _ax.axvline(_bound, color="#0969da", lw=0.8, ls="--")
    _ax.set_xlabel("log₁₀ QI residual")
    _ax.set_ylabel("normalised L∇B")
    _fig.colorbar(_h, ax=_ax, label="boundaries")
    for _side in ("top", "right"):
        _ax.spines[_side].set_visible(False)
    _fig.tight_layout()
    _best = metrics.loc[_x <= -3.5, "gradient_scale_length"]
    mo.vstack([
        _fig,
        mo.md(
            f"Of the {len(_best)} boundaries with log₁₀ QI ≤ −3.5 without the tolerance, the longest normalised gradient scale "
            f"length is {_best.max():.2f}; the paper's baseline reaches 8.61 on the simple-to-build problem."
        ),
    ])
    return


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

    The rows nearest to feasibility are still useful as starting points. The
    table lists them for the problem you choose; `closest(problem, n)` returns
    the same table from Python, and the `id` selects the boundary in the
    Hugging Face dataset:

    ```python
    from datasets import load_dataset

    ds = load_dataset("proxima-fusion/constellaration", split="train",
                      columns=["plasma_config_id", "boundary.json"])
    row = ds.filter(lambda r: r["plasma_config_id"] == "…")[0]
    ```
    """)
    return


@app.cell
def _(mo):
    problem = mo.ui.dropdown(options={v: k for k, v in TITLES.items()}, value="MHD-stable QI", label="Problem")
    problem
    return (problem,)


@app.cell
def _(metrics, mo, problem):
    mo.ui.table(closest(problem.value, 15, metrics), selection=None)
    return


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

    **Data.** `data/metrics.parquet` holds the metric columns of the `default`
    configuration of
    [proxima-fusion/constellaration](https://huggingface.co/datasets/proxima-fusion/constellaration)
    (revision `3346beff`, 26 October 2026, MIT licence) for every row with
    metrics, renamed and stored as 32-bit floats, without the boundaries.
    `optimiser` is VMEC++ where the row has VMEC++ optimisation settings and
    DESC otherwise. `extract.sql` in this folder writes the file.

    **Limits.** The stored metrics come from the forward model of the paper
    (the dataset's error column is `misc.has_neurips_2025_forward_model_error`),
    not from a fresh run of constellaration 0.2.6. Upstream describes the
    numerical differences between the paper's 0.2.1 and 0.2.6 as negligible
    apart from the sign of the rotational transform. The constraints here use
    its absolute value, as 0.2.6 does; 10,113 rows have negative rotational
    transform and would have failed under 0.2.1. A difference of a few
    percent between the stored metrics and a fresh evaluation would not
    change any count above, since the nearest row misses by about 20 %.

    Cadena, S. A. et al. (2025). ConStellaration: A dataset of QI-like
    stellarator plasma boundaries and optimization benchmarks. *NeurIPS 2025
    Datasets and Benchmarks Track*. [arXiv:2506.19583](https://arxiv.org/abs/2506.19583).

    Ran, N., Wang, Y., Zhang, X., Li, Z., Ran, Q., Li, W. and Allmendinger, R.
    (2025). ExLLM: Experience-enhanced LLM optimization for molecular design
    and beyond. [arXiv:2502.12845](https://arxiv.org/abs/2502.12845).
    """)
    return


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