"""TOML-based input helpers for the executable and driver scripts."""
from __future__ import annotations
import dataclasses
import difflib
import hashlib
from dataclasses import is_dataclass, replace
from typing import Any, Callable, cast
import os
from pathlib import Path
try: # Python 3.11+
import tomllib
except ModuleNotFoundError: # Python 3.10
import tomli as tomllib # type: ignore[import-not-found,no-redef]
from gkx.config import (
Case,
EndDampingReference,
RuntimeCollisionConfig,
RuntimeConfig,
RuntimeExpertConfig,
RuntimeNormalizationConfig,
RuntimeOutputConfig,
RuntimeParallelConfig,
RuntimePhysicsConfig,
RuntimeQuasilinearConfig,
RuntimeRunConfig,
RuntimeSpeciesConfig,
RuntimeTermsConfig,
)
RUNTIME_TOML_SCHEMA_VERSION = 1
EXECUTABLE_TOML_SHORTHAND_COMMANDS = {
"run",
"run-runtime-linear",
"scan-runtime-linear",
"run-runtime-nonlinear",
}
# Leading bytes of the binary formats a user is most likely to hand the CLI by
# mistake, mapped to what that file actually is.
_BINARY_SIGNATURES: tuple[tuple[bytes, str], ...] = (
(b"CDF\x01", "a NetCDF classic file"),
(b"CDF\x02", "a NetCDF 64-bit-offset file"),
(b"\x89HDF\r\n\x1a\n", "an HDF5 file, which is what NetCDF-4 uses"),
(b"PK\x03\x04", "a zip archive"),
(b"\x93NUMPY", "a NumPy array file"),
)
[docs]
def load_toml(path: str | Path) -> dict:
"""Load TOML, naming common binary formats instead of reporting byte offsets."""
path = Path(path)
try:
with path.open("rb") as f:
return tomllib.load(f)
except UnicodeDecodeError as exc:
binary = describe_binary_input(path)
if binary is not None:
raise ValueError(
f"{path} is {binary}, not a TOML input file. If this is a VMEC "
"equilibrium, it was not recognised as one: check that it "
"carries the wout variables (rmnc, zmns, xm, xn), or pass it "
"with --vmec to say so explicitly."
) from exc
raise ValueError(
f"{path} is not valid UTF-8, so it cannot be a TOML input file."
) from exc
except tomllib.TOMLDecodeError as exc:
raise ValueError(f"{path} is not valid TOML: {exc}") from exc
def _validate_runtime_schema_version(data: dict[str, Any]) -> None:
raw = data.get("schema_version", 0)
if isinstance(raw, bool) or not isinstance(raw, int):
raise ValueError("schema_version must be the integer 1")
if raw not in (0, RUNTIME_TOML_SCHEMA_VERSION):
raise ValueError(
f"unsupported GKX TOML schema_version {raw}; upgrade GKX or migrate "
"the deck using the input-schema documentation"
)
# Time keys that selected the removed Diffrax integrator family, mapped to the
# native replacement a deck should use instead. The dataclass merge silently
# drops keys it does not recognise, so a deck that still selects Diffrax would
# otherwise run on a different integrator without saying so.
_REMOVED_TIME_KEYS: dict[str, str] = {
"use_diffrax": (
"the native explicit integrator is now the only time-integration owner; "
'delete the key and select the scheme with method = "rk2"/"rk3"/"rk4"'
),
"diffrax_solver": (
'select the native tableau with method = "rk2", "rk3", "rk4", or "sspx3"'
),
"diffrax_adaptive": (
"the native linear owner is fixed-step; set dt from the reported CFL "
"bound, and use the nonlinear adaptive policy through fixed_dt = false"
),
"diffrax_rtol": "the native fixed-step owner is controlled by dt, not a tolerance",
"diffrax_atol": "the native fixed-step owner is controlled by dt, not a tolerance",
"diffrax_max_steps": "the native owner runs t_max / dt steps with no step ceiling",
}
[docs]
def _validate_removed_time_keys(data: dict[str, Any]) -> None:
"""Reject decks that still select the removed Diffrax integrator family."""
time_section = data.get("time")
if not isinstance(time_section, dict):
return
present = [key for key in _REMOVED_TIME_KEYS if key in time_section]
if not present:
return
details = "; ".join(f"{key}: {_REMOVED_TIME_KEYS[key]}" for key in present)
raise ValueError(
"[time] "
+ ", ".join(repr(key) for key in present)
+ " selected the Diffrax integrator, which GKX no longer ships. "
+ details
+ "."
)
[docs]
def toml_shorthand_command(data: dict[str, Any]) -> str:
"""Return the executable command used for direct TOML path shorthand."""
_ = data
return "run"
[docs]
def resolve_runtime_path(value: str | None, *, base_dir: Path) -> str | None:
"""Expand a config path and resolve it relative to ``base_dir``."""
if value is None:
return None
expanded = os.path.expanduser(os.path.expandvars(value))
if "$" in expanded:
return value
path = Path(expanded)
if not path.is_absolute():
path = (base_dir / path).resolve()
else:
path = path.resolve()
return str(path)
# Top-level tables and keys a runtime deck may carry. Anything else is a typo
# (``[collision]``) that would otherwise be dropped without a word.
RUNTIME_TOML_SECTIONS = frozenset(
{
"schema_version",
"species",
"grid",
"time",
"geometry",
"init",
"physics",
"collisions",
"normalization",
"terms",
"expert",
"output",
"quasilinear",
"parallel",
"damping_reference",
"run",
"scan",
"fit",
}
)
# Keys the command and case owners read from ``[run]`` and ``[scan]``.
RUNTIME_RUN_KEYS = frozenset(
{
"ky",
"kx",
"Nl",
"Nm",
"mu_vperp_max",
"solver",
"method",
"dt",
"steps",
"sample_stride",
"fit_signal",
"diagnostics",
"laguerre_mode",
}
)
RUNTIME_SCAN_KEYS = frozenset(
{
"ky",
"Nl",
"Nm",
"solver",
"method",
"dt",
"steps",
"sample_stride",
"fit_signal",
"warm_start",
}
)
[docs]
def reject_unknown_keys(section: str, keys: Any, allowed: Any) -> None:
"""Raise on deck keys no owner reads, naming the closest valid spelling."""
unknown = sorted(set(keys) - set(allowed))
if not unknown:
return
hints = []
for key in unknown:
match = difflib.get_close_matches(key, allowed, n=1, cutoff=0.5)
if not match:
folded = {name.lower(): name for name in allowed}
match = [folded[key.lower()]] if key.lower() in folded else []
hints.append(f"{key!r}" + (f" (did you mean {match[0]!r}?)" if match else ""))
raise ValueError(
f"{section}: unknown key(s) {', '.join(hints)}; a misspelled key would "
"otherwise run with the default value. Valid keys: "
+ ", ".join(sorted(allowed))
)
[docs]
def _merge_dataclass(base: Any, overrides: dict | None, section: str = "[deck]") -> Any:
"""Recursively merge a dict into a dataclass, returning a new instance."""
if overrides is None:
return base
if not is_dataclass(base) or isinstance(base, type):
raise TypeError("base must be a dataclass instance")
reject_unknown_keys(
section,
overrides,
base.__dataclass_fields__, # type: ignore[attr-defined]
)
updates = {}
for field in base.__dataclass_fields__.values(): # type: ignore[attr-defined]
name = field.name
if name not in overrides:
continue
value = overrides[name]
if value is None:
continue
current = getattr(base, name)
if is_dataclass(current) and isinstance(value, dict):
updates[name] = _merge_dataclass(current, value, f"{section}.{name}")
else:
updates[name] = value
return cast(Any, replace(base, **updates))
[docs]
def _normalize_geometry_overrides(overrides: dict | None) -> dict | None:
"""Return geometry overrides using the canonical runtime schema."""
if not isinstance(overrides, dict):
return overrides
return dict(overrides)
_IMEX_BETA_MAX = (
3.0e-4 # imex gains at beta 1e-4, loses at 1e-3 and 0.015 (A/B in #383)
)
[docs]
def _default_nonlinear_stepper(physics: Any) -> str:
"""Stepper of a nonlinear deck with no ``dt`` and no ``method``.
Kinetic electrons at low beta: adaptive ``imex-ark3`` (the electron streaming
CFL is gone; falls back to ``rk4-pi`` where the factor check refuses it).
Everything else: ``rk4-pi`` (about 2x fewer steps than ``rk4`` + CFL).
"""
ph = physics if isinstance(physics, dict) else {}
kinetic_electrons = not ph.get("adiabatic_electrons", True)
return (
"imex-ark3"
if kinetic_electrons and float(ph.get("beta", 0.0)) <= _IMEX_BETA_MAX
else "rk4-pi"
)
[docs]
def _normalize_time_overrides(
overrides: Any, physics: Any = None
) -> dict[str, Any] | None:
"""Pair the integrator with the step policy when the deck chose no step.
``ExplicitTimeConfig.dt`` is a required field and ``TimeConfig.dt`` carries
a default of 0.1. That is the whole of the real difference between the two
time-configuration surfaces: a caller of the library struct has always
chosen a step, a deck may not have. So the two ``fixed_dt`` defaults are
not the same question and must not simply be made equal -- ``fixed_dt=True``
is right for a step the user chose, and fourteen shipped decks and parity
fixtures omit ``fixed_dt`` and depend on it.
What is wrong is applying a *defaulted* step as a fixed one. Measured on
the shipped Cyclone deck at (Nl,Nm)=(4,8), ky=0.3, with the ``TimeConfig``
dataclass defaults, against that rung's certified adaptive eigenpair
gamma=0.10128645:
=================== ========== ======= ========== ==============
method / policy gamma steps RHS evals rel. error
=================== ========== ======= ========== ==============
rk2, fixed dt=0.1 FAILS -- -- FloatingPointError
rk4, fixed dt=0.1 FAILS -- -- FloatingPointError
rk2, CFL-controlled 0.10126899 7806 15612 -1.72e-04
rk3, CFL-controlled 0.10126041 4512 13536 -2.57e-04
rk4, CFL-controlled 0.10125984 2768 11072 -2.63e-04
=================== ========== ======= ========== ==============
The defaulted dt=0.1 is ~7.8x the CFL-stable 0.01281 for this deck, so both
fixed-step arms overflow whatever the scheme. With the controller on, all
three schemes agree to better than 3e-4 and rk4 reaches the same horizon in
29.1% fewer right-hand-side evaluations -- exactly 4/(2.82*2), because rk4's
CFL prefactor buys a 2.82x longer step for twice the work per step. rk4 is
therefore paired with the controller and only with it: at a *fixed* step
rk4 has no step-size compensation and is simply twice the cost of rk2, so
the dataclass default stays rk2 for a deck that chose its own dt.
A deck that sets ``dt`` keeps today's behaviour exactly, and every shipped
deck sets it, so this changes none of them. ``fixed_dt = false`` with ``dt``
takes the controller's stepper: adaptive rk2 grows every streaming mode.
"""
raw = dict(overrides) if isinstance(overrides, dict) else {}
if "dt" in raw and raw.setdefault("fixed_dt", True):
raw.setdefault("method", "rk2")
return raw
# No step was chosen: hand the step to the CFL controller and give it the
# scheme that makes the controller's step cheapest. An explicit ``method``
# or ``fixed_dt`` in the deck still wins.
raw.setdefault("fixed_dt", False)
nonlinear = isinstance(physics, dict) and bool(physics.get("nonlinear", False))
raw.setdefault(
"method", _default_nonlinear_stepper(physics) if nonlinear else "rk4"
)
return raw
[docs]
def _runtime_base_config(data: dict[str, Any]) -> RuntimeConfig:
"""Return a runtime config after applying common dataclass sections."""
cfg = RuntimeConfig()
sections = {
"grid": data.get("grid"),
"time": _normalize_time_overrides(data.get("time"), data.get("physics")),
"geometry": _normalize_geometry_overrides(data.get("geometry")),
"init": data.get("init"),
}
for name, overrides in sections.items():
cfg = replace(
cfg, **{name: _merge_dataclass(getattr(cfg, name), overrides, f"[{name}]")}
)
return cfg
[docs]
def _replace_runtime_section(
cfg: RuntimeConfig,
data: dict[str, Any],
key: str,
constructor: Callable[..., Any],
) -> RuntimeConfig:
"""Replace one runtime section when the TOML section is present."""
raw = data.get(key)
if not isinstance(raw, dict):
return cfg
return cast(RuntimeConfig, replace(cfg, **{key: constructor(**raw)}))
[docs]
def _apply_runtime_section_overrides(
cfg: RuntimeConfig,
data: dict[str, Any],
) -> RuntimeConfig:
"""Apply non-nested runtime config sections from TOML data."""
section_constructors: tuple[tuple[str, Callable[..., Any]], ...] = (
("physics", RuntimePhysicsConfig),
("collisions", RuntimeCollisionConfig),
("normalization", RuntimeNormalizationConfig),
("terms", RuntimeTermsConfig),
("expert", RuntimeExpertConfig),
("output", RuntimeOutputConfig),
("quasilinear", RuntimeQuasilinearConfig),
("parallel", RuntimeParallelConfig),
("damping_reference", EndDampingReference),
)
for key, constructor in section_constructors:
cfg = _replace_runtime_section(cfg, data, key, constructor)
return cfg
[docs]
def _runtime_species_from_toml(
species_raw: Any,
) -> tuple[RuntimeSpeciesConfig, ...] | None:
"""Parse optional ``[[species]]`` runtime entries."""
if species_raw is None:
return None
if not isinstance(species_raw, list):
raise TypeError("[[species]] entries must be provided as an array of tables")
species: list[RuntimeSpeciesConfig] = []
for item in species_raw:
if not isinstance(item, dict):
raise TypeError("Each [[species]] entry must be a table")
species.append(RuntimeSpeciesConfig(**item))
return tuple(species) if species else None
[docs]
def _resolve_runtime_config_paths(
cfg: RuntimeConfig, *, base_dir: Path
) -> RuntimeConfig:
"""Resolve every path-valued runtime field against the TOML directory."""
return replace(
cfg,
geometry=replace(
cfg.geometry,
vmec_file=resolve_runtime_path(cfg.geometry.vmec_file, base_dir=base_dir),
geometry_file=resolve_runtime_path(
cfg.geometry.geometry_file,
base_dir=base_dir,
),
),
init=replace(
cfg.init,
init_file=resolve_runtime_path(cfg.init.init_file, base_dir=base_dir),
),
output=replace(
cfg.output,
path=resolve_runtime_path(cfg.output.path, base_dir=base_dir),
restart_to_file=resolve_runtime_path(
cfg.output.restart_to_file,
base_dir=base_dir,
),
restart_from_file=resolve_runtime_path(
cfg.output.restart_from_file,
base_dir=base_dir,
),
),
quasilinear=replace(
cfg.quasilinear,
output_path=resolve_runtime_path(
cfg.quasilinear.output_path,
base_dir=base_dir,
),
),
)
[docs]
def load_runtime_from_toml(path: str | Path) -> tuple[RuntimeConfig, dict]:
"""Load unified runtime config from TOML, returning ``(cfg, data)``."""
path = Path(path)
data = load_toml(path)
_validate_runtime_schema_version(data)
_validate_removed_time_keys(data)
reject_unknown_keys(str(path), data, RUNTIME_TOML_SECTIONS)
for name, allowed in (("run", RUNTIME_RUN_KEYS), ("scan", RUNTIME_SCAN_KEYS)):
if isinstance(data.get(name), dict):
reject_unknown_keys(f"[{name}]", data[name], allowed)
base_dir = path.resolve().parent
cfg = _apply_runtime_section_overrides(_runtime_base_config(data), data)
species = _runtime_species_from_toml(data.get("species"))
if species is not None:
cfg = replace(cfg, species=species)
run = _runtime_run_from_toml(data.get("run"))
if run is not None:
cfg = replace(cfg, run=run)
return _resolve_runtime_config_paths(cfg, base_dir=base_dir), data
[docs]
def _runtime_run_from_toml(section: Any) -> RuntimeRunConfig | None:
"""Read the deck's ``[run]`` selections, or ``None`` when it has none.
Only the keys the runtime itself resolves are carried. Anything else in the
table belongs to a subcommand's own parser and is left where it is.
"""
if not isinstance(section, dict):
return None
fields = {f.name for f in dataclasses.fields(RuntimeRunConfig)}
present = {key: section[key] for key in fields if section.get(key) is not None}
if not present:
return None
return RuntimeRunConfig(**present)
[docs]
def load(path: str | Path) -> Case:
"""Load a resolved immutable GKX case from TOML."""
return load_runtime_from_toml(path)[0]
[docs]
def migrate_end_damping_reference(
path: str | Path, *, route: str, dt_step: float = 0.0
) -> RuntimeConfig:
"""Explicitly convert a legacy deck; write the returned case with to_toml().
dt_step is the actual fixed linear reference step, not an inferred input
default. timestep_free identifies RHS calls without a solver timestep.
"""
path = Path(path).resolve()
source = path.read_bytes()
cfg, _ = load_runtime_from_toml(path)
if source != path.read_bytes():
raise ValueError("source deck changed during migration")
if cfg.time.damp_ends_rate is not None or cfg.damping_reference is not None:
raise ValueError("source already has an explicit damping rate or reference")
if cfg.run.dt is not None and cfg.run.dt != cfg.time.dt:
raise ValueError("resolve conflicting run.dt and time.dt before migration")
if route == "fixed_linear" and not cfg.time.fixed_dt:
raise ValueError("adaptive linear damping has no equivalent constant rate")
if route == "timestep_free" and dt_step != 0.0:
raise ValueError("timestep_free conversion must not specify dt_step")
reference = EndDampingReference(
str(path),
hashlib.sha256(source).hexdigest(),
route,
0.1 if cfg.collisions.damp_ends_amp is None else cfg.collisions.damp_ends_amp,
cfg.time.dt,
dt_step,
cfg.collisions.damp_ends_scale_by_dt,
)
return replace(
cfg,
damping_reference=reference,
time=replace(cfg.time, damp_ends_rate=reference.rate),
collisions=replace(cfg.collisions, damp_ends_scale_by_dt=False),
)