"""Runtime executable command workflows."""
from __future__ import annotations
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence, cast
import numpy as np
from gkx.config import RuntimeConfig
from gkx.artifacts.io import (
write_quasilinear_artifacts,
write_runtime_linear_artifacts,
write_runtime_linear_scan_artifacts,
)
from gkx.workflows.linear import run_runtime_linear
from gkx.workflows.nonlinear import run_runtime_nonlinear
from gkx.workflows.runtime.orchestration_scan import run_runtime_scan
from gkx.workflows.runtime.toml import (
load_runtime_from_toml,
reject_unknown_keys,
resolve_runtime_path,
)
from gkx.workflows.runtime.startup import (
_RUNTIME_LINEAR_HL_FALLBACK,
_RUNTIME_NONLINEAR_HL_FALLBACK,
)
from gkx.workflows.runtime.warm_start import resolve_scan_warm_start
from gkx.workflows.runtime.results import (
RuntimeLinearResult,
)
from gkx.workflows.runtime.artifacts import (
print_linear_run_header,
print_nonlinear_command_outputs,
print_nonlinear_run_header,
print_nonlinear_run_summary,
run_runtime_nonlinear_with_artifacts,
write_linear_runtime_command_outputs,
write_scan_runtime_command_outputs,
)
[docs]
@dataclass(frozen=True)
class RuntimeLinearCommandOptions:
"""Resolved executable options for one linear runtime command."""
ky: float
Nl: int
Nm: int
solver: str
fit_signal: str
method: str | None
dt: float | None
steps: int | None
sample_stride: int
method_for_header: str
dt_for_header: float
steps_for_header: int
show_progress: bool
[docs]
@dataclass(frozen=True)
class RuntimeScanCommandOptions:
"""Resolved executable options for one linear scan command."""
ky_values: tuple[float, ...]
Nl: int
Nm: int
solver: str
fit_signal: str
method: str | None
dt: float | None
steps: int | None
sample_stride: int
batch_ky: bool
show_progress: bool
workers: int
parallel_executor: str
warm_start: bool | None
[docs]
@dataclass(frozen=True)
class RuntimeNonlinearCommandOptions:
"""Resolved executable options for one nonlinear runtime command."""
ky: float
Nl: int
Nm: int
method: str
dt: float
steps: int | None
sample_stride: int
diagnostics_stride: int | None
diagnostics: bool
laguerre_mode: str | None
show_progress: bool
# [time] run_to override from --until-saturated / --no-until-saturated,
# or None to keep the TOML value.
run_to: str | None = None
def _format_fitted_value(value: float, stderr: Any, precision: int) -> str:
"""Format 'value' or 'value+/-stderr' when the fit reports uncertainty."""
text = f"{value:.{precision}f}"
if stderr is not None and np.isfinite(stderr):
text += f"+/-{stderr:.{precision}f}"
return text
[docs]
def _arg_or_section(args: Any, section: dict[str, Any], name: str, default: Any) -> Any:
"""Return an executable flag override or a TOML section value."""
value = getattr(args, name, None)
if value is not None:
return value
return section.get(name, default)
[docs]
def _resolve_linear_fit_options(args: Any, section: dict[str, Any]) -> tuple[str, str]:
"""Resolve the linear eigensignal solver and fit signal."""
return (
str(_arg_or_section(args, section, "solver", "auto")),
str(_arg_or_section(args, section, "fit_signal", "auto")),
)
[docs]
def _resolve_grid_time_options(
args: Any,
section: dict[str, Any],
cfg: RuntimeConfig,
hl_fallback: tuple[int, int] = _RUNTIME_LINEAR_HL_FALLBACK,
) -> tuple[int, int, str | None, float | None, int | None, int]:
"""Resolve resolution (``hl_fallback`` from startup), time controls, cadence."""
method = _arg_or_section(args, section, "method", None)
dt = _arg_or_section(args, section, "dt", None)
steps = _arg_or_section(args, section, "steps", None)
return (
int(_arg_or_section(args, section, "Nl", hl_fallback[0])),
int(_arg_or_section(args, section, "Nm", hl_fallback[1])),
None if method is None else str(method),
None if dt is None else float(dt),
None if steps is None else int(steps),
int(_arg_or_section(args, section, "sample_stride", cfg.time.sample_stride)),
)
[docs]
def _runtime_fit_config(data: dict[str, Any]) -> dict[str, Any]:
"""Return the deck's ``[fit]`` options, refusing keys no fit owner reads."""
fit = dict(data.get("fit", {}))
reject_unknown_keys("[fit]", fit, RUNTIME_CASE_FIT_KEYS)
return fit
def _with_fit_signal(
section: dict[str, Any], fit_cfg: dict[str, Any]
) -> dict[str, Any]:
"""Move ``[fit] fit_signal`` into the command section that resolves it.
``fit_signal`` may be declared under ``[fit]`` (as the case helpers read
it) or under ``[run]``/``[scan]``; two different values are a conflict.
"""
if "fit_signal" not in fit_cfg:
return section
signal = fit_cfg.pop("fit_signal")
if section.get("fit_signal", signal) != signal:
raise ValueError(
f"fit_signal is {section['fit_signal']!r} in the command section but "
f"{signal!r} in [fit]; declare it once"
)
return {**section, "fit_signal": signal}
def _validate_linear_sampling(solver: str, steps: int, sample_stride: int) -> None:
"""Reject invalid explicit-time output cadence before solver setup."""
solver_key = solver.strip().lower().replace("-", "_")
if sample_stride < 1:
raise ValueError("sample_stride must be >= 1")
if solver_key in {"time", "explicit", "explicit_time"} and steps % sample_stride:
raise ValueError(
f"time solver steps ({steps}) must be divisible by sample_stride "
f"({sample_stride})"
)
[docs]
def should_show_progress(args: Any, configured: bool) -> bool:
"""Resolve progress output from executable flags, TOML config, and TTY state."""
import sys
if getattr(args, "progress", False):
return True
if getattr(args, "no_progress", False):
return False
return bool(configured or sys.stdout.isatty())
[docs]
def _resolve_linear_command_options(
args: Any,
cfg: RuntimeConfig,
run_cfg: dict[str, Any],
) -> RuntimeLinearCommandOptions:
"""Resolve linear command options from flags, TOML, and config defaults."""
Nl, Nm, method, dt, steps, sample_stride = _resolve_grid_time_options(
args, run_cfg, cfg
)
solver, fit_signal = _resolve_linear_fit_options(args, run_cfg)
dt_for_header = dt if dt is not None else float(cfg.time.dt)
steps_for_header = (
steps
if steps is not None
else int(round(float(cfg.time.t_max) / dt_for_header))
)
_validate_linear_sampling(solver, steps_for_header, sample_stride)
return RuntimeLinearCommandOptions(
ky=float(_arg_or_section(args, run_cfg, "ky", 0.3)),
Nl=Nl,
Nm=Nm,
solver=solver,
fit_signal=fit_signal,
method=method,
dt=dt,
steps=steps,
sample_stride=sample_stride,
method_for_header=str(method if method is not None else cfg.time.method),
dt_for_header=dt_for_header,
steps_for_header=steps_for_header,
show_progress=should_show_progress(args, bool(cfg.time.progress_bar)),
)
[docs]
def _parse_ky_values(args: Any, scan_cfg: dict[str, Any]) -> tuple[float, ...]:
"""Resolve scan ky values from CLI or TOML, failing closed on empty scans."""
raw = getattr(args, "ky_values", None)
if raw is not None:
values = tuple(float(x) for x in str(raw).split(",") if x.strip())
else:
values = tuple(float(x) for x in scan_cfg.get("ky", ()))
if not values:
raise ValueError("No ky values provided. Use --ky-values or [scan].ky in TOML.")
return values
[docs]
def _resolve_scan_command_options(
args: Any,
cfg: RuntimeConfig,
scan_cfg: dict[str, Any],
) -> RuntimeScanCommandOptions:
"""Resolve linear-scan command options from flags and TOML defaults."""
Nl, Nm, method, dt, steps, sample_stride = _resolve_grid_time_options(
args, scan_cfg, cfg
)
solver, fit_signal = _resolve_linear_fit_options(args, scan_cfg)
return RuntimeScanCommandOptions(
ky_values=_parse_ky_values(args, scan_cfg),
Nl=Nl,
Nm=Nm,
solver=solver,
fit_signal=fit_signal,
method=method,
dt=dt,
steps=steps,
sample_stride=sample_stride,
batch_ky=bool(getattr(args, "batch_ky", False)),
show_progress=should_show_progress(args, bool(cfg.time.progress_bar)),
workers=int(getattr(args, "workers", 1)),
parallel_executor=str(getattr(args, "parallel_executor", "thread")),
warm_start=resolve_scan_warm_start(args, scan_cfg),
)
[docs]
def _resolve_nonlinear_command_options(
args: Any,
cfg: RuntimeConfig,
run_cfg: dict[str, Any],
) -> RuntimeNonlinearCommandOptions:
"""Resolve nonlinear command options from flags, TOML, and config defaults."""
Nl, Nm, method, dt, steps, sample_stride = _resolve_grid_time_options(
args, run_cfg, cfg, _RUNTIME_NONLINEAR_HL_FALLBACK
)
if steps is not None:
nonlinear_steps: int | None = steps
elif bool(cfg.time.fixed_dt): # a --dt override sets the step count too
nonlinear_steps = int(round(cfg.time.t_max / (dt or cfg.time.dt)))
else:
nonlinear_steps = None
if getattr(args, "no_diagnostics", False):
diagnostics = False
elif getattr(args, "diagnostics", False):
diagnostics = True
else:
diagnostics = bool(run_cfg.get("diagnostics", cfg.time.diagnostics))
if getattr(args, "until_saturated", False):
run_to: str | None = "saturation"
elif getattr(args, "no_until_saturated", False):
run_to = "t_max"
else:
run_to = None
diagnostics_stride = getattr(args, "diagnostics_stride", None)
laguerre_mode = _arg_or_section(args, run_cfg, "laguerre_mode", None)
return RuntimeNonlinearCommandOptions(
ky=float(_arg_or_section(args, run_cfg, "ky", 0.3)),
Nl=Nl,
Nm=Nm,
dt=dt if dt is not None else float(cfg.time.dt),
steps=nonlinear_steps,
method=method if method is not None else str(cfg.time.method),
sample_stride=sample_stride,
diagnostics_stride=(
None if diagnostics_stride is None else int(diagnostics_stride)
),
diagnostics=diagnostics,
laguerre_mode=None if laguerre_mode is None else str(laguerre_mode),
show_progress=should_show_progress(args, bool(cfg.time.progress_bar)),
run_to=run_to,
)
_PRELOADED_RUNTIME_CONFIG_ATTR = "_gkx_preloaded_runtime_config"
_PRELOADED_RUNTIME_DATA_ATTR = "_gkx_preloaded_runtime_data"
[docs]
def attach_preloaded_runtime_config(
args: Any,
cfg: RuntimeConfig,
data: dict[str, Any],
) -> None:
"""Attach already-loaded runtime TOML data to a parser namespace.
The generic ``run`` dispatcher inspects a config to choose linear or
nonlinear execution. Passing the loaded object forward avoids a second TOML
parse while direct subcommands remain self-contained.
"""
setattr(args, _PRELOADED_RUNTIME_CONFIG_ATTR, cfg)
setattr(args, _PRELOADED_RUNTIME_DATA_ATTR, data)
[docs]
def load_runtime_command_config(
args: Any,
) -> tuple[RuntimeConfig, dict[str, Any]]:
"""Load runtime TOML data, reusing the generic-dispatch preload if present."""
cfg = getattr(args, _PRELOADED_RUNTIME_CONFIG_ATTR, None)
data = getattr(args, _PRELOADED_RUNTIME_DATA_ATTR, None)
if cfg is not None and data is not None:
return cast(RuntimeConfig, cfg), cast(dict[str, Any], data)
return load_runtime_from_toml(args.config)
[docs]
def _prepare_runtime_command_config(
args: Any,
*,
path_overrides: bool,
quasilinear_overrides: bool,
) -> tuple[RuntimeConfig, dict[str, Any]]:
"""Load runtime command config and apply the command-specific overrides."""
cfg, data = load_runtime_command_config(args)
if path_overrides:
cfg = apply_runtime_path_overrides(
cfg,
args,
resolve_runtime_path=resolve_runtime_path,
)
if quasilinear_overrides:
cfg = apply_quasilinear_overrides(cfg, args)
return cfg, data
[docs]
def runtime_output_path(args: Any, cfg: RuntimeConfig) -> str | None:
"""Return the executable output path override or TOML output path."""
if getattr(args, "out", None) is not None:
return str(args.out)
return cfg.output.path
[docs]
def apply_runtime_path_overrides(
cfg: RuntimeConfig,
args: Any,
*,
resolve_runtime_path: Callable[..., str | None],
) -> RuntimeConfig:
"""Apply cwd-resolved executable path overrides for geometry and init files."""
from dataclasses import replace
cwd = Path.cwd()
geometry = cfg.geometry
vmec_cli = getattr(args, "vmec_file", None)
geom_cli = getattr(args, "geometry_file", None)
if vmec_cli is not None:
geometry = replace(
geometry, vmec_file=resolve_runtime_path(str(vmec_cli), base_dir=cwd)
)
if geom_cli is not None:
geometry = replace(
geometry, geometry_file=resolve_runtime_path(str(geom_cli), base_dir=cwd)
)
init = cfg.init
init_cli = getattr(args, "init_file", None)
if init_cli is not None:
init = replace(
init, init_file=resolve_runtime_path(str(init_cli), base_dir=cwd)
)
return replace(cfg, geometry=geometry, init=init)
[docs]
def apply_quasilinear_overrides(cfg: RuntimeConfig, args: Any) -> RuntimeConfig:
"""Apply executable quasilinear diagnostic overrides."""
from dataclasses import replace
ql = cfg.quasilinear
updates: dict[str, object] = {}
if getattr(args, "quasilinear", False):
updates["enabled"] = True
mapping = {
"ql_mode": "mode",
"ql_saturation_rule": "saturation_rule",
"ql_csat": "csat",
"ql_normalization": "amplitude_normalization",
"ql_output": "output_path",
}
for arg_name, field_name in mapping.items():
value = getattr(args, arg_name, None)
if value is not None:
updates[field_name] = value
if not updates:
return cfg
return replace(cfg, quasilinear=replace(ql, **cast(Any, updates)))
def _status_printer(prefix: str) -> Callable[[str], None]:
def _emit(message: str) -> None:
print(f"{prefix}: {message}", flush=True)
return _emit
[docs]
def should_write_plots(args: Any, cfg: RuntimeConfig) -> bool:
"""Resolve whether a finished run draws its own figures.
``--no-plots`` overrides the configuration, which itself defaults to on.
"""
if bool(getattr(args, "no_plots", False)):
return False
return bool(getattr(cfg.output, "plots", True))
[docs]
def auto_plot_runtime_outputs(
kind: str,
paths: Mapping[str, str],
*,
enabled: bool,
) -> list[str]:
"""Draw and announce the figure set for a run that has finished writing.
The whole call is contained. The simulation is already on disk by the time
this runs, so a headless display, an unusable matplotlib backend, or an
import that fails outright must cost the pictures and nothing else -- the
caller's exit status never depends on what happens here.
"""
if not enabled or not paths:
return []
try:
from gkx.artifacts.run_figures import auto_plot_saved_run
written = auto_plot_saved_run(kind, paths)
except Exception as exc: # pragma: no cover - guards a broken plotting stack
print(f"warning: could not plot {kind} outputs: {exc}", flush=True)
return []
for figure in written:
print(f"saved {figure}")
return written
def _write_linear_runtime_command_outputs(
args: Any,
cfg: RuntimeConfig,
result: RuntimeLinearResult,
) -> dict[str, dict[str, str]]:
"""Write all optional artifacts produced by one linear runtime command."""
return write_linear_runtime_command_outputs(
linear_out_path=runtime_output_path(args, cfg),
quasilinear_out_path=(
getattr(args, "ql_output", None) or cfg.quasilinear.output_path
),
result=result,
linear_writer=write_runtime_linear_artifacts,
quasilinear_writer=write_quasilinear_artifacts,
)
def _write_scan_runtime_command_outputs(
args: Any,
cfg: RuntimeConfig,
scan: Any,
) -> dict[str, str]:
"""Write optional artifacts produced by one linear-scan runtime command."""
return write_scan_runtime_command_outputs(
runtime_output_path(args, cfg) or cfg.quasilinear.output_path,
scan,
writer=write_runtime_linear_scan_artifacts,
)
[docs]
def plot_saved_output_command(
argv: Sequence[str],
*,
plot_saved_output: Callable[..., Path],
) -> int:
"""Render a saved runtime artifact from the top-level ``--plot`` command."""
if len(argv) < 2:
print(
"usage: gkx plot OUTPUT_FILE [--out FIGURE.png]; legacy usage: gkx --plot OUTPUT_FILE [--out FIGURE.png]"
)
return 1
input_path = argv[1]
out_path = None
if len(argv) > 2:
if len(argv) == 4 and argv[2] == "--out":
out_path = argv[3]
else:
print(
"usage: gkx plot OUTPUT_FILE [--out FIGURE.png]; legacy usage: gkx --plot OUTPUT_FILE [--out FIGURE.png]"
)
return 1
rendered = plot_saved_output(input_path, out=out_path)
print(f"saved {rendered}")
from gkx.artifacts.run_figures import replot_nonlinear_bundle
# On GKX's own nonlinear bundle, rebuild the run's whole figure set.
for figure in replot_nonlinear_bundle(input_path):
print(f"saved {figure}")
return 0
[docs]
def run_runtime_linear_command(args: Any) -> int:
"""Execute the runtime-linear subcommand after parser dispatch."""
cfg, data = _prepare_runtime_command_config(
args,
path_overrides=True,
quasilinear_overrides=True,
)
fit_cfg = _runtime_fit_config(data)
run_cfg = _with_fit_signal(data.get("run", {}), fit_cfg)
opts = _resolve_linear_command_options(args, cfg, run_cfg)
print_linear_run_header(
label="runtime linear run",
config_path=str(args.config),
ky=opts.ky,
Nl=opts.Nl,
Nm=opts.Nm,
solver=opts.solver,
method=opts.method_for_header,
dt=opts.dt_for_header,
steps=opts.steps_for_header,
grid_shape=(int(cfg.grid.Nx), int(cfg.grid.Ny), int(cfg.grid.Nz)),
show_progress=opts.show_progress,
extra=(
f"model={cfg.physics.reduced_model} electrostatic={cfg.physics.electrostatic} "
f"electromagnetic={cfg.physics.electromagnetic} fit_signal={opts.fit_signal}"
),
)
kx = run_cfg.get("kx")
res = run_runtime_linear(
cfg,
ky_target=opts.ky,
kx_target=None if kx is None else float(kx),
Nl=opts.Nl,
Nm=opts.Nm,
solver=opts.solver,
method=opts.method,
dt=opts.dt,
steps=opts.steps,
sample_stride=opts.sample_stride,
fit_signal=opts.fit_signal,
show_progress=opts.show_progress,
status_callback=_status_printer("runtime"),
**fit_cfg,
)
gamma_text = _format_fitted_value(res.gamma, getattr(res, "gamma_stderr", None), 6)
omega_text = _format_fitted_value(res.omega, getattr(res, "omega_stderr", None), 6)
print(f"ky={res.ky:.4f} gamma={gamma_text} omega={omega_text}")
written = _write_linear_runtime_command_outputs(args, cfg, res)
auto_plot_runtime_outputs(
"linear",
written.get("linear", {}),
enabled=should_write_plots(args, cfg),
)
return 0
[docs]
def scan_runtime_linear_command(args: Any) -> int:
"""Execute the runtime-linear ky-scan subcommand after parser dispatch."""
cfg, data = _prepare_runtime_command_config(
args,
path_overrides=False,
quasilinear_overrides=True,
)
fit_cfg = _runtime_fit_config(data)
scan_cfg = _with_fit_signal(data.get("scan", {}), fit_cfg)
opts = _resolve_scan_command_options(args, cfg, scan_cfg)
scan = run_runtime_scan(
cfg,
list(opts.ky_values),
Nl=opts.Nl,
Nm=opts.Nm,
solver=opts.solver,
method=opts.method,
dt=opts.dt,
steps=opts.steps,
sample_stride=opts.sample_stride,
batch_ky=opts.batch_ky,
fit_signal=opts.fit_signal,
show_progress=opts.show_progress,
workers=opts.workers,
parallel_executor=opts.parallel_executor,
warm_start=opts.warm_start,
**fit_cfg,
)
for ky, g, w in zip(scan.ky, scan.gamma, scan.omega):
print(f"ky={ky:.4f} gamma={g:.6f} omega={w:.6f}")
written = _write_scan_runtime_command_outputs(args, cfg, scan)
auto_plot_runtime_outputs(
"linear_scan",
written,
enabled=should_write_plots(args, cfg),
)
return 0
[docs]
def run_runtime_nonlinear_command(args: Any) -> int:
"""Execute the runtime-nonlinear subcommand after parser dispatch."""
cfg, data = _prepare_runtime_command_config(
args,
path_overrides=True,
quasilinear_overrides=False,
)
run_cfg = data.get("run", {})
opts = _resolve_nonlinear_command_options(args, cfg, run_cfg)
if opts.run_to is not None:
cfg = replace(cfg, time=replace(cfg.time, run_to=opts.run_to))
print_nonlinear_run_header(
config_path=str(args.config),
runtime_config=cfg,
ky=opts.ky,
Nl=opts.Nl,
Nm=opts.Nm,
method=opts.method,
dt=opts.dt,
steps=opts.steps,
grid_shape=(int(cfg.grid.Nx), int(cfg.grid.Ny), int(cfg.grid.Nz)),
diagnostics=opts.diagnostics,
show_progress=opts.show_progress,
)
out_path = runtime_output_path(args, cfg)
kx = run_cfg.get("kx")
result, paths = run_runtime_nonlinear_with_artifacts(
cfg,
out=out_path,
ky_target=opts.ky,
kx_target=None if kx is None else float(kx),
Nl=opts.Nl,
Nm=opts.Nm,
dt=opts.dt,
steps=opts.steps,
method=opts.method,
sample_stride=opts.sample_stride,
diagnostics_stride=opts.diagnostics_stride,
laguerre_mode=opts.laguerre_mode,
diagnostics=opts.diagnostics,
show_progress=opts.show_progress,
status_callback=_status_printer("runtime"),
)
if not print_nonlinear_run_summary(result):
return 0
print_nonlinear_command_outputs(paths, enabled=out_path is not None)
auto_plot_runtime_outputs(
"nonlinear",
paths if out_path is not None else {},
enabled=should_write_plots(args, cfg),
)
return 0
__all__ = [
"RUNTIME_CASE_FIT_KEYS",
"RuntimeLinearCommandOptions",
"RuntimeNonlinearCommandOptions",
"RuntimeScanCommandOptions",
"_arg_or_section",
"_parse_ky_values",
"_prepare_runtime_command_config",
"_resolve_grid_time_options",
"_resolve_linear_command_options",
"_resolve_linear_fit_options",
"_resolve_nonlinear_command_options",
"_resolve_scan_command_options",
"_runtime_fit_config",
"_status_printer",
"apply_quasilinear_overrides",
"apply_runtime_path_overrides",
"attach_preloaded_runtime_config",
"auto_plot_runtime_outputs",
"load_runtime_command_config",
"plot_saved_output_command",
"print_linear_run_header",
"print_nonlinear_run_header",
"print_nonlinear_run_summary",
"run_runtime_linear_command",
"run_runtime_nonlinear_command",
"run_linear_case",
"run_nonlinear_case",
"runtime_output_path",
"scan_runtime_linear_command",
"should_show_progress",
"should_write_plots",
]
# Programmatic TOML case helpers share the same option-resolution owner.
RUNTIME_CASE_FIT_KEYS = {
"auto_window",
"tmin",
"tmax",
"window_fraction",
"min_points",
"start_fraction",
"growth_weight",
"require_positive",
"min_amp_fraction",
"window_method",
"mode_method",
"fit_signal",
}
_CASE_LINEAR_SPECS = (
("ky_target", "run", "ky", 0.3, float),
("kx_target", "run", "kx", None, None),
("Nl", "run", "Nl", _RUNTIME_LINEAR_HL_FALLBACK[0], int),
("Nm", "run", "Nm", _RUNTIME_LINEAR_HL_FALLBACK[1], int),
("solver", "run", "solver", "auto", str),
("method", "run", "method", None, None),
("dt", "run", "dt", None, None),
("steps", "run", "steps", None, None),
("sample_stride", "time", "sample_stride", None, None),
)
_CASE_NONLINEAR_SPECS = (
("ky_target", "run", "ky", 0.3, float),
("kx_target", "run", "kx", None, None),
("Nl", "run", "Nl", _RUNTIME_NONLINEAR_HL_FALLBACK[0], int),
("Nm", "run", "Nm", _RUNTIME_NONLINEAR_HL_FALLBACK[1], int),
("method", "run", "method", None, None),
("dt", "time", "dt", None, None),
("steps", "run", "steps", None, None),
("sample_stride", "time", "sample_stride", None, None),
("diagnostics_stride", "time", "diagnostics_stride", None, None),
)
def _case_run_kwargs(
raw: dict[str, Any],
overrides: Mapping[str, Any],
specs: tuple[tuple[str, str, str, Any, Callable[[Any], Any] | None], ...],
) -> dict[str, Any]:
"""Resolve explicit Python arguments before falling back to TOML/defaults."""
run_cfg = dict(raw.get("run", {}))
time_cfg = dict(raw.get("time", {}))
sections = {"run": run_cfg, "time": time_cfg}
resolved: dict[str, Any] = {}
for output_name, section_name, input_name, default, converter in specs:
value = overrides.get(input_name)
if value is None:
value = sections[section_name].get(input_name, default)
resolved[output_name] = converter(value) if converter is not None else value
return resolved
def _nonlinear_case_run_kwargs(
raw: dict[str, Any],
overrides: Mapping[str, Any],
) -> dict[str, Any]:
"""Resolve keyword arguments for one programmatic nonlinear TOML case."""
resolved = _case_run_kwargs(raw, overrides, _CASE_NONLINEAR_SPECS)
resolved["diagnostics"] = True
return resolved
def _linear_case_run_kwargs(
raw: dict[str, Any],
overrides: Mapping[str, Any],
) -> dict[str, Any]:
"""Resolve keyword arguments for one programmatic linear TOML case."""
return _case_run_kwargs(raw, overrides, _CASE_LINEAR_SPECS)
def _case_status_printer(message: str) -> None:
"""Print programmatic runtime-case progress messages."""
print(f"runtime: {message}")
[docs]
def run_linear_case(
config_path: str | Path,
*,
ky: float | None = None,
Nl: int | None = None,
Nm: int | None = None,
solver: str | None = None,
method: str | None = None,
dt: float | None = None,
steps: int | None = None,
sample_stride: int | None = None,
show_progress: bool = True,
) -> int:
"""Run a linear case from a runtime TOML with optional overrides."""
cfg, raw = load_runtime_from_toml(config_path)
run_kwargs = _linear_case_run_kwargs(
raw,
{
"ky": ky,
"Nl": Nl,
"Nm": Nm,
"solver": solver,
"method": method,
"dt": dt,
"steps": steps,
"sample_stride": sample_stride,
},
)
run_kwargs["show_progress"] = show_progress
result = run_runtime_linear(
cfg,
**run_kwargs,
**_runtime_fit_config(raw),
)
if cfg.output.path:
paths = write_runtime_linear_artifacts(cfg.output.path, result)
if "summary" in paths:
print(f"saved {paths['summary']}")
gamma_text = _format_fitted_value(
result.gamma, getattr(result, "gamma_stderr", None), 8
)
omega_text = _format_fitted_value(
result.omega, getattr(result, "omega_stderr", None), 8
)
print(f"ky={result.ky:.6f} gamma={gamma_text} omega={omega_text}")
return 0
[docs]
def run_nonlinear_case(
config_path: str | Path,
*,
ky: float | None = None,
Nl: int | None = None,
Nm: int | None = None,
method: str | None = None,
dt: float | None = None,
steps: int | None = None,
sample_stride: int | None = None,
diagnostics_stride: int | None = None,
show_progress: bool = True,
) -> int:
"""Run a nonlinear case from a runtime TOML with optional overrides."""
cfg, raw = load_runtime_from_toml(config_path)
run_kwargs = _nonlinear_case_run_kwargs(
raw,
{
"ky": ky,
"Nl": Nl,
"Nm": Nm,
"method": method,
"dt": dt,
"steps": steps,
"sample_stride": sample_stride,
"diagnostics_stride": diagnostics_stride,
},
)
run_kwargs["show_progress"] = show_progress
if cfg.output.path:
result, paths = run_runtime_nonlinear_with_artifacts(
cfg,
out=cfg.output.path,
**run_kwargs,
status_callback=_case_status_printer,
)
if "summary" in paths:
print(f"saved {paths['summary']}")
else:
result = run_runtime_nonlinear(
cfg,
resolved_diagnostics=False,
**run_kwargs,
status_callback=_case_status_printer,
)
if result.diagnostics is None or result.ky_selected is None:
print("completed without streamed diagnostics")
return 0
diag = result.diagnostics
print(
"ky={:.6f} Wg={:.8e} Wphi={:.8e} heat={:.8e} pflux={:.8e}".format(
float(result.ky_selected),
float(diag.Wg_t[-1]),
float(diag.Wphi_t[-1]),
float(diag.heat_flux_t[-1]),
float(diag.particle_flux_t[-1]),
)
)
return 0