"""Runtime startup and initialization helpers.
This module holds the geometry/loading/initial-condition logic used by the
public runtime entry points. It is intentionally kept separate from the solver
execution layer so startup behavior can be tested and refactored without
touching the time-integration control flow.
"""
from __future__ import annotations
from dataclasses import replace
from typing import Any, Callable, NoReturn, Sequence
import jax.numpy as jnp
from dataclasses import dataclass
from pathlib import Path
from typing import cast
import numpy as np
from gkx.core_velocity import mu_grid_laguerre_modes, mu_grid_nodes
from gkx.geometry import (
FluxTubeGeometryLike,
build_flux_tube_geometry,
ensure_flux_tube_geometry_data,
)
from gkx.operators.linear.params import (
DAMP_ENDS_LAYER_TRANSITS,
LinearParams,
LinearTerms,
_x64_enabled,
linear_terms_to_term_config,
Species,
build_linear_params,
)
from gkx.solvers_linear_krylov import KrylovConfig
from gkx.geometry.miller_eik import generate_runtime_miller_eik
from gkx.diagnostics.normalization import get_normalization_contract
from gkx.artifacts.io import load_netcdf_restart_state
from gkx.config import RuntimeConfig, RuntimeSpeciesConfig
from gkx.terms.config import TermConfig
from gkx.geometry.vmec_eik import generate_runtime_vmec_eik
from gkx.core_grid import SpectralGrid
from gkx.core_ky_layout import (
FULL,
HALF,
KyLayout,
negative_ky_block,
nyc_from_ny,
paired_row_limit,
source_ky_layout,
source_ny_full,
to_full,
to_half,
)
from gkx.operators.linear.cache_arrays import keep_moments
from gkx.operators.linear.cache_builder import build_linear_cache, mask_off_chain_rows
@dataclass(frozen=True)
class _PhiMomentAlgebra:
Jl: np.ndarray
jacobian: np.ndarray
charge: np.ndarray
density: np.ndarray
denom: np.ndarray
tau_e: float
[docs]
def _as_runtime_species_array(
value: float | jnp.ndarray, nspecies: int, name: str
) -> np.ndarray:
"""Return a length-``nspecies`` NumPy array for startup-only algebra."""
arr = np.asarray(value, dtype=float)
if arr.ndim == 0:
arr = arr.reshape(1)
arr = arr.reshape(-1)
if arr.size == 1:
return np.full((int(nspecies),), float(arr[0]), dtype=float)
if arr.size != int(nspecies):
raise ValueError(f"{name} must have length {nspecies} (got {arr.size})")
return arr.astype(float, copy=False)
def _validate_phi_target_species(
species_targets: tuple[int, ...],
*,
nspecies: int,
) -> None:
if not species_targets:
raise ValueError(
"init_field='phi' requires at least one kinetic target species"
)
if min(species_targets) < 0 or max(species_targets) >= nspecies:
raise ValueError("init_field='phi' target species index is out of range")
def _masked_gauge_phi_seed(
phi: np.ndarray,
*,
cache,
ky_i: int,
kx_i: int,
species_targets: tuple[int, ...],
) -> dict[int, np.ndarray] | None:
mask0 = np.asarray(cache.mask0, dtype=bool)
if not np.all(mask0[int(ky_i), int(kx_i), :]):
return None
if np.any(np.abs(phi) > 0.0):
raise ValueError(
"init_field='phi' cannot initialize the masked ky=0, kx=0 gauge mode"
)
return {int(idx): np.zeros_like(phi) for idx in species_targets}
def _phi_moment_algebra(
*,
cache,
params: LinearParams,
nspecies: int,
ky_i: int,
kx_i: int,
) -> _PhiMomentAlgebra:
Jl = np.asarray(cache.Jl, dtype=float)
jacobian = np.asarray(cache.jacobian, dtype=float)
charge = _as_runtime_species_array(params.charge_sign, nspecies, "charge_sign")
density = _as_runtime_species_array(params.density, nspecies, "density")
tz = _as_runtime_species_array(params.tz, nspecies, "tz")
zt = np.where(tz == 0.0, 0.0, 1.0 / tz)
tau_e_arr = np.asarray(params.tau_e, dtype=float).reshape(-1)
tau_e = float(tau_e_arr[0]) if tau_e_arr.size else 0.0
g0_species = np.sum(
Jl[:, :, int(ky_i), int(kx_i), :] * Jl[:, :, int(ky_i), int(kx_i), :], axis=1
)
qneut = np.sum(
density[:, None] * charge[:, None] * zt[:, None] * (1.0 - g0_species), axis=0
)
denom = tau_e + qneut
if not np.all(np.isfinite(denom)):
raise ValueError(
"init_field='phi' produced a non-finite quasineutrality denominator"
)
return _PhiMomentAlgebra(
Jl=Jl,
jacobian=jacobian,
charge=charge,
density=density,
denom=denom,
tau_e=tau_e,
)
def _density_average_for_target_phi(
phi: np.ndarray,
*,
algebra: _PhiMomentAlgebra,
cache,
ky_i: int,
kx_i: int,
) -> np.ndarray:
nbar = algebra.denom.astype(np.complex64) * phi
ky_is_zonal = np.isclose(float(np.asarray(cache.ky)[int(ky_i)]), 0.0)
boltzmann_ions = not np.any(algebra.charge > 0.0)
if algebra.tau_e <= 0.0 or boltzmann_ions or not ky_is_zonal or int(kx_i) <= 0:
return nbar
jac_sum = float(np.sum(algebra.jacobian))
phi_avg = (
np.sum(algebra.jacobian.astype(np.complex64) * phi) / jac_sum
if jac_sum != 0.0
else np.mean(phi)
)
return nbar - np.complex64(algebra.tau_e) * phi_avg
def _project_density_average_to_target_species(
nbar: np.ndarray,
*,
algebra: _PhiMomentAlgebra,
ky_i: int,
kx_i: int,
species_targets: tuple[int, ...],
) -> dict[int, np.ndarray]:
target_indices = np.asarray(species_targets, dtype=int)
coeff = (
algebra.density[target_indices, None]
* algebra.charge[target_indices, None]
* algebra.Jl[target_indices, 0, int(ky_i), int(kx_i), :]
)
coeff_norm = np.sum(coeff * coeff, axis=0)
tiny = 1.0e-30
bad = (coeff_norm <= tiny) & (np.abs(nbar) > tiny)
if np.any(bad):
raise ValueError(
"init_field='phi' cannot be represented by the selected density moments"
)
seed = np.where(
coeff_norm[None, :] > tiny,
(coeff / coeff_norm[None, :]).astype(np.complex64) * nbar[None, :],
np.complex64(0.0),
)
return {
int(s_idx): np.asarray(seed[i], dtype=np.complex64)
for i, s_idx in enumerate(target_indices)
}
[docs]
def _density_moments_for_target_phi(
phi_target: np.ndarray,
*,
cache,
params: LinearParams,
ky_i: int,
kx_i: int,
species_targets: tuple[int, ...],
) -> dict[int, np.ndarray]:
"""Invert the electrostatic field solve for a requested initial ``phi`` mode.
The runtime evolves Hermite-Laguerre moments, not fields. For literature
tests that prescribe an initial electrostatic-potential perturbation, seed
only the density moment with the algebraic moment profile whose immediate
quasineutrality solve returns ``phi_target``. The adiabatic-electron zonal
branch includes the same flux-surface-average correction used in
:mod:`gkx.terms.fields`.
"""
phi = np.asarray(phi_target, dtype=np.complex64)
nspecies = int(np.asarray(cache.Jl).shape[0])
_validate_phi_target_species(species_targets, nspecies=nspecies)
masked_seed = _masked_gauge_phi_seed(
phi,
cache=cache,
ky_i=ky_i,
kx_i=kx_i,
species_targets=species_targets,
)
if masked_seed is not None:
return masked_seed
algebra = _phi_moment_algebra(
cache=cache,
params=params,
nspecies=nspecies,
ky_i=ky_i,
kx_i=kx_i,
)
nbar = _density_average_for_target_phi(
phi,
algebra=algebra,
cache=cache,
ky_i=ky_i,
kx_i=kx_i,
)
return _project_density_average_to_target_species(
nbar,
algebra=algebra,
ky_i=ky_i,
kx_i=kx_i,
species_targets=species_targets,
)
def _species_to_linear(species_cfg: Sequence[RuntimeSpeciesConfig]) -> list[Species]:
kinetic = [s for s in species_cfg if bool(s.kinetic)]
if not kinetic:
raise ValueError(
"RuntimeConfig.species must include at least one kinetic species"
)
return [
Species(
charge=float(s.charge),
mass=float(s.mass),
density=float(s.density),
temperature=float(s.temperature),
tprim=float(s.tprim),
fprim=float(s.fprim),
nu=float(s.nu),
closure=float(s.hermite_closure == "reflectionless"),
)
for s in kinetic
]
def _species_moments(cfg: RuntimeConfig, Nm: int | None) -> tuple | None:
"""Return ``LinearParams.moments`` for the kinetic species, ``None`` if none override."""
kinetic, p = [s for s in cfg.species if s.kinetic], cfg.collisions.p_hyper_m
if all(s.Nl is None and s.Nm is None for s in kinetic):
return None
return tuple(
(
s.Nl,
s.Nm,
_default_hermite_hypercollision_exponent(s.Nm or Nm)
if p is None
else float(p),
)
for s in kinetic
)
def _default_hermite_hypercollision_exponent(nhermite: int | None) -> float:
"""Return the default Hermite hypercollision exponent."""
if nhermite is None:
return 20.0
return float(min(20, max(int(nhermite) // 2, 1)))
def _runtime_model_key(cfg: RuntimeConfig) -> str:
return cfg.physics.reduced_model.strip().lower()
def _raise_unsupported_reduced_model(cfg: RuntimeConfig) -> NoReturn:
"""Fail closed unless the promoted full gyrokinetic runtime is selected."""
raise ValueError(
f"Unknown physics.reduced_model={cfg.physics.reduced_model!r}. "
"Use physics.reduced_model='gyrokinetic' for promoted full-GK workflows."
)
def _runtime_default_krylov_config(cfg: RuntimeConfig) -> KrylovConfig:
"""Return ``auto``: Hermite-banded shift-invert where it fits, else adaptive.
``auto`` runs the banded direct route when its block bands fit in memory and
falls back to the residual-certified ``adaptive`` route if the operator is
not Hermite-banded or the pair misses the original-operator gate. Both gate
on the original operator, so no contract returns an uncertified pair. The
pr3-cm shift-invert defaults of ETG and KBM never converged their inner
solves (2026-10-15), which is why neither is the default.
"""
del cfg
return KrylovConfig()
# Velocity-space fallback for a linear call that names no (Nl, Nm). It was (24, 12),
# Hermite-starved and ~4.5% low in gamma on the shipped Cyclone deck; (12, 24) has the
# same Nl*Nm and puts the budget on the axis that needs it. Certified adaptive
# eigensolves at ky=0.3 against the golden gamma=0.09302951
# (src/gkx/data/cyclone_reference_adiabatic.csv), float32:
#
# Nl Nm gamma rel. to golden
# 8 24 0.09877566 +6.177%
# 8 32 0.09801412 +5.358%
# 24 12 0.08893196 -4.405% <- was
# 12 24 0.09340143 +0.400% <- is
#
# Nl = 8 at either Hermite count is worse than both, so this is a balance. A deck
# whose physics is the other way round (the ETG decks, Nl=24/Nm=8) must still say so.
_RUNTIME_LINEAR_HL_FALLBACK = (12, 24)
# A nonlinear run without Nl/Nm starts coarse; resolution is the deck's call.
_RUNTIME_NONLINEAR_HL_FALLBACK = (4, 8)
# Kinetic-electron decks that name no resolution take these instead, (slower species,
# faster species) -- docs/guidance.rst, "Velocity space per species".
_SPECIES_HL_FALLBACK = {
_RUNTIME_LINEAR_HL_FALLBACK: ((12, 32), (6, 48)),
_RUNTIME_NONLINEAR_HL_FALLBACK: ((4, 16), (4, 16)),
}
def _resolve_runtime_hl_dims(
cfg: RuntimeConfig,
*,
Nl: int | None,
Nm: int | None,
nonlinear: bool = False,
) -> tuple[RuntimeConfig, int, int]:
"""Resolve model-native Hermite/Laguerre dimensions (and per-species fallbacks)."""
model = _runtime_model_key(cfg)
if model in {"", "gyrokinetic", "full", "full-gk"}:
fallback = (
_RUNTIME_NONLINEAR_HL_FALLBACK if nonlinear else _RUNTIME_LINEAR_HL_FALLBACK
)
return cfg.with_species_fallback(
int(fallback[0] if Nl is None else Nl),
int(fallback[1] if Nm is None else Nm),
_SPECIES_HL_FALLBACK,
)
_raise_unsupported_reduced_model(cfg)
def _require_full_gk_runtime_model(cfg: RuntimeConfig) -> None:
"""Reject non-promoted reduced-model configs before full-GK execution."""
model = _runtime_model_key(cfg)
if model in {"", "gyrokinetic", "full", "full-gk"}:
return
_raise_unsupported_reduced_model(cfg)
[docs]
def runtime_geometry_config_for_builder(
cfg: RuntimeConfig,
*,
vmec_eik_builder: Callable[[RuntimeConfig], Any],
miller_eik_builder: Callable[[RuntimeConfig], Any],
) -> Any:
"""Return the geometry config consumed by the flux-tube builder."""
model = cfg.geometry.model.strip().lower()
if model == "vmec":
eik_path = vmec_eik_builder(cfg)
return replace(cfg.geometry, model="vmec-eik", geometry_file=str(eik_path))
if model == "miller":
eik_path = miller_eik_builder(cfg)
return replace(cfg.geometry, model="imported-eik", geometry_file=str(eik_path))
return cfg.geometry
[docs]
def build_runtime_geometry(cfg: RuntimeConfig) -> FluxTubeGeometryLike:
"""Resolve runtime geometry, generating `*.eik.nc` geometry when needed."""
return build_flux_tube_geometry(
runtime_geometry_config_for_builder(
cfg,
vmec_eik_builder=generate_runtime_vmec_eik,
miller_eik_builder=generate_runtime_miller_eik,
)
)
[docs]
def build_runtime_linear_params(
cfg: RuntimeConfig,
*,
Nm: int | None = None,
geom: FluxTubeGeometryLike | None = None,
) -> LinearParams:
"""Build `LinearParams` from a unified runtime config."""
if (
cfg.damping_reference is not None
and cfg.time.damp_ends_rate != cfg.damping_reference.rate
):
raise ValueError(
"damp_ends_rate disagrees with damping_reference; clear stale provenance"
)
if cfg.collisions.damp_ends_scale_by_dt:
raise ValueError(
"damp_ends_scale_by_dt=true removed; use [time] damp_ends_rate. "
"Remove flag. Old rates: linear A/(dt_input*dt_step), timestep-free "
"A/dt_input. Adaptive steps: no constant equivalent. See docs/inputs.rst."
)
_require_full_gk_runtime_model(cfg)
if geom is None:
geom = build_runtime_geometry(cfg)
contract = get_normalization_contract(cfg.normalization.contract)
rho_star, omega_d_scale, omega_star_scale = (
float(default if override is None else override)
for default, override in (
(contract.rho_star, cfg.normalization.rho_star),
(contract.omega_d_scale, cfg.normalization.omega_d_scale),
(contract.omega_star_scale, cfg.normalization.omega_star_scale),
)
)
species = _species_to_linear(cfg.species)
has_kinetic_electron = any(float(s.charge) < 0.0 for s in species)
has_kinetic_ion = any(float(s.charge) > 0.0 for s in species)
if cfg.physics.adiabatic_electrons and cfg.physics.adiabatic_ions:
raise ValueError(
"adiabatic_electrons and adiabatic_ions are mutually exclusive"
)
if cfg.physics.adiabatic_electrons and has_kinetic_electron:
raise ValueError(
"adiabatic_electrons=True conflicts with kinetic electron species"
)
if cfg.physics.adiabatic_ions and has_kinetic_ion:
raise ValueError("adiabatic_ions=True conflicts with kinetic ion species")
# ``tau_e`` is the field solver's historical name for the Boltzmann-species
# quasineutrality coefficient; it applies to either adiabatic species.
has_boltzmann_species = (
cfg.physics.adiabatic_electrons or cfg.physics.adiabatic_ions
)
tau_e = float(cfg.physics.tau_e) if has_boltzmann_species else 0.0
beta = float(cfg.physics.beta) if cfg.physics.electromagnetic else 0.0
fapar = (
1.0
if (cfg.physics.electromagnetic and cfg.physics.use_apar and beta > 0.0)
else 0.0
)
p_hyper_m = cfg.collisions.p_hyper_m
if p_hyper_m is None:
p_hyper_m = _default_hermite_hypercollision_exponent(Nm)
params = build_linear_params(
species,
tau_e=tau_e,
kpar_scale=float(geom.gradpar()),
omega_d_scale=omega_d_scale,
omega_star_scale=omega_star_scale,
rho_star=rho_star,
beta=beta,
fapar=fapar,
nu_hyper=float(cfg.collisions.nu_hyper),
p_hyper=float(cfg.collisions.p_hyper),
nu_hyper_l=float(cfg.collisions.nu_hyper_l),
nu_hyper_m=float(cfg.collisions.nu_hyper_m),
nu_hyper_lm=float(cfg.collisions.nu_hyper_lm),
p_hyper_l=float(cfg.collisions.p_hyper_l),
p_hyper_m=float(p_hyper_m),
p_hyper_lm=float(cfg.collisions.p_hyper_lm),
D_hyper=float(cfg.collisions.D_hyper),
p_hyper_kperp=float(cfg.collisions.p_hyper_kperp),
hypercollisions_const=float(cfg.collisions.hypercollisions_const),
hypercollisions_kz=float(cfg.collisions.hypercollisions_kz),
)
return replace(
params,
moments=_species_moments(cfg, Nm),
nu_hyper_m_const=(
None
if cfg.collisions.nu_hyper_m_const is None
else float(cfg.collisions.nu_hyper_m_const)
),
nu_hermite=float(cfg.collisions.nu_hermite),
nu_laguerre=float(cfg.collisions.nu_laguerre),
damp_ends_widthfrac=float(cfg.collisions.damp_ends_widthfrac),
damp_ends_rate=_runtime_damp_ends_rate(cfg),
mu_vperp_max=cfg.run.mu_vperp_max,
collision_model=(
"landau"
if str(cfg.time.collision_operator).strip().lower() in ("landau", "coulomb")
else "dougherty"
),
)
def _runtime_damp_ends_rate(cfg: RuntimeConfig) -> float | None:
"""Return ``[time] damp_ends_rate``, refusing the ambiguous legacy spelling.
``None`` selects the default rate of :meth:`LinearParams.end_damping_rate`.
"""
rate, amp = cfg.time.damp_ends_rate, cfg.collisions.damp_ends_amp
if rate is not None or amp is None:
return rate
if amp == 0.0:
return 0.0
if cfg.grid.boundary == "periodic":
return None
dt = float(cfg.time.dt)
raise ValueError(
f"[collisions] damp_ends_amp = {amp:g} without [time] damp_ends_rate is "
"ambiguous: GKX < 2.6 damped the field-line ends at damp_ends_amp/dt on "
f"the time route ({amp / dt:g} at dt = {dt:g}) and at damp_ends_amp on "
"the timestep-free and nonlinear routes. Set [time] damp_ends_rate to "
"the rate you mean, or delete damp_ends_amp to use the default rate "
"(docs/inputs.rst, end damping)."
)
[docs]
def end_damping_status(params: LinearParams, grid: Any, boundary: str) -> str:
"""Describe the end-damping rate a run uses, for the run's status line."""
if boundary == "periodic":
return "end damping: off (periodic boundary)"
z = np.asarray(grid.z, dtype=float)
z_extent = float(z.size * (z[1] - z[0])) if z.size > 1 else 0.0
rate = float(params.end_damping_rate(z_extent, jnp.float32))
if params.damp_ends_rate is not None:
return f"end damping: rate {rate:.6g} ([time] damp_ends_rate)"
return (
f"end damping: rate {rate:.6g} (default: {DAMP_ENDS_LAYER_TRANSITS:g} "
"thermal transits of the fastest kinetic species across the "
f"{float(params.damp_ends_widthfrac):g} end layers of Lz = {z_extent:.6g})"
)
[docs]
def build_runtime_linear_terms(
cfg: RuntimeConfig, *, custom_collisions: bool = False
) -> LinearTerms:
"""Build `LinearTerms` from unified toggles.
All-zero species ``nu`` drop the built-in collision term; a caller passing
a custom operator with its own rates sets ``custom_collisions`` (#324).
"""
em_on = bool(cfg.physics.electromagnetic)
use_apar = em_on and bool(cfg.physics.use_apar)
use_bpar = em_on and bool(cfg.physics.use_bpar)
collisions_on = bool(cfg.physics.collisions) and (
custom_collisions or any(float(sp.nu) != 0.0 for sp in cfg.species)
)
hyper_on = bool(cfg.physics.hypercollisions)
return LinearTerms(
streaming=float(cfg.terms.streaming),
mirror=float(cfg.terms.mirror),
curvature=float(cfg.terms.curvature),
gradb=float(cfg.terms.gradb),
diamagnetic=float(cfg.terms.diamagnetic),
collisions=float(cfg.terms.collisions if collisions_on else 0.0),
hypercollisions=float(cfg.terms.hypercollisions if hyper_on else 0.0),
hyperdiffusion=float(cfg.terms.hyperdiffusion),
end_damping=float(cfg.terms.end_damping),
apar=float(cfg.terms.apar if use_apar else 0.0),
bpar=float(cfg.terms.bpar if use_bpar else 0.0),
)
[docs]
def build_runtime_term_config(cfg: RuntimeConfig) -> TermConfig:
"""Build nonlinear-ready `TermConfig` from unified toggles."""
lin_terms = build_runtime_linear_terms(cfg)
nonlinear_on = float(cfg.terms.nonlinear if cfg.physics.nonlinear else 0.0)
return linear_terms_to_term_config(lin_terms, nonlinear=nonlinear_on)
# ---- merged from workflows/runtime/initial_conditions.py ----
# That module had exactly one consumer -- this one -- which imported ten of
# its private names. A boundary that wide is a split, not an interface.
_GLIBC_RAND_MAX = float((1 << 31) - 1)
_INITIAL_FIELD_MOMENTS: dict[str, tuple[int, int]] = {
"density": (0, 0),
"upar": (0, 1),
"tpar": (0, 2),
"tperp": (1, 0),
"qpar": (0, 3),
"qperp": (1, 1),
}
_ALL_FIELD_SCALES: dict[str, float] = {
"density": 1.0,
"upar": 1.0,
"tpar": 1.0 / np.sqrt(2.0),
"tperp": 1.0,
"qpar": 1.0 / np.sqrt(6.0),
"qperp": 1.0,
}
_VALID_INIT_FIELDS = {"all", "phi", *_INITIAL_FIELD_MOMENTS.keys()}
def _build_gaussian_profile(
z: np.ndarray,
*,
kx: float,
ky: float,
s_hat: float,
width: float,
envelope_constant: float,
envelope_sine: float,
) -> np.ndarray:
if ky == 0.0:
return np.zeros_like(z)
theta0 = kx / (s_hat * ky)
env = envelope_constant + envelope_sine * np.sin(z - theta0)
return env * np.exp(-(((z - theta0) / width) ** 2))
[docs]
def _build_single_phi_gaussian_profile(
z: np.ndarray,
*,
kx: float,
ky: float,
s_hat: float,
width: float,
envelope_constant: float,
envelope_sine: float,
) -> np.ndarray:
"""Return a single-mode Gaussian potential profile along the flux tube.
The W7-X zonal-flow benchmark prescribes a Gaussian electrostatic-potential
perturbation centered in the middle of the flux tube. A multi-mode
ballooning-angle Gaussian initializer is undefined for ``ky=0`` because its
center uses ``kx / (s_hat * ky)``; for the zonal case the physically stated
center is therefore the tube midpoint, ``z=0``.
"""
if ky != 0.0 and s_hat != 0.0:
return _build_gaussian_profile(
z,
kx=kx,
ky=ky,
s_hat=s_hat,
width=width,
envelope_constant=envelope_constant,
envelope_sine=envelope_sine,
)
center = 0.0
env = envelope_constant + envelope_sine * np.sin(z - center)
return env * np.exp(-(((z - center) / width) ** 2))
def _reshape_netcdf_state(
raw: np.ndarray,
*,
nspec: int,
nl: int,
nm: int,
nyc: int,
nx: int,
nz: int,
) -> np.ndarray:
nR = nyc * nx * nz
arr = raw.reshape((nspec, nm, nl, nR)).transpose(0, 2, 1, 3)
ky_idx = np.arange(nyc)[:, None, None]
kx_idx = np.arange(nx)[None, :, None]
z_idx = np.arange(nz)[None, None, :]
idxyz = ky_idx + nyc * (kx_idx + nx * z_idx)
arr_reordered = arr[..., idxyz.ravel()]
return arr_reordered.reshape((nspec, nl, nm, nyc, nx, nz))
[docs]
def _expand_ky(arr: np.ndarray, *, ny_full: int) -> np.ndarray:
"""Widen a stored ``ky >= 0`` init state onto this grid's two-sided axis.
``ny_full`` is passed in rather than inferred from ``Nyc``: the two map
many-to-one (``gkx.core_ky_layout``), and inferring the even branch
silently produced an ``Ny - 1`` row state for an odd-``Ny`` grid.
"""
ny = int(ny_full)
if ny <= 1 or arr.shape[-3] == ny:
return arr
return to_full(arr, ny_full=ny)
[docs]
def _enforce_full_ky_hermitian(arr: np.ndarray) -> np.ndarray:
"""Mirror positive-`ky` content into the negative branch for full FFT grids."""
state = np.asarray(arr, dtype=np.complex64)
ny = int(state.shape[-3])
if ny <= 1:
return state
nyc = nyc_from_ny(ny)
if paired_row_limit(ny) <= 1:
return state
state[..., nyc:, :, :] = negative_ky_block(state, ny_full=ny)
return state
[docs]
def _load_initial_state_from_file(
path: Path,
*,
nspecies: int,
Nl: int,
Nm: int,
ny: int,
nx: int,
nz: int,
ky_layout: KyLayout = FULL,
) -> np.ndarray:
"""Load a restart or seed file onto ``ky_layout``.
``ny`` is the two-sided axis length in both layouts, which is what the
file's own contents are indexed by; ``ky_layout`` decides how many rows
come back. Both file forms keep working on both axes: a binary file of
``Nyc`` rows is the ``ky >= 0`` block and is widened only when the caller
asked for the two-sided axis, and a two-sided file is narrowed only when
the caller asked for the half one. Neither conversion loses anything --
the discarded rows are the conjugates of the kept ones (plan 5.3 N3).
"""
ky_layout = cast(KyLayout, str(ky_layout).strip().lower())
half = ky_layout == HALF
if path.suffix.lower() == ".nc":
return load_netcdf_restart_state(
path,
nspecies=nspecies,
Nl=Nl,
Nm=Nm,
ny=ny,
nx=nx,
nz=nz,
ky_layout=ky_layout,
)
raw = np.fromfile(path, dtype=np.complex64)
nyc = nyc_from_ny(ny)
expected_nyc = nspecies * Nl * Nm * nyc * nx * nz
expected_full = nspecies * Nl * Nm * ny * nx * nz
if raw.size == expected_nyc:
arr = _reshape_netcdf_state(
raw, nspec=nspecies, nl=Nl, nm=Nm, nyc=nyc, nx=nx, nz=nz
)
return arr if half else _expand_ky(arr, ny_full=ny)
if raw.size == expected_full:
full = raw.reshape((nspecies, Nl, Nm, ny, nx, nz))
return to_half(full, ny_full=ny) if half else full
raise ValueError(
f"init_file size {raw.size} does not match expected {expected_nyc} (nyc) or {expected_full} (full)"
)
[docs]
def _centered_glibc_random_pairs(seed: int, count: int) -> np.ndarray:
"""Return centered random pairs using glibc `rand()` semantics."""
if count <= 0:
return np.empty((0, 2), dtype=np.float64)
seed_use = 1 if int(seed) == 0 else int(seed)
state = np.zeros(344 + 2 * count, dtype=np.uint64)
state[0] = np.uint64(seed_use)
for i in range(1, 31):
state[i] = np.uint64((16807 * int(state[i - 1])) % int(_GLIBC_RAND_MAX))
for i in range(31, 34):
state[i] = state[i - 31]
for i in range(34, state.size):
state[i] = (state[i - 31] + state[i - 3]) & np.uint64(0xFFFFFFFF)
rand_vals = (state[344:] >> np.uint64(1)).astype(np.float64, copy=False)
half = 0.5 * _GLIBC_RAND_MAX
inv = 1.0 / _GLIBC_RAND_MAX
pairs = np.empty((count, 2), dtype=np.float64)
for i in range(count):
pairs[i, 0] = (rand_vals[2 * i] - half) * inv
pairs[i, 1] = (rand_vals[2 * i + 1] - half) * inv
return pairs
[docs]
def _dealiased_initial_mode_pairs(grid: SpectralGrid) -> list[tuple[int, int]]:
"""Return the dealiased startup-loop `(kx, ky)` pairs for multimode initial conditions.
The binormal index 0 is skipped only when it really is the zonal mode. A
linear run selects one ``k_y`` before seeding, so its grid holds a single
nonzero binormal entry at index 0; skipping it by position left the whole
state at zero and the eigensolver was handed a null seed.
"""
nx = int(np.asarray(grid.kx).size)
ky_values = np.asarray(grid.ky)
rows = int(ky_values.size)
# The seeded band is the dealiased non-negative rows of the *two-sided*
# axis: ``1 + (Ny - 1) // 3`` is where ``|ky| < Ny/3`` stops, in both
# layouts. Reading it off the stored count would seed only
# ``1 + (Nyc - 1) // 3`` rows on a half-spectrum grid, a third of the band,
# with no error (gkx.core_ky_layout).
ny = source_ny_full(grid)
kx_max = 1 + (nx - 1) // 3
ky_max = min(1 + (ny - 1) // 3, rows)
ky_indices = [int(ky_i) for ky_i in range(ky_max) if float(ky_values[ky_i]) != 0.0]
return [(int(kx_i), ky_i) for kx_i in range(kx_max) for ky_i in ky_indices]
[docs]
def _periodic_zp_from_grid(z: np.ndarray) -> float:
"""Return periodic `Zp` from the discrete theta grid."""
z_arr = np.asarray(z, dtype=float)
if z_arr.size <= 1:
return 1.0
dz = float(z_arr[1] - z_arr[0])
period = abs(dz) * float(z_arr.size)
if period <= 0.0:
return 1.0
return period / (2.0 * np.pi)
def _validate_initialization(cfg: RuntimeConfig) -> tuple[str, str]:
init_field = cfg.init.init_field.lower()
if init_field not in _VALID_INIT_FIELDS:
raise ValueError(
"init_field must be one of {'density','upar','tpar','tperp','qpar','qperp','all','phi'}"
)
if cfg.init.gaussian_width <= 0.0:
raise ValueError("gaussian_width must be > 0")
init_file_mode = cfg.init.init_file_mode.strip().lower()
if init_file_mode not in {"replace", "add"}:
raise ValueError("init_file_mode must be one of {'replace', 'add'}")
return init_field, init_file_mode
def _scaled_restart_state(
cfg: RuntimeConfig,
grid: SpectralGrid,
*,
Nl: int,
Nm: int,
nspecies: int,
) -> np.ndarray | None:
if cfg.init.init_file is None:
return None
loaded_state = _load_initial_state_from_file(
Path(cfg.init.init_file),
nspecies=nspecies,
Nl=Nl,
Nm=Nm,
ny=source_ny_full(grid),
nx=grid.kx.size,
nz=grid.z.size,
ky_layout=source_ky_layout(grid),
)
return np.asarray(loaded_state, dtype=np.complex64) * np.complex64(
float(cfg.init.init_file_scale)
)
def _species_targets(cfg: RuntimeConfig, nspecies: int) -> tuple[int, ...]:
if nspecies == 1:
return (0,)
if not cfg.init.init_electrons_only:
return tuple(range(nspecies))
electron_indices = tuple(
i for i, sp in enumerate(cfg.species[:nspecies]) if float(sp.charge) < 0.0
)
return electron_indices or (nspecies - 1,)
def _single_mode_values(
grid: SpectralGrid,
geom: FluxTubeGeometryLike,
cfg: RuntimeConfig,
*,
init_field: str,
ky_index: int,
kx_index: int,
) -> np.ndarray:
z = np.asarray(grid.z)
z_period = _periodic_zp_from_grid(z)
z_phase = np.cos(float(cfg.init.kpar_init) * z / z_period)
amp = float(cfg.init.init_amp)
if cfg.init.init_single and cfg.init.gaussian_init:
profile = _build_single_phi_gaussian_profile(
z,
kx=float(grid.kx[kx_index]),
ky=float(grid.ky[ky_index]),
s_hat=float(geom.s_hat),
width=float(cfg.init.gaussian_width),
envelope_constant=float(cfg.init.gaussian_envelope_constant),
envelope_sine=float(cfg.init.gaussian_envelope_sine),
)
phase = 1.0 if init_field == "phi" else (1.0 + 1.0j)
return amp * phase * profile.astype(np.complex64, copy=False)
return amp * z_phase.astype(np.complex64, copy=False)
def _validate_named_moment_resolution(init_field: str, *, Nl: int, Nm: int) -> None:
if init_field in {"all", "phi"}:
return
l_idx, m_idx = _INITIAL_FIELD_MOMENTS[init_field]
if l_idx >= Nl or m_idx >= Nm:
raise ValueError("init_field moment exceeds (Nl, Nm) resolution")
@dataclass
class _InitialConditionBuilder:
grid: SpectralGrid
geom: FluxTubeGeometryLike
cfg: RuntimeConfig
Nl: int
Nm: int
state: np.ndarray
species_targets: tuple[int, ...]
phi_seed_context: tuple[object, LinearParams] | None = None
def set_mode(
self, l_idx: int, m_idx: int, ky_i: int, kx_i: int, vals_k: np.ndarray
) -> None:
if l_idx >= self.Nl or m_idx >= self.Nm:
return
rows: int | slice = l_idx
if (vperp_max := self.cfg.run.mu_vperp_max) is not None: # nodal mu grid
theta = jnp.asarray(self.grid.z)
bmag = np.asarray(
ensure_flux_tube_geometry_data(self.geom, theta).bmag(theta)
)
nodes = mu_grid_nodes(self.Nl, float(vperp_max))
modes = mu_grid_laguerre_modes(*nodes, bmag, l_idx + 1)[l_idx]
vals_k, rows = modes * np.asarray(vals_k)[None], slice(None)
for s_idx in self.species_targets:
self.state[s_idx, rows, m_idx, ky_i, kx_i, :] = vals_k
def set_named_mode_scaled(
self, field_name: str, ky_i: int, kx_i: int, vals_k: np.ndarray
) -> None:
l_idx, m_idx = _INITIAL_FIELD_MOMENTS[field_name]
self.set_mode(l_idx, m_idx, ky_i, kx_i, vals_k * _ALL_FIELD_SCALES[field_name])
def set_named_mode_raw(
self, field_name: str, ky_i: int, kx_i: int, vals_k: np.ndarray
) -> None:
l_idx, m_idx = _INITIAL_FIELD_MOMENTS[field_name]
self.set_mode(l_idx, m_idx, ky_i, kx_i, vals_k)
def set_phi_mode(self, ky_i: int, kx_i: int, vals_k: np.ndarray) -> None:
if self.Nl < 1 or self.Nm < 1:
raise ValueError(
"init_field='phi' requires at least one Laguerre and one Hermite moment"
)
if self.phi_seed_context is None:
phi_params = build_runtime_linear_params(
self.cfg, Nm=self.Nm, geom=self.geom
)
self.phi_seed_context = (
build_linear_cache(self.grid, self.geom, phi_params, self.Nl, self.Nm),
phi_params,
)
cache, phi_params = self.phi_seed_context
seeds = _density_moments_for_target_phi(
np.asarray(vals_k, dtype=np.complex64),
cache=cache,
params=phi_params,
ky_i=int(ky_i),
kx_i=int(kx_i),
species_targets=self.species_targets,
)
for s_idx, seed_vals in seeds.items():
self.state[s_idx, 0, 0, ky_i, kx_i, :] = seed_vals
def seed_field(
self, init_field: str, ky_i: int, kx_i: int, vals_k: np.ndarray
) -> None:
if init_field == "all":
for field_name in _INITIAL_FIELD_MOMENTS:
self.set_named_mode_scaled(field_name, ky_i, kx_i, vals_k)
elif init_field == "phi":
self.set_phi_mode(ky_i, kx_i, vals_k)
else:
self.set_named_mode_raw(init_field, ky_i, kx_i, vals_k)
def _seed_gaussian_multimode(
builder: _InitialConditionBuilder,
*,
init_field: str,
amp: float,
) -> None:
z = np.asarray(builder.grid.z)
nx = builder.grid.kx.size
for kx_i, ky_i in _dealiased_initial_mode_pairs(builder.grid):
ky_k = float(builder.grid.ky[ky_i])
if ky_k == 0.0:
continue
profile_k = _build_gaussian_profile(
z,
kx=abs(float(builder.grid.kx[kx_i])),
ky=ky_k,
s_hat=float(builder.geom.s_hat),
width=float(builder.cfg.init.gaussian_width),
envelope_constant=float(builder.cfg.init.gaussian_envelope_constant),
envelope_sine=float(builder.cfg.init.gaussian_envelope_sine),
)
vals_k = amp * profile_k * (1.0 + 1.0j)
builder.seed_field(init_field, ky_i, kx_i, vals_k)
if kx_i != 0:
builder.seed_field(init_field, ky_i, int(nx - kx_i), vals_k)
def _seed_random_multimode(
builder: _InitialConditionBuilder,
*,
init_field: str,
amp: float,
) -> None:
_validate_named_moment_resolution(init_field, Nl=builder.Nl, Nm=builder.Nm)
z = np.asarray(builder.grid.z)
z_phase = np.cos(float(builder.cfg.init.kpar_init) * z / _periodic_zp_from_grid(z))
nx = builder.grid.kx.size
active_modes = _dealiased_initial_mode_pairs(builder.grid)
rand_pairs = amp * _centered_glibc_random_pairs(
int(builder.cfg.init.random_seed), len(active_modes)
)
for (kx_i, ky_i), (ra, rb) in zip(active_modes, rand_pairs, strict=True):
vals_k = ((rb + 1j * ra) if kx_i == 0 else (ra + 1j * rb)) * z_phase
builder.seed_field(init_field, ky_i, kx_i, vals_k)
if kx_i != 0:
vals_neg = (rb + 1j * ra) * z_phase
builder.seed_field(init_field, ky_i, int(nx - kx_i), vals_neg)
[docs]
def runtime_state_dtype() -> np.dtype:
"""Return the complex dtype the runtime's own initial state is handed in.
Every seed in this module is assembled in ``complex64``: the seed is an
arbitrary small perturbation, so its own rounding is immaterial, and holding
the assembly at one width keeps the ky-Hermitian and off-chain masking code
free of a dtype branch. What is *not* immaterial is the width the seed
imposes downstream. ``jnp.asarray`` does not promote, so a ``complex64``
seed pinned the whole run to ``complex64`` even under ``JAX_ENABLE_X64``:
the Krylov basis, the eigenvector and therefore the certification gate,
which ``certifiable_residual_tolerance`` floors at ``1e3 * eps(dtype)``
-- 1.19e-4 in float32 against the 1e-9 the float64 gate applies. Decks that
instruct the reader to run under ``JAX_ENABLE_X64`` (the shipped Cyclone
deck is one) silently got the float32 answer and a float32 gate.
The seed is therefore widened once, here, to the working precision. This
does not change the default: without ``JAX_ENABLE_X64`` the dtype is
``complex64`` exactly as before. It matches what the ``initial_state`` API
already promises for a caller-supplied seed, which preserves an explicit
``complex128`` array including any prior initialization rounding.
"""
return np.dtype(np.complex128) if _x64_enabled() else np.dtype(np.complex64)
def _finalize_initial_state(
grid: SpectralGrid,
state: np.ndarray,
*,
loaded_state: np.ndarray | None,
init_file_mode: str,
) -> jnp.ndarray:
if grid.ky.size > 1 and np.any(np.asarray(grid.ky) < 0.0):
state = _enforce_full_ky_hermitian(state)
dtype = runtime_state_dtype()
if loaded_state is None:
return jnp.asarray(state, dtype=dtype)
if init_file_mode == "replace":
return jnp.asarray(loaded_state, dtype=dtype)
return jnp.asarray(cast(np.ndarray, loaded_state + state), dtype=dtype)
[docs]
def _build_initial_condition(
grid: SpectralGrid,
geom: FluxTubeGeometryLike,
cfg: RuntimeConfig,
*,
ky_index: int,
kx_index: int,
Nl: int,
Nm: int,
nspecies: int,
) -> jnp.ndarray:
"""Build the runtime initial state for one selected mode or seed policy."""
init_field, init_file_mode = _validate_initialization(cfg)
state: np.ndarray = np.zeros(
(nspecies, Nl, Nm, grid.ky.size, grid.kx.size, grid.z.size),
dtype=np.complex64,
)
loaded_state = _scaled_restart_state(cfg, grid, Nl=Nl, Nm=Nm, nspecies=nspecies)
if loaded_state is not None:
# A restart or init file is state this runtime did not build, so it can
# carry rows the linked chains never reach. Zero them here, once, on
# read. Nothing this module seeds is outside the chains, so the
# runtime's own initial conditions are untouched, and periodic decks
# have no mask at all.
loaded_state = np.asarray(
mask_off_chain_rows(
loaded_state,
grid,
geom,
build_runtime_linear_params(cfg, Nm=Nm, geom=geom),
),
dtype=np.complex64,
)
amp = float(cfg.init.init_amp)
builder = _InitialConditionBuilder(
grid=grid,
geom=geom,
cfg=cfg,
Nl=Nl,
Nm=Nm,
state=state,
species_targets=_species_targets(cfg, nspecies),
)
if cfg.init.gaussian_init and not cfg.init.init_single:
_seed_gaussian_multimode(builder, init_field=init_field, amp=amp)
elif not cfg.init.init_single and not cfg.init.gaussian_init:
_seed_random_multimode(builder, init_field=init_field, amp=amp)
else:
_validate_named_moment_resolution(init_field, Nl=Nl, Nm=Nm)
vals = _single_mode_values(
grid,
geom,
cfg,
init_field=init_field,
ky_index=ky_index,
kx_index=kx_index,
)
builder.seed_field(init_field, ky_index, kx_index, vals)
state0 = _finalize_initial_state(
grid,
state,
loaded_state=loaded_state,
init_file_mode=init_file_mode,
)
return keep_moments(state0, _species_moments(cfg, Nm))
__all__ = [
"_build_gaussian_profile",
"_build_initial_condition",
"_build_single_phi_gaussian_profile",
"_as_runtime_species_array",
"_centered_glibc_random_pairs",
"_dealiased_initial_mode_pairs",
"_density_moments_for_target_phi",
"_enforce_full_ky_hermitian",
"_expand_ky",
"_load_initial_state_from_file",
"_periodic_zp_from_grid",
"_reshape_netcdf_state",
"build_runtime_geometry",
"build_runtime_linear_params",
"end_damping_status",
"build_runtime_linear_terms",
"build_runtime_term_config",
"load_netcdf_restart_state",
"runtime_geometry_config_for_builder",
"runtime_state_dtype",
]