# /// 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_ylim(bottom=0) _ax.set_xlabel("tolerance on the normalised constraint violation") _ax.set_ylabel("rows meeting every constraint") _ax.legend(frameon=False, loc="lower right") 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()