"""Linear parameter, term-toggle, and validation policy helpers."""
from __future__ import annotations
import warnings
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING, Any, Callable, Iterable
import jax
import jax.numpy as jnp
import numpy as np
if TYPE_CHECKING:
from gkx.terms.config import TermConfig
__all__ = [
"BUILTIN_COLLISION_OPERATORS",
"COLLISION_OPERATOR_NAMES",
"DAMP_ENDS_LAYER_TRANSITS",
"DEPRECATED_LINEAR_PARAM_ALIASES",
"linear_params_for_geometry",
"LinearParams",
"LinearTerms",
"Species",
"Preconditioner",
"PreconditionerSpec",
"_as_species_array",
"_check_nonnegative",
"_check_positive",
"_is_tracer",
"_resolve_implicit_preconditioner",
"_x64_enabled",
"linear_terms_to_term_config",
"build_linear_params",
"term_config_to_linear_terms",
]
[docs]
@dataclass(frozen=True)
class Species:
r"""Dimensionless kinetic-species inputs used to build solver parameters.
``tprim`` and ``fprim`` are :math:`a/L_T` and :math:`a/L_n`, copied through
to :class:`LinearParams` unchanged; see its docstring for the conversion to
:math:`R/L`.
"""
charge: float
mass: float
density: float
temperature: float
tprim: float
fprim: float
nu: float = 0.0
closure: float = 0.0
#: Retired :class:`LinearParams` field names, mapped to what replaced them. The
#: old names said ``R/L``; the values they held were always ``a/L``. Kept
#: working -- as constructor keywords and as attributes, both warning -- because
#: ``LinearParams`` is public API.
DEPRECATED_LINEAR_PARAM_ALIASES: dict[str, str] = {
"R_over_LTi": "tprim",
"R_over_Ln": "fprim",
"R_over_LTe": "tprim_e",
}
#: Default end-absorber strength: e-folds per thermal transit of the fastest
#: kinetic species across one damping layer (:meth:`LinearParams.end_damping_rate`).
DAMP_ENDS_LAYER_TRANSITS = 30.0
def _warn_alias(old: str, new: str) -> None:
warnings.warn(
f"LinearParams.{old} is deprecated; use {new}. The operator consumes "
f"a/L gradients (the TOML 'tprim'/'fprim'), never R/L, so the old name "
f"asserted a normalization that was never applied -- a value passed as "
f"R/L needs dividing by R/a.",
DeprecationWarning,
stacklevel=3,
)
def _with_deprecated_aliases(aliases: dict[str, str]):
"""Keep retired field names working as constructor keywords and attributes."""
def decorate(cls):
base_init = cls.__init__
def __init__(self, *args, **kwargs):
for old, new in aliases.items():
if old in kwargs:
_warn_alias(old, new)
kwargs[new] = kwargs.pop(old)
base_init(self, *args, **kwargs)
cls.__init__ = __init__
for old, new in aliases.items():
def read(self, _old=old, _new=new):
_warn_alias(_old, _new)
return getattr(self, _new)
setattr(cls, old, property(read))
return cls
return decorate
[docs]
@jax.tree_util.register_pytree_node_class
@_with_deprecated_aliases(DEPRECATED_LINEAR_PARAM_ALIASES)
@dataclass(frozen=True)
class LinearParams:
r"""Parameters for the linear gyrokinetic operator (supports multi-species arrays).
``tprim`` and ``fprim`` are the gradient drives, in the units the operator
actually consumes: :math:`a/L_T` and :math:`a/L_n` with :math:`L_{ref}=a`,
which is what the TOML ``tprim``/``fprim`` keys carry and what
:func:`build_linear_params` copies across unchanged. They are *not*
:math:`R/L_T` and :math:`R/L_n`; multiply by :math:`R/a` to quote a result
in those. The defaults are the Cyclone base case, matching
``benchmarks/cases/cyclone_nonlinear_t400.toml``
and :class:`gkx.config.ModelConfig`, so
:math:`R/L_T = 2.49 \times 2.77778 = 6.92`.
``tprim`` is the per-species array the RHS diamagnetic drive reads.
``tprim_e`` does not enter the RHS at all -- it only widens the frequency
scale used for eigenvalue branch selection, so a kinetic-electron case must
put the electron drive in ``tprim``, not here.
"""
charge_sign: float | jnp.ndarray = 1.0
density: float | jnp.ndarray = 1.0
mass: float | jnp.ndarray = 1.0
temp: float | jnp.ndarray = 1.0
tau_e: float = 1.0
vth: float | jnp.ndarray = 1.0
rho: float | jnp.ndarray = 1.0
kpar_scale: float = 1.0
fprim: float | jnp.ndarray = 0.8
tprim: float | jnp.ndarray = 2.49
tprim_e: float | jnp.ndarray = 0.0
omega_d_scale: float = 1.0
omega_star_scale: float = 1.0
energy_const: float = 0.0
energy_par_coef: float = 0.5
energy_perp_coef: float = 1.0
nu: float | jnp.ndarray = 0.0
nu_hermite: float = 1.0
nu_laguerre: float = 2.0
nu_hyper: float = 0.0
p_hyper: float = 4.0
nu_hyper_l: float = 0.0
nu_hyper_m: float = 1.0
nu_hyper_m_const: float | None = None
nu_hyper_lm: float = 0.0
p_hyper_l: float = 6.0
p_hyper_m: float = 20.0
p_hyper_lm: float = 6.0
hypercollisions_const: float = 1.0
hypercollisions_kz: float = 0.0
#: Per species, 1 closes the Hermite hierarchy with the reflectionless
#: sink on the last moment (:func:`hermite_closure_profile`) and drops that
#: species' ``|k_z|`` hypercollisions; 0 keeps the hypercollisions.
closure: float | jnp.ndarray = 0.0
D_hyper: float = 0.0
p_hyper_kperp: float = 2.0
damp_ends_widthfrac: float | jnp.ndarray = 0.125
tz: float | jnp.ndarray = 1.0
rho_star: float = 1.0
beta: float = 0.0
fapar: float = 0.0
apar_beta_scale: float = 0.5
ampere_g0_scale: float = 0.5
bpar_beta_scale: float = 0.5
damp_ends_rate: float | jnp.ndarray | None = None
#: Built-in collision kernel: ``"landau"`` (linearized Landau, drift-kinetic
#: limit; the default, as in the runtime) or ``"dougherty"`` (conserving
#: Lenard-Bernstein form).
collision_model: str = "landau"
#: Perpendicular velocity representation. ``None`` keeps the Laguerre
#: expansion in ``mu B``. A number selects the nodal magnetic-moment grid
#: (:func:`gkx.core_velocity.mu_grid_nodes`) with that ``vperp_max``
#: (``inf``: Gauss-Laguerre); the Laguerre axis then holds the ``Nl`` nodes.
mu_vperp_max: float | None = None
#: Moments each species keeps, one ``(Nl, Nm, p_hyper_m)`` per species, at most the
#: state's ``(Nl, Nm)``; ``None`` in a slot keeps the state's there and ``None`` for
#: the field gives every species the state's. The RHS runs each run of equal species
#: at its own size, so the moments above it stay zero and its hypercollisions are
#: normalized to its own count (:attr:`LinearCache.groups`).
moments: tuple[tuple[int | None, int | None, float], ...] | None = None
[docs]
def species_slice(self, a: int, b: int, **overrides: Any) -> "LinearParams":
"""The parameters of species ``a:b`` of a ``moments`` run, with ``overrides``."""
ns = len(self.moments or ())
rows: dict[str, Any] = {
n: _as_species_array(getattr(self, n), ns, n)[a:b]
for n in _SPECIES_PARAM_NAMES
}
return replace(self, **{**rows, "moments": None, **overrides})
[docs]
def const_branch_nu_hyper_m(self):
"""Return the Hermite hypercollision rate the *constant* branch uses.
``nu_hyper_m`` drives both hypercollision branches, so a run that wants
the shipped ``|k_z|`` Hermite model *and* a constant-branch Laguerre
sink would otherwise damp Hermite twice. ``nu_hyper_m_const`` splits
the two: ``None`` keeps the historical behaviour (the constant branch
reuses ``nu_hyper_m``), and ``0.0`` gives a pure Laguerre sink beside an
unchanged ``|k_z|`` branch. This is the only knob the Laguerre
regularization contract (:ref:`plan 0.5 <velocity-regularization>`)
needs that the operator could not express.
"""
if self.nu_hyper_m_const is None:
return self.nu_hyper_m
return self.nu_hyper_m_const
[docs]
def end_damping_rate(self, z_extent, dtype):
r"""Return the end-absorber rate :math:`\nu_{end}`, the same on every route.
An explicit ``damp_ends_rate`` wins. Otherwise
:math:`\nu_{end} = C\,\max_s v_{ts}\,k_{\parallel,s}/L_w`, with
:math:`k_{\parallel,s}` = ``kpar_scale`` (the geometry's ``gradpar``) and
:math:`L_w` = ``damp_ends_widthfrac`` times ``z_extent``: ``C`` e-folds
per thermal transit of the fastest kinetic species across the layer
(:data:`DAMP_ENDS_LAYER_TRANSITS`). No timestep enters.
"""
if self.damp_ends_rate is not None:
return jnp.asarray(self.damp_ends_rate, dtype=dtype)
width = jnp.asarray(self.damp_ends_widthfrac, dtype=dtype) * z_extent
transit = jnp.max(jnp.asarray(self.vth, dtype=dtype)) * jnp.abs(
jnp.asarray(self.kpar_scale, dtype=dtype)
)
return DAMP_ENDS_LAYER_TRANSITS * transit / jnp.where(width > 0, width, 1.0)
def tree_flatten(self):
children = (
self.charge_sign,
self.density,
self.mass,
self.temp,
self.tau_e,
self.vth,
self.rho,
self.kpar_scale,
self.fprim,
self.tprim,
self.tprim_e,
self.omega_d_scale,
self.omega_star_scale,
self.energy_const,
self.energy_par_coef,
self.energy_perp_coef,
self.nu,
self.nu_hermite,
self.nu_laguerre,
self.nu_hyper,
self.p_hyper,
self.nu_hyper_l,
self.nu_hyper_m,
self.nu_hyper_m_const,
self.nu_hyper_lm,
self.p_hyper_l,
self.p_hyper_m,
self.p_hyper_lm,
self.hypercollisions_const,
self.hypercollisions_kz,
self.closure,
self.D_hyper,
self.p_hyper_kperp,
self.damp_ends_widthfrac,
self.tz,
self.rho_star,
self.beta,
self.fapar,
self.apar_beta_scale,
self.ampere_g0_scale,
self.bpar_beta_scale,
self.damp_ends_rate,
)
return children, (self.collision_model, self.mu_vperp_max, self.moments)
@classmethod
def tree_unflatten(cls, aux_data, children):
return cls(
*children,
collision_model=aux_data[0],
mu_vperp_max=aux_data[1],
moments=aux_data[2],
)
_SPECIES_PARAM_NAMES = (
"charge_sign",
"density",
"mass",
"temp",
"vth",
"rho",
"fprim",
"tprim",
"tprim_e",
"nu",
"tz",
"closure",
)
[docs]
def linear_params_for_geometry(geometry, **overrides) -> LinearParams:
r"""Build :class:`LinearParams` whose parallel scale matches the geometry.
``kpar_scale`` multiplies the parallel derivative, so it has to carry the
geometry's ``gradpar`` :math:`=1/(qR_0)` for s-alpha. It defaults to ``1.0``
because :class:`LinearParams` knows nothing about geometry, which means a
hand-built ``LinearParams()`` silently streams at the wrong rate -- by a
factor of 8.4 for a typical Cyclone-like case. The runtime path sets it from
the geometry; this is the equivalent for the Python API.
Pass ``kpar_scale`` explicitly to override, for instance when reproducing a
published normalization that folds :math:`qR_0` in elsewhere.
"""
if "kpar_scale" not in overrides:
overrides["kpar_scale"] = float(np.asarray(geometry.gradpar()))
return LinearParams(**overrides)
[docs]
def build_linear_params(
species: Iterable[Species],
**overrides: Any,
) -> LinearParams:
"""Build multi-species :class:`LinearParams` from physical inputs.
Note that ``kpar_scale`` is not inferred here either; see
:func:`linear_params_for_geometry`.
"""
rows = tuple(species)
if not rows:
raise ValueError("species must contain at least one kinetic species")
def array(name: str) -> jnp.ndarray:
values = [getattr(row, name) for row in rows]
return jnp.asarray(np.asarray(values, dtype=float))
charge = array("charge")
mass = array("mass")
temperature = array("temperature")
return LinearParams(
charge_sign=charge,
density=array("density"),
mass=mass,
temp=temperature,
vth=jnp.sqrt(temperature / mass),
rho=jnp.sqrt(temperature * mass) / jnp.abs(charge),
tz=temperature / charge,
tprim=array("tprim"),
fprim=array("fprim"),
nu=array("nu"),
closure=array("closure"),
**overrides,
)
[docs]
@jax.tree_util.register_pytree_node_class
@dataclass(frozen=True)
class LinearTerms:
"""Switches for linear-operator components (1.0 = on, 0.0 = off)."""
streaming: float = 1.0
mirror: float = 1.0
curvature: float = 1.0
gradb: float = 1.0
diamagnetic: float = 1.0
collisions: float = 1.0
hypercollisions: float = 1.0
hyperdiffusion: float = 0.0
end_damping: float = 1.0
apar: float = 1.0
bpar: float = 1.0
def tree_flatten(self):
children = (
self.streaming,
self.mirror,
self.curvature,
self.gradb,
self.diamagnetic,
self.collisions,
self.hypercollisions,
self.hyperdiffusion,
self.end_damping,
self.apar,
self.bpar,
)
return children, None
@classmethod
def tree_unflatten(cls, aux_data, children):
return cls(*children)
[docs]
def linear_terms_to_term_config(
terms: LinearTerms | None,
*,
nonlinear: float = 0.0,
) -> TermConfig:
"""Convert :class:`LinearTerms` into the modular :class:`TermConfig`."""
from gkx.terms.config import TermConfig
term_weights = terms if terms is not None else LinearTerms()
return TermConfig(
streaming=term_weights.streaming,
mirror=term_weights.mirror,
curvature=term_weights.curvature,
gradb=term_weights.gradb,
diamagnetic=term_weights.diamagnetic,
collisions=term_weights.collisions,
hypercollisions=term_weights.hypercollisions,
hyperdiffusion=term_weights.hyperdiffusion,
end_damping=term_weights.end_damping,
apar=term_weights.apar,
bpar=term_weights.bpar,
nonlinear=nonlinear,
)
[docs]
def term_config_to_linear_terms(term_cfg: TermConfig | None) -> LinearTerms:
"""Convert modular :class:`TermConfig` into linear-only term weights."""
from gkx.terms.config import TermConfig
cfg = term_cfg if term_cfg is not None else TermConfig()
return LinearTerms(
streaming=cfg.streaming,
mirror=cfg.mirror,
curvature=cfg.curvature,
gradb=cfg.gradb,
diamagnetic=cfg.diamagnetic,
collisions=cfg.collisions,
hypercollisions=cfg.hypercollisions,
hyperdiffusion=cfg.hyperdiffusion,
end_damping=cfg.end_damping,
apar=cfg.apar,
bpar=cfg.bpar,
)
def _is_tracer(x) -> bool:
return isinstance(x, jax.core.Tracer)
def _x64_enabled() -> bool:
return bool(getattr(jax.config, "jax_enable_x64", False))
def _host_values_or_none(x) -> np.ndarray | None:
"""Return ``x`` as host values, or ``None`` when it is genuinely traced.
Read the argument itself, never a ``jnp`` copy of it. Inside a trace every
``jnp`` call stages out, so ``jnp.asarray(-1.0)`` is a tracer even for a
Python constant; a guard that asks after that round trip sees a tracer for
every input and stops validating anything at all under ``jit``. ``numpy``
answers from the buffer, so a concrete argument is still checked inside a
trace and only a value that really is traced goes unchecked -- which it
must, since its sign is not known until the trace is run.
"""
if _is_tracer(x):
return None
return np.asarray(x)
def _check_positive(x, name: str) -> None:
arr = _host_values_or_none(x)
if arr is None:
return
if np.any(arr <= 0.0):
raise ValueError(f"{name} must be > 0")
def _check_nonnegative(x, name: str) -> None:
arr = _host_values_or_none(x)
if arr is None:
return
if np.any(arr < 0.0):
raise ValueError(f"{name} must be >= 0")
[docs]
def _as_species_array(value: float | jnp.ndarray, ns: int, name: str) -> jnp.ndarray:
"""Ensure a parameter is a 1D array of length ns for multi-species handling."""
arr = jnp.asarray(value)
if arr.ndim == 0:
arr = arr[None]
if arr.size == 1:
return jnp.broadcast_to(arr, (ns,))
if int(arr.size) != int(ns):
raise ValueError(f"{name} must have length {ns} (got {arr.size})")
return arr
Preconditioner = Callable[[jnp.ndarray], jnp.ndarray]
PreconditionerSpec = Preconditioner | str | None
def _resolve_implicit_preconditioner(
preconditioner: PreconditionerSpec,
) -> PreconditionerSpec:
if preconditioner is None:
return "auto"
if isinstance(preconditioner, str):
return preconditioner.strip().lower()
return preconditioner
#: Operators applied inside the linear RHS from the cache (every solver path);
#: the others are moment operators handed to the fixed-step integrator.
#: ``coulomb`` is a compatibility alias of ``landau``, which reproduces the
#: retired eight-moment table of Frei, Ernst & Ricci (2022) at half its nu.
BUILTIN_COLLISION_OPERATORS: tuple[str, ...] = (
"none",
"lenard_bernstein",
"landau",
"coulomb",
)
COLLISION_OPERATOR_NAMES: tuple[str, ...] = (
"none",
"lenard_bernstein",
"sugama",
"improved_sugama",
"coulomb",
"coulomb_finite_kperp",
"landau",
)