"""Linear-operator cache data model."""
from __future__ import annotations
from dataclasses import dataclass, field as dataclass_field
from typing import Any
import jax
import jax.numpy as jnp
[docs]
@dataclass(frozen=True)
class HermiteWindow:
"""Global Hermite coordinates of a locally held slab of moments.
A Hermite-sharded route hands the serial kernels a slab that is a window
onto the global ``m`` axis, optionally widened by ghost rows. Every operator
whose coefficient is a function of the *global* moment index -- the field
drives at ``m = 0..3``, the closure at the last moment, the hypercollision
normalization -- needs the global coordinates rather than the slab's own
``arange``. ``index`` may be a tracer (the offset is a per-shard quantity),
``total`` is static because it selects Python branches.
"""
index: jnp.ndarray
total: int
[docs]
def hermite_index_of(window: HermiteWindow | None, local_count: int) -> jnp.ndarray:
"""Return the global Hermite index of every row of a local slab."""
if window is None:
return jnp.arange(int(local_count), dtype=jnp.int32)
return jnp.asarray(window.index, dtype=jnp.int32)
[docs]
def hermite_total_of(window: HermiteWindow | None, local_count: int) -> int:
"""Return the global Hermite extent behind a local slab."""
return int(local_count) if window is None else int(window.total)
[docs]
@jax.tree_util.register_pytree_node_class
@dataclass(frozen=True)
class LinearCache:
"""Precomputed arrays for the linear operator."""
Jl: jnp.ndarray
b: jnp.ndarray
kperp2: jnp.ndarray
kperp2_bmag: bool
bmag: jnp.ndarray
omega_d: jnp.ndarray
cv_d: jnp.ndarray
gb_d: jnp.ndarray
bgrad: jnp.ndarray
jacobian: jnp.ndarray
mask0: jnp.ndarray
dz: jnp.ndarray
kz: jnp.ndarray
ky: jnp.ndarray
kx: jnp.ndarray
kx_grid: jnp.ndarray
ky_grid: jnp.ndarray
dealias_mask: jnp.ndarray
kxfac: jnp.ndarray
lb_lam: jnp.ndarray
collision_lam: jnp.ndarray
hyper_ratio: jnp.ndarray
ratio_l: jnp.ndarray
ratio_m: jnp.ndarray
ratio_lm: jnp.ndarray
mask_const: jnp.ndarray
mask_kz: jnp.ndarray
m_pow: jnp.ndarray
m_norm_kz_factor: jnp.ndarray
damp_profile: jnp.ndarray
linked_damp_profile: jnp.ndarray
l: jnp.ndarray # noqa: E741 - public cache field for the Laguerre index.
m: jnp.ndarray
l4: jnp.ndarray
sqrt_m: jnp.ndarray
sqrt_m_p1: jnp.ndarray
sqrt_p: jnp.ndarray
sqrt_m_ladder: jnp.ndarray
JlB: jnp.ndarray
laguerre_to_grid: jnp.ndarray
laguerre_to_spectral: jnp.ndarray
laguerre_roots: jnp.ndarray
laguerre_j0: jnp.ndarray
laguerre_j1_over_alpha: jnp.ndarray
kx_link_plus: jnp.ndarray
kx_link_minus: jnp.ndarray
kx_link_mask_plus: jnp.ndarray
kx_link_mask_minus: jnp.ndarray
linked_inverse_permutation: jnp.ndarray = dataclass_field(
default_factory=lambda: jnp.asarray([], dtype=jnp.int32)
)
linked_gather_map: jnp.ndarray = dataclass_field(
default_factory=lambda: jnp.asarray([], dtype=jnp.int32)
)
linked_gather_mask: jnp.ndarray = dataclass_field(
default_factory=lambda: jnp.asarray([], dtype=bool)
)
linked_full_cover: bool = False
linked_use_gather: bool = False
linked_indices: tuple[jnp.ndarray, ...] = ()
linked_kz: tuple[jnp.ndarray, ...] = ()
use_twist_shift: bool = False
jtwist: int = 0
#: Length of the two-sided ``ky`` axis behind ``ky``, or ``None`` when the
#: cached rows are not a complete axis. Carried from the grid because a
#: half-spectrum block cannot say how long its own full axis is, and the
#: reduction weight of an even grid's Nyquist row depends on the answer
#: (:mod:`gkx.core_ky_layout`).
ny_full: int | None = None
#: The parent grid's hyperdiffusion ``ky`` cutoff in cache units, or
#: ``None`` when ``ky`` holds that row itself (see
#: :attr:`gkx.core_grid.SpectralGrid.ky_cut`).
ky_cut: jnp.ndarray | None = None
#: Non-twisting flux tube radial index shift ``m0(ky, z)``: row ``kx`` at
#: ``z`` holds ballooning mode ``kx + m0``. ``None`` off that geometry.
ntft_m0: jnp.ndarray | None = None
#: Block-sparse unit-frequency linearized Landau factors
#: (:class:`~gkx.operators.linear.collisions.LandauBlocks`); ``None`` when
#: the run uses the Dougherty kernel.
landau_pairs: Any = None
#: True when the Laguerre axis holds nodes of the mu grid and
#: ``laguerre_roots`` their ``mu_hat`` values (``LinearParams.mu_vperp_max``).
mu_grid: bool = False
#: Nodal Lenard-Bernstein perpendicular rate ``(Nmu, Nmu, Nz)`` (mu grid only).
mu_lb: Any = None
#: Per-species moment counts (``LinearParams.moments``): one exact-size cache per run of
#: species sharing a ``(Nl, Nm)``, in species order; ``None`` when every species keeps
#: the state's. The RHS runs each group at its own size.
groups: tuple["LinearCache", ...] | None = None
def tree_flatten(self):
children = (
self.Jl,
self.b,
self.kperp2,
self.bmag,
self.omega_d,
self.cv_d,
self.gb_d,
self.bgrad,
self.jacobian,
self.mask0,
self.dz,
self.kz,
self.ky,
self.kx,
self.kx_grid,
self.ky_grid,
self.dealias_mask,
self.kxfac,
self.lb_lam,
self.collision_lam,
self.hyper_ratio,
self.ratio_l,
self.ratio_m,
self.ratio_lm,
self.mask_const,
self.mask_kz,
self.m_pow,
self.m_norm_kz_factor,
self.damp_profile,
self.linked_damp_profile,
self.l,
self.m,
self.l4,
self.sqrt_m,
self.sqrt_m_p1,
self.sqrt_p,
self.sqrt_m_ladder,
self.JlB,
self.laguerre_to_grid,
self.laguerre_to_spectral,
self.laguerre_roots,
self.laguerre_j0,
self.laguerre_j1_over_alpha,
self.kx_link_plus,
self.kx_link_minus,
self.kx_link_mask_plus,
self.kx_link_mask_minus,
self.linked_inverse_permutation,
self.linked_gather_map,
self.linked_gather_mask,
)
linked_idx = self.linked_indices or ()
linked_kz = self.linked_kz or ()
children = (
children
+ tuple(linked_idx)
+ tuple(linked_kz)
+ (self.ky_cut, self.ntft_m0, self.landau_pairs, self.mu_lb, self.groups)
)
aux_data = (
self.kperp2_bmag,
self.use_twist_shift,
self.jtwist,
len(linked_idx),
len(linked_kz),
self.linked_full_cover,
self.linked_use_gather,
self.ny_full,
self.mu_grid,
)
return children, aux_data
@classmethod
def tree_unflatten(cls, aux_data, children):
(
kperp2_bmag,
use_twist_shift,
jtwist,
n_linked_idx,
n_linked_kz,
linked_full_cover,
linked_use_gather,
ny_full,
mu_grid,
) = aux_data
base_count = 50
base_children = list(children[:base_count])
base_children.insert(3, kperp2_bmag)
linked_idx = tuple(children[base_count : base_count + n_linked_idx])
linked_kz = tuple(
children[
base_count + n_linked_idx : base_count + n_linked_idx + n_linked_kz
]
)
return cls(
*base_children,
linked_indices=linked_idx,
linked_kz=linked_kz,
use_twist_shift=use_twist_shift,
jtwist=jtwist,
linked_full_cover=linked_full_cover,
linked_use_gather=linked_use_gather,
ny_full=ny_full,
mu_grid=mu_grid,
ky_cut=children[-5],
ntft_m0=children[-4],
landau_pairs=children[-3],
mu_lb=children[-2],
groups=children[-1],
)