"""Runtime ky-scan orchestration and combined-batch scan policy."""
from __future__ import annotations
from dataclasses import dataclass, replace
from dataclasses import asdict
from typing import Any, Callable, Mapping, Sequence
import numpy as np
from gkx.benchmarking_shared import _midplane_index
from gkx.config import RuntimeConfig
from gkx.core_grid import build_spectral_grid, select_ky_grid
from gkx.diagnostics.growth_rates import (
fit_growth_rate,
fit_growth_rate_auto,
fit_growth_rate_auto_with_stats,
fit_growth_rate_with_stats,
)
from gkx.diagnostics.modes import (
ModeSelection,
extract_mode_time_series,
select_ky_index,
)
from gkx.diagnostics.normalization import apply_diagnostic_normalization
from gkx.geometry import apply_geometry_grid_defaults
from gkx.parallel import independent_map
from gkx.solvers_linear_integrators import integrate_linear_diagnostics
from gkx.solvers_linear_krylov import KrylovConfig
from gkx.workflows.linear import _normalize_linear_solver_name, run_runtime_linear
from gkx.workflows.runtime.diagnostics import (
_RuntimeLinearFitOptions,
_fit_signal_key,
ensure_finite_linear_history,
)
from gkx.workflows.runtime.results import (
RuntimeLinearResult,
RuntimeLinearScanResult,
RuntimeParameterScanResult,
)
from gkx.workflows.runtime.warm_start import (
WarmStartPolicy,
carry_state,
linear_scan_warm_start_refusal,
relative_change,
scan_visit_order,
)
from gkx.workflows.runtime.startup import (
_build_initial_condition,
_resolve_runtime_hl_dims,
build_runtime_geometry,
build_runtime_linear_params,
build_runtime_linear_terms,
)
[docs]
@dataclass(frozen=True)
class RuntimeIndependentParallelPlan:
"""Resolved independent-worker policy for runtime scan workloads."""
requested_workers: int
effective_workers: int
executor: str
strategy: str
axis: str
source: str
problem_size: int
@property
def enabled(self) -> bool:
"""Whether the resolved plan uses more than one independent worker."""
return self.effective_workers > 1
[docs]
def to_dict(self) -> dict[str, Any]:
"""Return a JSON-friendly policy payload for runtime artifacts."""
payload = asdict(self)
payload["enabled"] = self.enabled
return payload
def _parallel_requests_combined_ky_scan(cfg: RuntimeConfig) -> bool:
"""Return whether runtime parallel config requests the combined-ky scan path."""
parallel = getattr(cfg, "parallel", None)
if parallel is None:
return False
return (
str(getattr(parallel, "strategy", "serial")).lower() == "combined_ky"
and str(getattr(parallel, "axis", "ky")).lower() == "ky"
)
def _normalize_independent_executor(backend: str, fallback: str) -> str:
backend_key = str(backend).strip().lower().replace("-", "_")
fallback_key = str(fallback).strip().lower().replace("-", "_")
aliases = {
"thread": "thread",
"threads": "thread",
"process": "process",
"processes": "process",
}
if backend_key in {"", "auto"}:
try:
return aliases[fallback_key]
except KeyError as exc:
raise ValueError("parallel_executor must be 'thread' or 'process'") from exc
try:
return aliases[backend_key]
except KeyError as exc:
raise ValueError(
"runtime [parallel] backend for independent scans must be "
"'auto', 'thread', or 'process'"
) from exc
def _runtime_independent_parallel_plan(
cfg: RuntimeConfig,
*,
problem_size: int,
workers: int,
executor: str,
) -> RuntimeIndependentParallelPlan:
"""Resolve independent ``k_y`` worker policy from arguments and config."""
size = int(problem_size)
if size < 0:
raise ValueError("problem_size must be non-negative")
requested = int(workers)
if requested < 1:
raise ValueError("workers must be >= 1")
executor_key = _normalize_independent_executor("auto", executor)
source = "arguments"
strategy = "serial"
axis = "ky"
parallel = getattr(cfg, "parallel", None)
if parallel is not None:
strategy = str(getattr(parallel, "strategy", "serial")).strip().lower()
axis = str(getattr(parallel, "axis", "ky")).strip().lower()
if requested == 1 and strategy == "batch":
if axis != "ky":
raise ValueError(
"runtime [parallel] strategy='batch' is supported only for axis='ky'"
)
configured_workers = getattr(parallel, "num_devices", None)
if configured_workers is None:
configured_workers = getattr(parallel, "batch_size", None)
requested = max(int(configured_workers or 1), 1)
executor_key = _normalize_independent_executor(
str(getattr(parallel, "backend", "auto")), executor
)
source = "runtime_config"
effective = 0 if size == 0 else min(requested, size)
return RuntimeIndependentParallelPlan(
requested_workers=requested,
effective_workers=effective,
executor=executor_key,
strategy=strategy,
axis=axis,
source=source,
problem_size=size,
)
@dataclass(frozen=True)
class _BatchScanSetup:
geom: Any
grid: Any
params: Any
terms: Any
ky_indices: np.ndarray
nspecies: int
@dataclass(frozen=True)
class _BatchDiagnostics:
phi_t: np.ndarray
density_t: np.ndarray
time: np.ndarray
@dataclass(frozen=True)
class _RuntimeScanOptions(_RuntimeLinearFitOptions):
method: str | None
dt: float | None
steps: int | None
sample_stride: int | None
fit_signal: str
show_progress: bool
def task_fields(self) -> dict[str, Any]:
return dict(vars(self))
[docs]
def run_runtime_scan_ky_task(task: dict[str, Any]) -> RuntimeLinearResult:
"""Run one independent ky point for ordered scan-worker execution."""
return run_runtime_linear(
task["cfg"],
ky_target=float(task["ky"]),
Nl=int(task["Nl"]),
Nm=int(task["Nm"]),
solver=str(task["solver"]),
method=task["method"],
dt=task["dt"],
steps=task["steps"],
sample_stride=task["sample_stride"],
auto_window=bool(task["auto_window"]),
tmin=task["tmin"],
tmax=task["tmax"],
window_fraction=float(task["window_fraction"]),
min_points=int(task["min_points"]),
start_fraction=float(task["start_fraction"]),
growth_weight=float(task["growth_weight"]),
require_positive=bool(task["require_positive"]),
min_amp_fraction=float(task["min_amp_fraction"]),
window_method=str(task.get("window_method", "stationary")),
krylov_cfg=task["krylov_cfg"],
mode_method=str(task["mode_method"]),
fit_signal=str(task["fit_signal"]),
show_progress=bool(task["show_progress"]),
initial_state=task.get("initial_state"),
return_state=bool(task.get("return_state", False)),
)
def _combined_ky_scan_requested(
*,
cfg: RuntimeConfig,
batch_ky: bool,
solver_key: str,
) -> bool:
requested = bool(batch_ky or _parallel_requests_combined_ky_scan(cfg))
if requested and solver_key == "krylov":
raise ValueError("batch_ky is only supported for time integration")
if requested and bool(getattr(cfg.quasilinear, "enabled", False)):
raise NotImplementedError(
"quasilinear scan artifacts currently require serial ky evaluation"
)
return requested
# ``[parallel] strategy`` answers two different questions depending on which
# layer reads it, so a per-ky worker cannot simply inherit the scan's copy.
#
# A SCAN-level strategy says how the ky points are dispatched, and the scan
# honours it before any worker starts: ``combined_ky`` routes to the combined-ky
# solver, and the rest resolve through ``runtime_independent_parallel_plan`` into
# an independent worker pool. Each worker then runs an ordinary single-ky solve,
# so forwarding the strategy sent that inner solve down the sharded linear-RHS
# path, which implements only ``strategy="velocity"``. That is how the shipped
# examples/11_parallel_scan deck died in the worker with NotImplementedError under
# ``strategy="batch"``; ``device_batch``, ``pmap`` and ``pjit`` were measured
# failing in the same worker with the same error, so they are neutralised the
# same way. ``combined_ky`` belongs here too and is not merely defensive: it
# routes away only for ``axis="ky"``, so ``axis="hermite"`` reaches this helper
# and was measured failing here as well.
_SCAN_LEVEL_WORKER_STRATEGIES = frozenset(
{"batch", "combined_ky", "device_batch", "pjit", "pmap"}
)
# A SOLVER-level strategy says how one solve is decomposed, so it has to reach
# the inner solve untouched. ``velocity`` is the one route the linear RHS
# implements. ``state`` and ``shard_map`` name velocity-space decompositions the
# linear RHS has no route for yet -- ``shard_map`` with ``axis="species_hermite"``
# is the production nonlinear route -- so forwarding them makes the solver refuse
# on its own terms, which is the honest answer. Neutralising them would answer a
# sharding request with an unsharded run and report success.
_SOLVER_LEVEL_WORKER_STRATEGIES = frozenset(
{"serial", "shard_map", "state", "velocity"}
)
def _worker_parallel_strategy(strategy: str) -> str:
"""Return the ``[parallel] strategy`` one per-ky worker may run with.
Forwarding is a whitelist, not a blacklist, because neither default is safe
for a strategy nobody has classified: forwarding an unclassified scan-level
strategy repeats the bug above, and neutralising an unclassified
solver-level one silently drops the decomposition it asked for. So an
unclassified strategy raises here rather than picking one of the two.
"""
if strategy in _SOLVER_LEVEL_WORKER_STRATEGIES:
return strategy
if strategy in _SCAN_LEVEL_WORKER_STRATEGIES:
return "serial"
raise NotImplementedError(
f"[parallel] strategy='{strategy}' is not classified as scan-level or "
"solver-level for independent ky workers; add it to "
"_SCAN_LEVEL_WORKER_STRATEGIES (the scan dispatches it and the worker "
"runs an ordinary single-ky solve) or to "
"_SOLVER_LEVEL_WORKER_STRATEGIES (the worker's own solve honours it)"
)
def _scan_worker_tasks(
cfg: RuntimeConfig,
ky_arr: np.ndarray,
*,
Nl: int,
Nm: int,
solver: str,
krylov_cfg: Any,
options: _RuntimeScanOptions,
) -> list[dict[str, Any]]:
strategy = str(getattr(cfg.parallel, "strategy", "serial"))
worker_strategy = _worker_parallel_strategy(strategy)
worker_cfg = (
cfg
if worker_strategy == strategy
else replace(cfg, parallel=replace(cfg.parallel, strategy=worker_strategy))
)
return [
{
"cfg": worker_cfg,
"ky": float(ky),
"Nl": Nl,
"Nm": Nm,
"solver": solver,
"krylov_cfg": krylov_cfg,
**options.task_fields(),
}
for ky in ky_arr
]
def _scan_result_from_worker_results(
ky_arr: np.ndarray,
results: list[Any],
parallel_plan: Any,
) -> RuntimeLinearScanResult:
gamma = np.zeros_like(ky_arr)
omega = np.zeros_like(ky_arr)
ql_payloads: list[dict[str, Any]] = []
for i, res in enumerate(results):
gamma[i] = float(res.gamma)
omega[i] = float(res.omega)
if res.quasilinear is not None:
ql_payloads.append(res.quasilinear)
parallel_payload = parallel_plan.to_dict()
parallel_payload.update(
{
"identity_contract": "independent ky workers must preserve serial ky ordering and values",
"quasilinear_state_extraction": bool(ql_payloads),
}
)
return RuntimeLinearScanResult(
ky=ky_arr,
gamma=gamma,
omega=omega,
quasilinear=tuple(ql_payloads) if ql_payloads else None,
parallel=parallel_payload,
)
def _run_combined_ky_scan(
cfg: RuntimeConfig,
ky_arr: np.ndarray,
*,
Nl: int,
Nm: int,
options: _RuntimeScanOptions,
) -> RuntimeLinearScanResult:
return run_runtime_scan_batch(
cfg,
ky_arr,
Nl=Nl,
Nm=Nm,
**options.task_fields(),
)
def _run_warm_serial_ky_scan(
ky_arr: np.ndarray,
tasks: list[dict[str, Any]],
parallel_plan: Any,
) -> RuntimeLinearScanResult:
"""Walk a ky scan in neighbour order, seeding each point from the last one.
The converged state of one ky is an excellent guess at the next, so the
scan is visited in monotone ky order and each point starts from its
neighbour's answer instead of the cold initial condition. Results are
written back into the requested ky order, so the reported arrays are
identical in shape and ordering to the cold scan. A point whose state comes
back unusable simply seeds the next point cold rather than poisoning it.
"""
order = scan_visit_order(ky_arr)
results: list[Any] = [None] * int(ky_arr.size)
carried: np.ndarray | None = None
warm_points = 0
for position in order:
index = int(position)
result = run_runtime_scan_ky_task(
{**tasks[index], "initial_state": carried, "return_state": True}
)
results[index] = result
warm_points += int(carried is not None)
carried = carry_state(result.state)
scan = _scan_result_from_worker_results(ky_arr, results, parallel_plan)
return replace(
scan,
warm_start={
"enabled": True,
"visit_order": [int(value) for value in order],
"warm_points": int(warm_points),
"cold_points": int(ky_arr.size) - int(warm_points),
},
)
def _run_independent_ky_scan(
cfg: RuntimeConfig,
ky_arr: np.ndarray,
*,
Nl: int,
Nm: int,
solver: str,
krylov_cfg: Any,
options: _RuntimeScanOptions,
workers: int,
parallel_executor: str,
warm_policy: WarmStartPolicy,
solver_key: str,
) -> RuntimeLinearScanResult:
tasks = _scan_worker_tasks(
cfg,
ky_arr,
Nl=Nl,
Nm=Nm,
solver=solver,
krylov_cfg=krylov_cfg,
options=options,
)
parallel_plan = _runtime_independent_parallel_plan(
cfg, problem_size=int(ky_arr.size), workers=workers, executor=parallel_executor
)
# The plan is resolved first because it can promote the worker count from
# ``[parallel]``: warm start yields to any request for real independence
# rather than quietly serializing it.
refusal = linear_scan_warm_start_refusal(
policy=warm_policy,
solver_key=solver_key,
workers=parallel_plan.requested_workers,
)
if refusal is None and ky_arr.size > 1:
return _run_warm_serial_ky_scan(ky_arr, tasks, parallel_plan)
results = independent_map(
run_runtime_scan_ky_task,
tasks,
workers=parallel_plan.requested_workers,
executor=parallel_plan.executor,
)
return _scan_result_from_worker_results(ky_arr, results, parallel_plan)
def _batch_scan_setup(
cfg: RuntimeConfig,
ky_arr: np.ndarray,
*,
Nl: int,
Nm: int,
) -> _BatchScanSetup:
geom = build_runtime_geometry(cfg)
grid_cfg = apply_geometry_grid_defaults(geom, cfg.grid)
full_grid = build_spectral_grid(grid_cfg)
full_ky_indices = np.asarray(
[select_ky_index(np.asarray(full_grid.ky), ky) for ky in ky_arr],
dtype=int,
)
# Linear scans retain every requested mode. The parent mask is only for
# nonlinear convolution dealiasing and would otherwise suppress high ky.
grid = select_ky_grid(full_grid, full_ky_indices)
params = build_runtime_linear_params(cfg, Nm=Nm, geom=geom)
terms = build_runtime_linear_terms(cfg)
ky_indices = np.arange(ky_arr.size, dtype=int)
nspecies = max(len([s for s in cfg.species if s.kinetic]), 1)
return _BatchScanSetup(
geom=geom,
grid=grid,
params=params,
terms=terms,
ky_indices=ky_indices,
nspecies=nspecies,
)
def _combined_batch_initial_condition(
cfg: RuntimeConfig,
setup: _BatchScanSetup,
*,
Nl: int,
Nm: int,
) -> Any:
g0 = None
for ky_idx in setup.ky_indices:
g0_local = _build_initial_condition(
setup.grid,
setup.geom,
cfg,
ky_index=int(ky_idx),
kx_index=0,
Nl=Nl,
Nm=Nm,
nspecies=setup.nspecies,
)
g0 = g0_local if g0 is None else g0 + g0_local
if g0 is None:
raise ValueError("No ky values provided for batch scan")
return g0
def _batch_time_config(
cfg: RuntimeConfig,
*,
options: _RuntimeScanOptions,
) -> Any:
tcfg = cfg.time
if options.method is not None:
tcfg = replace(tcfg, method=str(options.method))
if options.dt is not None:
tcfg = replace(tcfg, dt=float(options.dt))
if options.steps is not None:
tcfg = replace(tcfg, t_max=float(options.steps) * float(tcfg.dt))
if options.sample_stride is not None:
tcfg = replace(tcfg, sample_stride=int(options.sample_stride))
return tcfg
def _run_batch_diagnostics(
setup: _BatchScanSetup,
g0: Any,
tcfg: Any,
*,
show_progress: bool,
) -> _BatchDiagnostics:
from gkx.solvers_time_runners import _resolve_config_collision_operator
steps_val = int(round(tcfg.t_max / tcfg.dt))
# The runtime linear scan is the path the executable takes, so the TOML
# collision_operator selection has to be resolved here too.
collision_operator = _resolve_config_collision_operator(tcfg, setup.params, g0)
diag = integrate_linear_diagnostics(
g0,
setup.grid,
setup.geom,
setup.params,
dt=tcfg.dt,
steps=steps_val,
method=tcfg.method,
terms=setup.terms,
sample_stride=tcfg.sample_stride,
species_index=0,
record_hl_energy=False,
show_progress=show_progress,
collision_operator=collision_operator,
)
phi_t_np = np.asarray(diag[1])
dens_t_np = np.asarray(diag[2])
ensure_finite_linear_history("field", phi_t_np)
ensure_finite_linear_history("density", dens_t_np)
t_arr = (
float(tcfg.dt)
* float(tcfg.sample_stride)
* (np.arange(phi_t_np.shape[0], dtype=float) + 1.0)
)
return _BatchDiagnostics(phi_t=phi_t_np, density_t=dens_t_np, time=t_arr)
def _auto_fit_scan_candidate(
signal: Any,
time: np.ndarray,
*,
options: _RuntimeScanOptions,
) -> tuple[float, float, float]:
"""Fit and score one channel of a combined-ky scan point.
``auto`` selects the diagnostic channel, never the fit window: with
``auto_window`` off, both channels are fitted over the requested window so
the winner reports what an explicit ``fit_signal`` request would.
"""
if options.auto_window:
gamma, omega, _tmin, _tmax, r2, r2_phase = fit_growth_rate_auto_with_stats(
time,
signal,
window_fraction=options.window_fraction,
min_points=options.min_points,
start_fraction=options.start_fraction,
growth_weight=options.growth_weight,
require_positive=options.require_positive,
min_amp_fraction=options.min_amp_fraction,
window_method=options.window_method,
)
else:
gamma, omega, r2, r2_phase = fit_growth_rate_with_stats(
time,
signal,
tmin=options.tmin,
tmax=options.tmax,
)
return gamma, omega, r2 + 0.2 * r2_phase + options.growth_weight * gamma
def _auto_fit_scan_point(
diagnostics: _BatchDiagnostics,
sel: ModeSelection,
*,
options: _RuntimeScanOptions,
) -> tuple[float, float]:
gamma_phi, omega_phi, score_phi = _auto_fit_scan_candidate(
extract_mode_time_series(diagnostics.phi_t, sel, method=options.mode_method),
diagnostics.time,
options=options,
)
gamma_den, omega_den, score_den = _auto_fit_scan_candidate(
extract_mode_time_series(
diagnostics.density_t, sel, method=options.mode_method
),
diagnostics.time,
options=options,
)
return (gamma_phi, omega_phi) if score_phi >= score_den else (gamma_den, omega_den)
def _fit_batch_scan_point(
diagnostics: _BatchDiagnostics,
sel: ModeSelection,
*,
fit_key: str,
options: _RuntimeScanOptions,
) -> tuple[float, float]:
if fit_key == "auto":
return _auto_fit_scan_point(
diagnostics,
sel,
options=options,
)
signal = extract_mode_time_series(
diagnostics.density_t if fit_key == "density" else diagnostics.phi_t,
sel,
method=options.mode_method,
)
if options.auto_window:
g_val, o_val, _tmin, _tmax = fit_growth_rate_auto(
diagnostics.time,
signal,
window_fraction=options.window_fraction,
min_points=options.min_points,
start_fraction=options.start_fraction,
growth_weight=options.growth_weight,
require_positive=options.require_positive,
min_amp_fraction=options.min_amp_fraction,
window_method=options.window_method,
)
return g_val, o_val
return fit_growth_rate(
diagnostics.time, signal, tmin=options.tmin, tmax=options.tmax
)
def _fit_batch_scan_result(
cfg: RuntimeConfig,
ky_arr: np.ndarray,
setup: _BatchScanSetup,
diagnostics: _BatchDiagnostics,
*,
options: _RuntimeScanOptions,
) -> RuntimeLinearScanResult:
"""Fit each requested ky from one combined-ky diagnostic time history."""
gamma = np.zeros_like(ky_arr, dtype=float)
omega = np.zeros_like(ky_arr, dtype=float)
fit_key = _fit_signal_key(options.fit_signal)
for i, ky_idx in enumerate(setup.ky_indices):
sel = ModeSelection(
ky_index=int(ky_idx), kx_index=0, z_index=_midplane_index(setup.grid)
)
g_val, o_val = _fit_batch_scan_point(
diagnostics,
sel,
fit_key=fit_key,
options=options,
)
gamma[i], omega[i] = apply_diagnostic_normalization(
g_val,
o_val,
rho_star=float(np.asarray(setup.params.rho_star)),
diagnostic_norm=cfg.normalization.diagnostic_norm,
)
return RuntimeLinearScanResult(ky=ky_arr, gamma=gamma, omega=omega)
[docs]
def run_runtime_scan(
cfg: RuntimeConfig,
ky_values: Sequence[float],
*,
Nl: int | None = None,
Nm: int | None = None,
solver: str = "auto",
method: str | None = None,
dt: float | None = None,
steps: int | None = None,
sample_stride: int | None = None,
batch_ky: bool = False,
auto_window: bool = True,
tmin: float | None = None,
tmax: float | None = None,
window_fraction: float = 0.4,
min_points: int = 40,
start_fraction: float = 0.2,
growth_weight: float = 0.2,
require_positive: bool = True,
min_amp_fraction: float = 0.0,
window_method: str = "stationary",
krylov_cfg: KrylovConfig | None = None,
mode_method: str = "project",
fit_signal: str = "auto",
show_progress: bool = False,
workers: int = 1,
parallel_executor: str = "thread",
warm_start: bool | None = None,
) -> RuntimeLinearScanResult:
"""Run a ky scan; ``warm_start`` overrides the case output policy.
Coordinates serial, independent-worker, or combined-ky runtime scans.
"""
ky_arr = np.asarray(ky_values, dtype=float)
cfg, Nl_use, Nm_use = _resolve_runtime_hl_dims(cfg, Nl=Nl, Nm=Nm)
solver_key = _normalize_linear_solver_name(solver)
options = _RuntimeScanOptions(
method=method,
dt=dt,
steps=steps,
sample_stride=sample_stride,
auto_window=auto_window,
tmin=tmin,
tmax=tmax,
window_fraction=window_fraction,
min_points=min_points,
start_fraction=start_fraction,
growth_weight=growth_weight,
require_positive=require_positive,
min_amp_fraction=min_amp_fraction,
window_method=window_method,
mode_method=mode_method,
fit_signal=fit_signal,
show_progress=show_progress,
)
if _combined_ky_scan_requested(cfg=cfg, batch_ky=batch_ky, solver_key=solver_key):
return _run_combined_ky_scan(
cfg,
ky_arr,
Nl=Nl_use,
Nm=Nm_use,
options=options,
)
return _run_independent_ky_scan(
cfg,
ky_arr,
Nl=Nl_use,
Nm=Nm_use,
solver=solver,
krylov_cfg=krylov_cfg,
options=options,
workers=workers,
parallel_executor=parallel_executor,
warm_policy=WarmStartPolicy.from_config(cfg, override=warm_start),
solver_key=solver_key,
)
[docs]
def run_runtime_scan_batch(
cfg: RuntimeConfig,
ky_arr: np.ndarray,
*,
Nl: int,
Nm: int,
method: str | None,
dt: float | None,
steps: int | None,
sample_stride: int | None,
auto_window: bool,
tmin: float | None,
tmax: float | None,
window_fraction: float,
min_points: int,
start_fraction: float,
growth_weight: float,
require_positive: bool,
min_amp_fraction: float,
mode_method: str,
fit_signal: str,
show_progress: bool,
window_method: str = "stationary",
) -> RuntimeLinearScanResult:
"""Batch a ky scan using one time integration over the full grid."""
if ky_arr.size == 0:
raise ValueError("ky_values must not be empty")
setup = _batch_scan_setup(cfg, ky_arr, Nl=Nl, Nm=Nm)
g0 = _combined_batch_initial_condition(cfg, setup, Nl=Nl, Nm=Nm)
options = _RuntimeScanOptions(
method=method,
dt=dt,
steps=steps,
sample_stride=sample_stride,
auto_window=auto_window,
tmin=tmin,
tmax=tmax,
window_fraction=window_fraction,
min_points=min_points,
start_fraction=start_fraction,
growth_weight=growth_weight,
require_positive=require_positive,
min_amp_fraction=min_amp_fraction,
window_method=window_method,
mode_method=mode_method,
fit_signal=fit_signal,
show_progress=show_progress,
)
tcfg = _batch_time_config(cfg, options=options)
diagnostics = _run_batch_diagnostics(
setup,
g0,
tcfg,
show_progress=options.show_progress,
)
return _fit_batch_scan_result(cfg, ky_arr, setup, diagnostics, options=options)
[docs]
def run_runtime_parameter_scan(
cfg: RuntimeConfig,
parameter_values: Sequence[float],
*,
parameter_name: str,
update_config: Callable[[RuntimeConfig, float, int], RuntimeConfig],
ky_target: float = 0.3,
linear_options: Mapping[str, Any] | None = None,
point_options: Callable[[float, int, RuntimeLinearResult | None], Mapping[str, Any]]
| None = None,
candidate_options: Callable[
[float, int, RuntimeLinearResult | None], Sequence[Mapping[str, Any]]
]
| None = None,
select_candidate: Callable[
[float, int, tuple[RuntimeLinearResult, ...], RuntimeLinearResult | None], int
]
| None = None,
continuation: bool = False,
warm_start: bool = False,
warm_start_max_step: float = 0.25,
) -> RuntimeParameterScanResult:
"""Run a scalar scan and optionally continue a selected solution branch.
``continuation`` follows one branch and requires every point to hand back a
state. ``warm_start`` is the weaker, opportunistic form: the previous
point's state is reused only while the parameter moves by no more than
``warm_start_max_step`` relative to the point before it, and only when that
state is finite and non-degenerate. A larger step, or an unusable state,
silently restores the cold initial condition rather than seeding the solver
from somewhere it has no reason to be. Both compile once and reuse the same
executable for every point, because the shapes never change.
"""
name = str(parameter_name).strip()
if not name:
raise ValueError("parameter_name must be nonempty")
values = np.asarray(parameter_values, dtype=float)
if values.ndim != 1 or values.size == 0:
raise ValueError("parameter_values must be a nonempty one-dimensional array")
shared = dict(linear_options or {})
runs: list[RuntimeLinearResult] = []
previous: RuntimeLinearResult | None = None
for index, value in enumerate(values):
point_cfg = update_config(cfg, float(value), index)
if not isinstance(point_cfg, RuntimeConfig):
raise TypeError("update_config must return RuntimeConfig")
options = dict(shared)
if point_options is not None:
options.update(point_options(float(value), index, previous))
if continuation or warm_start or candidate_options is not None:
options["return_state"] = True
if continuation and previous is not None:
if previous.state is None:
raise ValueError("continuation requires each point to return state")
options["initial_state"] = previous.state
elif warm_start and previous is not None:
step = relative_change([values[index - 1]], [value])
seed = (
carry_state(previous.state)
if step <= float(warm_start_max_step)
else None
)
if seed is not None:
options["initial_state"] = seed
overrides = (
tuple(candidate_options(float(value), index, previous))
if candidate_options is not None
else ({},)
)
if not overrides:
raise ValueError("candidate_options must return at least one candidate")
candidates = tuple(
run_runtime_linear(point_cfg, ky_target=ky_target, **(options | dict(item)))
for item in overrides
)
selected = (
int(select_candidate(float(value), index, candidates, previous))
if select_candidate is not None
else 0
)
if selected < 0 or selected >= len(candidates):
raise IndexError("select_candidate returned an out-of-range index")
previous = candidates[selected]
runs.append(previous)
return RuntimeParameterScanResult(
parameter_name=name,
values=values,
gamma=np.asarray([result.gamma for result in runs], dtype=float),
omega=np.asarray([result.omega for result in runs], dtype=float),
runs=tuple(runs),
)
__all__ = [
"RuntimeIndependentParallelPlan",
"run_runtime_scan_ky_task",
"run_runtime_scan_batch",
"run_runtime_scan",
"run_runtime_parameter_scan",
]