"""Runtime executable command workflows."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Mapping, Sequence, cast
from gkx.workflows.runtime.config import RuntimeConfig
from gkx.workflows.runtime.results import (
RuntimeLinearResult,
RuntimeNonlinearResult,
)
from gkx.workflows.runtime.orchestration_artifacts import (
print_linear_run_header,
print_nonlinear_command_outputs,
print_nonlinear_run_header,
print_nonlinear_run_summary,
write_linear_runtime_command_outputs,
write_scan_runtime_command_outputs,
)
RUNTIME_COMMAND_FIT_KEYS = {
"auto_window",
"tmin",
"tmax",
"window_fraction",
"min_points",
"start_fraction",
"growth_weight",
"require_positive",
"min_amp_fraction",
"mode_method",
}
[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
[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
[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
) -> tuple[int, int, str | None, float | None, int | None, int]:
"""Resolve resolution, optional time controls, and output 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", 24)),
int(_arg_or_section(args, section, "Nm", 12)),
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 fit options supported by runtime executable commands."""
return {
k: v for k, v in data.get("fit", {}).items() if k in RUNTIME_COMMAND_FIT_KEYS
}
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")),
)
[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
)
if steps is not None:
nonlinear_steps: int | None = steps
elif bool(cfg.time.fixed_dt):
nonlinear_steps = int(round(cfg.time.t_max / 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))
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)),
)
_PRELOADED_RUNTIME_CONFIG_ATTR = "_gkx_preloaded_runtime_config"
_PRELOADED_RUNTIME_DATA_ATTR = "_gkx_preloaded_runtime_data"
[docs]
@dataclass(frozen=True)
class RuntimeCommandDeps:
"""Patchable dependencies for executable runtime subcommands."""
load_runtime_from_toml: Callable[[str | Path], tuple[RuntimeConfig, dict[str, Any]]]
run_runtime_linear: Callable[..., RuntimeLinearResult]
run_runtime_scan: Callable[..., Any]
run_runtime_nonlinear_with_artifacts: Callable[
..., tuple[RuntimeNonlinearResult, dict[str, str]]
]
write_runtime_linear_artifacts: Callable[
[str | Path, RuntimeLinearResult], dict[str, str]
]
write_runtime_linear_scan_artifacts: Callable[[str | Path, Any], dict[str, str]]
write_quasilinear_artifacts: Callable[[str | Path, dict[str, Any]], dict[str, str]]
resolve_runtime_path: Callable[..., str | None]
[docs]
def build_runtime_command_deps(facade: Any) -> RuntimeCommandDeps:
"""Build runtime command dependencies from a patchable executable facade."""
return RuntimeCommandDeps(
load_runtime_from_toml=facade.load_runtime_from_toml,
run_runtime_linear=facade.run_runtime_linear,
run_runtime_scan=facade.run_runtime_scan,
run_runtime_nonlinear_with_artifacts=facade.run_runtime_nonlinear_with_artifacts,
write_runtime_linear_artifacts=facade.write_runtime_linear_artifacts,
write_runtime_linear_scan_artifacts=facade.write_runtime_linear_scan_artifacts,
write_quasilinear_artifacts=facade.write_quasilinear_artifacts,
resolve_runtime_path=facade.resolve_runtime_path,
)
[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,
*,
deps: RuntimeCommandDeps,
) -> 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 deps.load_runtime_from_toml(args.config)
[docs]
def _prepare_runtime_command_config(
args: Any,
*,
deps: RuntimeCommandDeps,
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, deps=deps)
if path_overrides:
cfg = apply_runtime_path_overrides(
cfg,
args,
resolve_runtime_path=deps.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
def _write_linear_runtime_command_outputs(
args: Any,
cfg: RuntimeConfig,
result: RuntimeLinearResult,
*,
deps: RuntimeCommandDeps,
) -> 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=deps.write_runtime_linear_artifacts,
quasilinear_writer=deps.write_quasilinear_artifacts,
)
def _write_scan_runtime_command_outputs(
args: Any,
cfg: RuntimeConfig,
scan: Any,
*,
deps: RuntimeCommandDeps,
) -> 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=deps.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]")
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]")
return 1
rendered = plot_saved_output(input_path, out=out_path)
print(f"saved {rendered}")
return 0
[docs]
def run_runtime_linear_command(args: Any, *, deps: RuntimeCommandDeps) -> int:
"""Execute the runtime-linear subcommand after parser dispatch."""
cfg, data = _prepare_runtime_command_config(
args,
deps=deps,
path_overrides=True,
quasilinear_overrides=True,
)
run_cfg = data.get("run", {})
fit_cfg = _runtime_fit_config(data)
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}"
),
)
res = deps.run_runtime_linear(
cfg,
ky_target=opts.ky,
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,
)
print(f"ky={res.ky:.4f} gamma={res.gamma:.6f} omega={res.omega:.6f}")
_write_linear_runtime_command_outputs(args, cfg, res, deps=deps)
return 0
[docs]
def scan_runtime_linear_command(args: Any, *, deps: RuntimeCommandDeps) -> int:
"""Execute the runtime-linear ky-scan subcommand after parser dispatch."""
cfg, data = _prepare_runtime_command_config(
args,
deps=deps,
path_overrides=False,
quasilinear_overrides=True,
)
scan_cfg = data.get("scan", {})
fit_cfg = _runtime_fit_config(data)
opts = _resolve_scan_command_options(args, cfg, scan_cfg)
scan = deps.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,
**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}")
_write_scan_runtime_command_outputs(args, cfg, scan, deps=deps)
return 0
[docs]
def run_runtime_nonlinear_command(args: Any, *, deps: RuntimeCommandDeps) -> int:
"""Execute the runtime-nonlinear subcommand after parser dispatch."""
cfg, data = _prepare_runtime_command_config(
args,
deps=deps,
path_overrides=True,
quasilinear_overrides=False,
)
run_cfg = data.get("run", {})
opts = _resolve_nonlinear_command_options(args, cfg, run_cfg)
print_nonlinear_run_header(
config_path=str(args.config),
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)
result, paths = deps.run_runtime_nonlinear_with_artifacts(
cfg,
out=out_path,
ky_target=opts.ky,
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)
return 0
__all__ = [
"RUNTIME_CASE_FIT_KEYS",
"RUNTIME_COMMAND_FIT_KEYS",
"RuntimeCommandDeps",
"RuntimeCaseDeps",
"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",
"build_runtime_command_deps",
"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",
]
# 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",
"mode_method",
"fit_signal",
}
_CASE_LINEAR_SPECS = (
("ky_target", "run", "ky", 0.3, float),
("Nl", "run", "Nl", 24, int),
("Nm", "run", "Nm", 12, 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),
("Nl", "run", "Nl", 4, int),
("Nm", "run", "Nm", 8, 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),
)
[docs]
@dataclass(frozen=True)
class RuntimeCaseDeps:
"""Patchable dependencies for runtime TOML case workflows."""
load_runtime_from_toml: Callable[[str | Path], tuple[RuntimeConfig, dict[str, Any]]]
run_runtime_linear: Callable[..., RuntimeLinearResult]
run_runtime_nonlinear: Callable[..., RuntimeNonlinearResult]
write_runtime_linear_artifacts: Callable[[str | Path, Any], dict[str, str]]
run_runtime_nonlinear_with_artifacts: Callable[
..., tuple[RuntimeNonlinearResult, dict[str, str]]
]
def default_runtime_case_deps() -> RuntimeCaseDeps:
"""Build default executable workflow dependencies."""
from gkx.workflows.runtime.toml import load_runtime_from_toml
from gkx.runtime import run_runtime_linear, run_runtime_nonlinear
from gkx.workflows.runtime.artifacts import (
run_runtime_nonlinear_with_artifacts,
write_runtime_linear_artifacts,
)
return RuntimeCaseDeps(
load_runtime_from_toml=load_runtime_from_toml,
run_runtime_linear=run_runtime_linear,
run_runtime_nonlinear=run_runtime_nonlinear,
write_runtime_linear_artifacts=write_runtime_linear_artifacts,
run_runtime_nonlinear_with_artifacts=run_runtime_nonlinear_with_artifacts,
)
def _runtime_case_fit_config(raw: dict[str, Any]) -> dict[str, Any]:
"""Return fit options accepted by programmatic runtime-case helpers."""
return {k: v for k, v in raw.get("fit", {}).items() if k in RUNTIME_CASE_FIT_KEYS}
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,
deps: RuntimeCaseDeps | None = None,
) -> int:
"""Run a linear case from a runtime TOML with optional overrides."""
case_deps = default_runtime_case_deps() if deps is None else deps
cfg, raw = case_deps.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 = case_deps.run_runtime_linear(
cfg,
**run_kwargs,
**_runtime_case_fit_config(raw),
)
if cfg.output.path:
paths = case_deps.write_runtime_linear_artifacts(cfg.output.path, result)
if "summary" in paths:
print(f"saved {paths['summary']}")
print(f"ky={result.ky:.6f} gamma={result.gamma:.8f} omega={result.omega:.8f}")
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,
deps: RuntimeCaseDeps | None = None,
) -> int:
"""Run a nonlinear case from a runtime TOML with optional overrides."""
case_deps = default_runtime_case_deps() if deps is None else deps
cfg, raw = case_deps.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 = case_deps.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 = case_deps.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