sciencehub-team/constellarationPublic

Edit a copy

3 definitions in dataset.py that other code can import and call. Each has the notebook's setup cell, and nothing else of it.

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.

Parameter	Type	Default
path	str	DATA

dataset.load_metrics()

Source

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)

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.

Parameter	Type	Default
problem	str	required
metrics	pd.DataFrame	required

dataset.violations(problem=..., metrics=...)

Source

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

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.

Parameter	Type	Default
problem	str	required
n	int	10
metrics	pd.DataFrame \| None	None

dataset.closest(problem=...)

Source

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)