"""Spatial and velocity-ladder kernels for parallel streaming."""
from __future__ import annotations
import jax
import jax.numpy as jnp
import numpy as np
from gkx.core_ky_layout import is_half, nyc_from_ny
from gkx.core_velocity import hermite_ladder_coeffs
# One positivity guard for the whole linear operator. The local copy asked
# whether a ``jnp`` round trip of its argument was traced, which is true of every
# argument inside a trace, so the streaming kernels stopped checking ``dz`` and
# ``vth`` under ``jit`` at all. Sharing the guard keeps one answer to the
# question rather than two that can drift apart again.
from gkx.operators.linear.params import _check_positive, _is_tracer # noqa: F401
def _fft_ik_multiplier(kz: jnp.ndarray, like: jnp.ndarray) -> jnp.ndarray:
"""Return ``1j * kz`` with the same complex dtype as an FFT array."""
real_dtype = jnp.real(like).dtype
kz_cast = jnp.asarray(kz, dtype=real_dtype).astype(like.dtype)
return jnp.asarray(1j, dtype=like.dtype) * kz_cast
def _fft_abs_multiplier(kz: jnp.ndarray, like: jnp.ndarray) -> jnp.ndarray:
"""Return ``abs(kz)`` as a complex multiplier matching an FFT array."""
real_dtype = jnp.real(like).dtype
return jnp.asarray(jnp.abs(kz), dtype=real_dtype).astype(like.dtype)
def _ifft_last(x: jnp.ndarray) -> jnp.ndarray:
"""``jnp.fft.ifft`` over the last axis as a forward FFT of the conjugate.
XLA:GPU follows every inverse transform with a separate scaling kernel; the
``1 / n`` here fuses into the neighbouring elementwise work instead.
"""
return jnp.conj(jnp.fft.fft(jnp.conj(x), axis=-1)) * (1.0 / x.shape[-1])
[docs]
def grad_z_periodic(
f: jnp.ndarray,
dz: float | jnp.ndarray | None = None,
kz: jnp.ndarray | None = None,
*,
ny_full: int | None = None,
) -> jnp.ndarray:
"""Spectral periodic derivative along the last axis.
``ny_full`` marks a complete real-FFT layout: see
:func:`_real_fft_nyquist_derivative`.
"""
if kz is None:
if dz is None:
raise ValueError("Either dz or kz must be provided")
_check_positive(dz, "dz")
n = f.shape[-1]
if kz is None:
dz_val = jnp.asarray(dz, dtype=jnp.real(f).dtype)
kz = 2.0 * jnp.pi * jnp.fft.fftfreq(n, d=dz_val)
f_hat = jnp.fft.fft(f, axis=-1)
df_hat = _fft_ik_multiplier(kz, f_hat) * f_hat
return _real_fft_nyquist_derivative(_ifft_last(df_hat), ny_full)
def _real_fft_nyquist_derivative(out: jnp.ndarray, ny_full: int | None) -> jnp.ndarray:
"""Zero the even-``Nz`` Nyquist derivative on self-conjugate real-FFT rows.
Rows ``ky=0`` (and ``ky=Ny/2``) are real in ``(x, z)``: the sampled Nyquist
cosine has zero derivative, but ``i k_nyq`` makes it imaginary (#328).
Generic complex arrays (``ny_full=None``) keep the signed derivative.
"""
if (
ny_full is None
or out.ndim < 3
or out.shape[-1] % 2
or out.shape[-3] not in (int(ny_full), nyc_from_ny(ny_full))
):
return out
ny = int(ny_full)
alternating = jnp.asarray((-1.0) ** np.arange(out.shape[-1]), jnp.real(out).dtype)
rows = [0] + ([ny // 2] if ny % 2 == 0 and ny // 2 < out.shape[-3] else [])
def component(part):
return jnp.mean(part * alternating, axis=-1, keepdims=True) * alternating
def scatter(out): # XLA:CPU turns a concatenate here into a slow copy fusion
return out.at[..., rows, :, :].add(-component(out[..., rows, :, :]))
def concatenate(out):
pieces, start = [], 0
for row in rows:
part = out[..., row : row + 1, :, :]
pieces += [out[..., start:row, :, :], part - component(part)]
start = row + 1
return jnp.concatenate([*pieces, out[..., start:, :, :]], axis=-3)
return jax.lax.platform_dependent(out, cpu=scatter, default=concatenate)
def _shift_kx_linked(
f: jnp.ndarray,
kx_link: jnp.ndarray,
kx_mask: jnp.ndarray,
) -> jnp.ndarray:
"""Shift along kx for each ky using precomputed link indices."""
f_ky = jnp.moveaxis(f, -2, 0)
kx_mask = kx_mask.astype(f.dtype)
def _gather_ky(
f_slice: jnp.ndarray, idx: jnp.ndarray, mask: jnp.ndarray
) -> jnp.ndarray:
gathered = jnp.take(f_slice, idx, axis=-1)
return gathered * mask
shifted = jax.vmap(_gather_ky, in_axes=(0, 0, 0))(f_ky, kx_link, kx_mask)
return jnp.moveaxis(shifted, 0, -2)
def _grad_z_linked_fd(
f: jnp.ndarray,
dz: float | jnp.ndarray,
kx_link_plus: jnp.ndarray,
kx_link_minus: jnp.ndarray,
kx_mask_plus: jnp.ndarray,
kx_mask_minus: jnp.ndarray,
) -> jnp.ndarray:
"""Finite-difference z-derivative with twist-shift kx linking at the ends."""
_check_positive(dz, "dz")
dz_val = jnp.asarray(dz, dtype=jnp.real(f).dtype)
f_z0 = f[..., 0]
f_zm1 = f[..., -1]
f_z0_shift = _shift_kx_linked(f_z0, kx_link_plus, kx_mask_plus)
f_zm1_shift = _shift_kx_linked(f_zm1, kx_link_minus, kx_mask_minus)
f_roll_p1 = jnp.concatenate([f[..., 1:], f_z0_shift[..., None]], axis=-1)
f_roll_m1 = jnp.concatenate([f_zm1_shift[..., None], f[..., :-1]], axis=-1)
return (f_roll_p1 - f_roll_m1) / (2.0 * dz_val)
def _reverse_from_one(x: jnp.ndarray, axis: int) -> jnp.ndarray:
"""Return ``x`` reindexed as ``[0, n-1, n-2, ..., 1]`` along ``axis``.
This is the conjugate partner index for a real-FFT half spectrum. Written as
``jnp.take`` with that index vector, XLA lowers it as a general gather and
materialises the result: profiling on 2026-09-01 attributed a transpose, a
full-array copy, the gather and a transpose back to each such call. Written
as a slice, a reverse and a concatenate the same permutation fuses, and
measures 9.6 ms against 13.0 ms on a ``(2,4,8,96,96,48)`` complex64 state
with byte-identical output.
"""
head = jax.lax.slice_in_dim(x, 0, 1, axis=axis)
tail = jax.lax.rev(jax.lax.slice_in_dim(x, 1, x.shape[axis], axis=axis), (axis,))
return jnp.concatenate([head, tail], axis=axis)
def _restore_linked_real_fft_conjugates(
out: jnp.ndarray,
*,
covered_rows: jnp.ndarray,
ny_full: int | None = None,
) -> jnp.ndarray:
"""Restore the conjugate ``-ky`` rows on a full real-FFT spectral grid.
The linked-FFT chains are built on the unique dealiased positive-``ky``
block. When the runtime carries the full real-FFT-expanded ``ky`` layout, the
untouched negative rows must be reconstructed by real-FFT conjugate
symmetry so the linked derivative acts on the physical Hermitian state.
On a half-spectrum state (plan 5.3 N3) there are no negative rows and this
is the identity, returning the *same object* so the graph is unchanged. The
fill must not merely be skipped but recognized as inapplicable: ``(-j) % Nyc``
is a different positive row, not a partner, so running the two-sided rule on
a half axis would conjugate-mirror one physical mode onto another. Telling
the two apart needs ``ny_full``; see :mod:`gkx.core_ky_layout`.
"""
Ny = out.shape[-3]
if Ny <= 1 or is_half(Ny, ny_full):
return out
row_idx = jnp.arange(Ny, dtype=jnp.int32)
src_rows = jnp.mod(-row_idx, Ny)
covered = jnp.asarray(covered_rows, dtype=bool)
source_covered = jnp.take(covered, src_rows, axis=0)
fill_mask = (~covered) & source_covered & (row_idx != 0)
mirrored = _reverse_from_one(out, axis=out.ndim - 3)
Nx = out.shape[-2]
if Nx > 1:
mirrored = _reverse_from_one(mirrored, axis=mirrored.ndim - 2)
mirrored = jnp.conj(mirrored)
mask_shape = (1,) * (out.ndim - 3) + (Ny, 1, 1)
return jnp.where(fill_mask.reshape(mask_shape), mirrored, out)
def _validate_linked_fft_inputs(
linked_indices: tuple[jnp.ndarray, ...],
linked_kz: tuple[jnp.ndarray, ...],
*,
operator: str | tuple[str, ...],
) -> None:
operators = (operator,) if isinstance(operator, str) else tuple(operator)
if not operators or any(op not in {"grad", "abs"} for op in operators):
raise ValueError(f"unsupported linked FFT operator {operator!r}")
if len(linked_indices) != len(linked_kz):
raise ValueError("linked_indices and linked_kz must have the same length")
if not linked_indices:
suffix = "derivative" if operator == "grad" else "operator"
raise ValueError(f"linked_indices cannot be empty for linked FFT {suffix}")
def _flatten_linked_fft_state(
f: jnp.ndarray,
) -> tuple[jnp.ndarray, tuple[int, ...], int, int, int]:
"""Flatten ``(ky, kx)`` into one mode axis in the state's own row order.
Row ``p`` is mode ``(ky, kx) = (p // Nx, p % Nx)``. A reshape keeps the
state buffer where it is; swapping ``kx`` in front of ``ky`` first, as the
chain index convention ``ky + Ny * kx`` would suggest, made XLA copy the
whole state before every gather and again after every scatter. The chain
maps are translated to row order by :func:`_state_mode_rows` instead.
"""
Ny = f.shape[-3]
Nx = f.shape[-2]
Nz = f.shape[-1]
lead_shape = f.shape[:-3]
return f.reshape(*lead_shape, Ny * Nx, Nz), lead_shape, Ny, Nx, Nz
def _state_mode_rows(idx_flat: jnp.ndarray, *, Ny: int, Nx: int) -> jnp.ndarray:
"""Map chain mode indices ``ky + Ny * kx`` to state rows ``ky * Nx + kx``."""
idx = jnp.asarray(idx_flat, dtype=jnp.int32)
return jnp.mod(idx, Ny) * Nx + idx // Ny
def _state_row_table(table: jnp.ndarray, *, Ny: int, Nx: int) -> jnp.ndarray:
"""Reorder a per-mode table indexed ``ky + Ny * kx`` into state row order."""
rows = np.arange(Ny * Nx)
return jnp.take(table, rows // Nx + Ny * (rows % Nx), axis=0)
def _scatter_unique_linked_modes(
target: jnp.ndarray,
idx_flat: jnp.ndarray,
updates: jnp.ndarray,
) -> jnp.ndarray:
"""Scatter chain updates into unique linked ``(kx, ky)`` rows."""
idx = jnp.asarray(idx_flat, dtype=jnp.int32)
target_t = jnp.moveaxis(target, -2, 0)
updates_t = jnp.moveaxis(updates, -2, 0)
idx = idx[:, None]
dnums = jax.lax.ScatterDimensionNumbers(
update_window_dims=tuple(range(1, updates_t.ndim)),
inserted_window_dims=(0,),
scatter_dims_to_operand_dims=(0,),
)
out_t = jax.lax.scatter(
target_t,
idx,
updates_t,
dnums,
unique_indices=True,
)
return jnp.moveaxis(out_t, 0, -2)
def _linked_fft_multiplier(
kz_link: jnp.ndarray,
f_hat: jnp.ndarray,
operator: str | tuple[str, ...],
) -> jnp.ndarray:
"""Return the spectral multiplier of one chain class.
A tuple names one operator per slot of the leading axis of ``f_hat``, so
stacked operands share the class's gather, transforms and scatter while
each slot keeps its own ``i k_z`` or ``|k_z|``.
"""
def single(op: str) -> jnp.ndarray:
if op == "grad":
return _fft_ik_multiplier(kz_link, f_hat)
return _fft_abs_multiplier(kz_link, f_hat)
if isinstance(operator, str):
return single(operator)
stacked = jnp.stack([single(op) for op in operator])
return stacked.reshape((len(operator),) + (1,) * (f_hat.ndim - 2) + (-1,))
def _linked_fft_chain_update(
f_flat: jnp.ndarray | tuple[jnp.ndarray, ...],
idx_map: jnp.ndarray,
kz_link: jnp.ndarray,
*,
lead_shape: tuple[int, ...],
Ny: int,
Nx: int,
Nz: int,
operator: str | tuple[str, ...],
) -> tuple[jnp.ndarray, jnp.ndarray]:
"""Apply one linked-chain FFT derivative or absolute-``k_z`` operator.
A tuple of flattened operands is gathered one by one and stacked in chain
layout, so the class issues one FFT and one IFFT for all of them without
first writing the stacked full states.
"""
if idx_map.ndim != 2:
raise ValueError("linked index maps must have shape (nChains, nLinks)")
nChains, nLinks = idx_map.shape
idx_flat = idx_map.reshape(-1)
rows = _state_mode_rows(idx_flat, Ny=Ny, Nx=Nx)
if isinstance(f_flat, tuple):
dtype = f_flat[0].dtype
f_link = jnp.stack([jnp.take(part, rows, axis=-2) for part in f_flat])
else:
dtype = f_flat.dtype
f_link = jnp.take(f_flat, rows, axis=-2)
f_link = f_link.reshape(*lead_shape, nChains, nLinks * Nz)
f_hat = jnp.fft.fft(f_link, axis=-1)
multiplier = _linked_fft_multiplier(kz_link, f_hat, operator)
df_hat = multiplier * f_hat
df_link = _ifft_last(df_hat)
df_link = df_link.reshape(*lead_shape, nChains * nLinks, Nz)
return idx_flat, jnp.asarray(df_link, dtype=dtype)
def _linked_fft_chain_outputs(
f_flat: jnp.ndarray | tuple[jnp.ndarray, ...],
linked_indices: tuple[jnp.ndarray, ...],
linked_kz: tuple[jnp.ndarray, ...],
*,
lead_shape: tuple[int, ...],
Ny: int,
Nx: int,
Nz: int,
operator: str | tuple[str, ...],
) -> tuple[list[jnp.ndarray], list[jnp.ndarray]]:
chain_updates: list[jnp.ndarray] = []
chain_indices: list[jnp.ndarray] = []
for idx_map, kz_link in zip(linked_indices, linked_kz):
idx_flat, df_link = _linked_fft_chain_update(
f_flat,
idx_map,
kz_link,
lead_shape=lead_shape,
Ny=Ny,
Nx=Nx,
Nz=Nz,
operator=operator,
)
chain_indices.append(idx_flat)
chain_updates.append(df_link)
return chain_indices, chain_updates
def _linked_fft_covered_rows(
chain_indices: list[jnp.ndarray],
*,
Ny: int,
) -> jnp.ndarray:
idx_cat = (
chain_indices[0]
if len(chain_indices) == 1
else jnp.concatenate(chain_indices, axis=0)
)
return jnp.zeros((Ny,), dtype=bool).at[jnp.mod(idx_cat, Ny)].set(True)
def _linked_fft_gather_output(
chain_updates: list[jnp.ndarray],
*,
linked_gather_map: jnp.ndarray | None,
linked_gather_mask: jnp.ndarray | None,
lead_shape: tuple[int, ...],
Nx: int,
Ny: int,
Nz: int,
) -> jnp.ndarray:
updates_cat = jnp.concatenate(chain_updates, axis=-2)
gather_map = _state_row_table(
jnp.asarray(linked_gather_map, dtype=jnp.int32), Ny=Ny, Nx=Nx
)
gather_mask = _state_row_table(
jnp.asarray(linked_gather_mask, dtype=updates_cat.dtype), Ny=Ny, Nx=Nx
)
updates_full = jnp.take(updates_cat, gather_map, axis=-2)
mask_shape = (1,) * (updates_full.ndim - 2) + (gather_mask.shape[0], 1)
updates_full = updates_full * gather_mask.reshape(mask_shape)
return updates_full.reshape(*lead_shape, Ny, Nx, Nz)
def _linked_fft_full_cover_output(
chain_updates: list[jnp.ndarray],
*,
linked_inverse_permutation: jnp.ndarray | None,
lead_shape: tuple[int, ...],
Nx: int,
Ny: int,
Nz: int,
) -> jnp.ndarray:
if linked_inverse_permutation is None:
raise ValueError(
"linked_inverse_permutation required when linked_full_cover is True"
)
updates_cat = jnp.concatenate(chain_updates, axis=-2)
inv = _state_row_table(
jnp.asarray(linked_inverse_permutation, dtype=jnp.int32), Ny=Ny, Nx=Nx
)
df_flat = jnp.take(updates_cat, inv, axis=-2)
return df_flat.reshape(*lead_shape, Ny, Nx, Nz)
def _linked_fft_scatter_output(
chain_indices: list[jnp.ndarray],
chain_updates: list[jnp.ndarray],
*,
lead_shape: tuple[int, ...],
Nx: int,
Ny: int,
Nz: int,
) -> jnp.ndarray:
df_flat = jnp.zeros((*lead_shape, Ny * Nx, Nz), dtype=chain_updates[0].dtype)
for idx_flat, df_link in zip(chain_indices, chain_updates):
df_flat = _scatter_unique_linked_modes(
df_flat, _state_mode_rows(idx_flat, Ny=Ny, Nx=Nx), df_link
)
return df_flat.reshape(*lead_shape, Ny, Nx, Nz)
def _shift_kx_modes(f: jnp.ndarray, shift: jnp.ndarray) -> jnp.ndarray:
"""Return ``out[ky, kx, z] = f[ky, kx + shift[ky, z], z]`` in signed modes.
Radial mode numbers follow the FFT order; a source outside the resolved
band ``|n| <= (Nx - 1) // 2`` contributes zero.
"""
nx = f.shape[-2]
modes = np.fft.fftfreq(nx, d=1.0 / nx).astype(np.int32)
source = modes[None, :, None] + jnp.asarray(shift, jnp.int32)[:, None, :]
valid = jnp.abs(source) <= (nx - 1) // 2
index = jnp.broadcast_to(jnp.mod(source, nx), f.shape)
out = jnp.take_along_axis(f, index, axis=-2)
return jnp.where(valid, out, jnp.zeros((), dtype=f.dtype))
def _linked_fft_apply(
f: jnp.ndarray | tuple[jnp.ndarray, ...],
linked_indices: tuple[jnp.ndarray, ...],
linked_kz: tuple[jnp.ndarray, ...],
*,
operator: str | tuple[str, ...],
linked_inverse_permutation: jnp.ndarray | None = None,
linked_full_cover: bool = False,
linked_gather_map: jnp.ndarray | None = None,
linked_gather_mask: jnp.ndarray | None = None,
linked_use_gather: bool = False,
ny_full: int | None = None,
ntft_m0: jnp.ndarray | None = None,
) -> jnp.ndarray:
"""Apply linked-chain spectral operators; a tuple ``f`` returns them stacked.
On a non-twisting flux tube the state row ``kx`` at ``z`` holds the
ballooning mode ``kx + m0(ky, z)``, so a parallel chain keeps
``kx + m0`` fixed (not ``kx``): the operand is moved to that frame before
the chain transform and the result moved back.
With a tuple of equally shaped operands and one operator per operand, the
result carries a new leading axis whose slot ``i`` equals
``_linked_fft_apply(f[i], ..., operator=operator[i])``: every chain class
then issues one gather-stack, FFT, IFFT and output write for all operands.
"""
_validate_linked_fft_inputs(linked_indices, linked_kz, operator=operator)
if ntft_m0 is not None:
shifted = _linked_fft_apply(
tuple(_shift_kx_modes(part, -ntft_m0) for part in f)
if isinstance(f, tuple)
else _shift_kx_modes(f, -ntft_m0),
linked_indices,
linked_kz,
operator=operator,
linked_inverse_permutation=linked_inverse_permutation,
linked_full_cover=linked_full_cover,
linked_gather_map=linked_gather_map,
linked_gather_mask=linked_gather_mask,
linked_use_gather=linked_use_gather,
ny_full=ny_full,
)
return _shift_kx_modes(shifted, ntft_m0)
if isinstance(f, tuple):
if isinstance(operator, str) or len(operator) != len(f):
raise ValueError("stacked linked FFT operands need one operator each")
if len({(part.shape, part.dtype) for part in f}) != 1:
raise ValueError("stacked linked FFT operands must share shape and dtype")
parts = [_flatten_linked_fft_state(part) for part in f]
f_flat: jnp.ndarray | tuple[jnp.ndarray, ...] = tuple(p[0] for p in parts)
_flat0, part_lead, Ny, Nx, Nz = parts[0]
lead_shape = (len(f), *part_lead)
else:
f_flat, lead_shape, Ny, Nx, Nz = _flatten_linked_fft_state(f)
chain_indices, chain_updates = _linked_fft_chain_outputs(
f_flat,
linked_indices,
linked_kz,
lead_shape=lead_shape,
Ny=Ny,
Nx=Nx,
Nz=Nz,
operator=operator,
)
covered_rows = _linked_fft_covered_rows(chain_indices, Ny=Ny)
if linked_use_gather:
out = _linked_fft_gather_output(
chain_updates,
linked_gather_map=linked_gather_map,
linked_gather_mask=linked_gather_mask,
lead_shape=lead_shape,
Nx=Nx,
Ny=Ny,
Nz=Nz,
)
elif linked_full_cover:
out = _linked_fft_full_cover_output(
chain_updates,
linked_inverse_permutation=linked_inverse_permutation,
lead_shape=lead_shape,
Nx=Nx,
Ny=Ny,
Nz=Nz,
)
else:
out = _linked_fft_scatter_output(
chain_indices,
chain_updates,
lead_shape=lead_shape,
Nx=Nx,
Ny=Ny,
Nz=Nz,
)
# Materialized before the conjugate restore: unbarriered, XLA:CPU fuses the
# output gather into the restore and the derivative and recomputes it (2.4x
# on the stacked streaming and hypercollision pair). On GPU it costs 30%.
out = jax.lax.platform_dependent(
out, cpu=jax.lax.optimization_barrier, default=lambda x: x
)
out = _restore_linked_real_fft_conjugates(
out, covered_rows=covered_rows, ny_full=ny_full
)
if isinstance(operator, str):
return _real_fft_nyquist_derivative(out, ny_full) if operator == "grad" else out
return jnp.stack(
[
_real_fft_nyquist_derivative(part, ny_full) if op == "grad" else part
for part, op in zip(out, operator)
]
)
def abs_z_periodic(
f: jnp.ndarray, dz: float | jnp.ndarray | None = None, kz: jnp.ndarray | None = None
) -> jnp.ndarray:
"""Apply ``|k_z|`` spectrally along the last axis, periodic case.
The companion of :func:`grad_z_periodic` for operators that need the
magnitude of the parallel wavenumber rather than the signed derivative,
such as the reflectionless Hermite closure.
"""
if kz is None:
if dz is None:
raise ValueError("Either dz or kz must be provided")
_check_positive(dz, "dz")
kz = (
2.0
* jnp.pi
* jnp.fft.fftfreq(f.shape[-1], d=jnp.asarray(dz, dtype=jnp.real(f).dtype))
)
f_hat = jnp.fft.fft(f, axis=-1)
return _ifft_last(_fft_abs_multiplier(kz, f_hat) * f_hat)
def grad_z_linked_fft(
f: jnp.ndarray,
dz: float | jnp.ndarray,
linked_indices: tuple[jnp.ndarray, ...],
linked_kz: tuple[jnp.ndarray, ...],
linked_inverse_permutation: jnp.ndarray | None = None,
linked_full_cover: bool = False,
linked_gather_map: jnp.ndarray | None = None,
linked_gather_mask: jnp.ndarray | None = None,
linked_use_gather: bool = False,
ny_full: int | None = None,
ntft_m0: jnp.ndarray | None = None,
) -> jnp.ndarray:
"""Spectral z-derivative using linked-chain FFT modes."""
_check_positive(dz, "dz")
return _linked_fft_apply(
f,
linked_indices,
linked_kz,
operator="grad",
linked_inverse_permutation=linked_inverse_permutation,
linked_full_cover=linked_full_cover,
linked_gather_map=linked_gather_map,
linked_gather_mask=linked_gather_mask,
linked_use_gather=linked_use_gather,
ny_full=ny_full,
ntft_m0=ntft_m0,
)
def abs_z_linked_fft(
f: jnp.ndarray,
linked_indices: tuple[jnp.ndarray, ...],
linked_kz: tuple[jnp.ndarray, ...],
linked_inverse_permutation: jnp.ndarray | None = None,
linked_full_cover: bool = False,
linked_gather_map: jnp.ndarray | None = None,
linked_gather_mask: jnp.ndarray | None = None,
linked_use_gather: bool = False,
ny_full: int | None = None,
ntft_m0: jnp.ndarray | None = None,
) -> jnp.ndarray:
"""Apply |kz| in linked-FFT space."""
return _linked_fft_apply(
f,
linked_indices,
linked_kz,
operator="abs",
linked_inverse_permutation=linked_inverse_permutation,
linked_full_cover=linked_full_cover,
linked_gather_map=linked_gather_map,
linked_gather_mask=linked_gather_mask,
linked_use_gather=linked_use_gather,
ny_full=ny_full,
ntft_m0=ntft_m0,
)
[docs]
def shift_axis(arr: jnp.ndarray, offset: int, axis: int) -> jnp.ndarray:
"""Shift an array along an axis with zero padding (non-periodic)."""
axis = axis % arr.ndim
n = arr.shape[axis]
k = min(abs(offset), n)
start = min(max(offset, 0), n)
body = jax.lax.slice_in_dim(arr, start, start + n - k, axis=axis)
pad = [(0, 0)] * arr.ndim
pad[axis] = (0, k) if offset > 0 else (k, 0)
return jnp.pad(body, pad)
[docs]
def apply_hermite_v(G: jnp.ndarray) -> jnp.ndarray:
"""Multiply Hermite coefficients by v_parallel (ladder form)."""
axis_m = -4
Nm = G.shape[axis_m]
sqrt_p, sqrt_m = hermite_ladder_coeffs(Nm - 1)
sqrt_p = sqrt_p[:Nm]
sqrt_m = sqrt_m[:Nm]
G_plus = shift_axis(G, 1, axis_m)
G_minus = shift_axis(G, -1, axis_m)
shape = [1] * G.ndim
shape[axis_m] = Nm
sqrt_p = sqrt_p.reshape(shape)
sqrt_m = sqrt_m.reshape(shape)
return sqrt_p * G_plus + sqrt_m * G_minus
[docs]
def apply_hermite_v2(G: jnp.ndarray) -> jnp.ndarray:
"""Multiply Hermite coefficients by v_parallel^2."""
return apply_hermite_v(apply_hermite_v(G))
[docs]
def apply_laguerre_x(G: jnp.ndarray) -> jnp.ndarray:
"""Multiply Laguerre coefficients by the perpendicular energy variable.
This helper acts in the plain Laguerre basis, where
``x L_l = (2l+1) L_l - (l+1) L_{l+1} - l L_{l-1}`` gives negative
off-diagonal couplings. The runtime state stores Laguerre coefficients in
the alternating ``(-1)**l`` basis (the convention of
``gkx.core_velocity.J_l_all``), which flips the sign of both off-diagonal
terms; composing with runtime states requires the alternating-basis
counterparts used by the production RHS (see the positive-neighbor
couplings in ``gkx.terms.linear_terms``, e.g.
``curvature_gradb_contribution``). It is not used by the production RHS.
"""
axis_l = -5
Nl = G.shape[axis_l]
ell = jnp.arange(Nl)
G_plus = shift_axis(G, 1, axis_l)
G_minus = shift_axis(G, -1, axis_l)
l_shape = [1] * G.ndim
l_shape[axis_l] = Nl
l_col = ell.reshape(l_shape)
return (2.0 * l_col + 1.0) * G - (l_col + 1.0) * G_plus - l_col * G_minus
def streaming_ladder_term(
H: jnp.ndarray,
kz: jnp.ndarray,
vth: float | jnp.ndarray,
sqrt_p: jnp.ndarray,
sqrt_m: jnp.ndarray,
*,
dz: float | jnp.ndarray | None = None,
kx_link_plus: jnp.ndarray | None = None,
kx_link_minus: jnp.ndarray | None = None,
kx_mask_plus: jnp.ndarray | None = None,
kx_mask_minus: jnp.ndarray | None = None,
linked_indices: tuple[jnp.ndarray, ...] | None = None,
linked_kz: tuple[jnp.ndarray, ...] | None = None,
linked_inverse_permutation: jnp.ndarray | None = None,
linked_full_cover: bool = False,
linked_gather_map: jnp.ndarray | None = None,
linked_gather_mask: jnp.ndarray | None = None,
linked_use_gather: bool = False,
ny_full: int | None = None,
ntft_m0: jnp.ndarray | None = None,
use_twist_shift: bool = False,
) -> jnp.ndarray:
"""Apply streaming with precomputed Hermite-ladder coefficients."""
_check_positive(vth, "vth")
if use_twist_shift:
if dz is None:
raise ValueError("dz must be provided for twist-shift boundaries")
if linked_indices is not None and linked_kz is not None:
dH_dz = grad_z_linked_fft(
H,
dz=dz,
linked_indices=linked_indices,
linked_kz=linked_kz,
linked_inverse_permutation=linked_inverse_permutation,
linked_full_cover=linked_full_cover,
linked_gather_map=linked_gather_map,
linked_gather_mask=linked_gather_mask,
linked_use_gather=linked_use_gather,
ny_full=ny_full,
ntft_m0=ntft_m0,
)
else:
if kx_link_plus is None or kx_link_minus is None:
raise ValueError(
"kx_link arrays must be provided for twist-shift boundaries"
)
if kx_mask_plus is None or kx_mask_minus is None:
raise ValueError(
"kx_link masks must be provided for twist-shift boundaries"
)
dH_dz = _grad_z_linked_fd(
H,
dz=dz,
kx_link_plus=kx_link_plus,
kx_link_minus=kx_link_minus,
kx_mask_plus=kx_mask_plus,
kx_mask_minus=kx_mask_minus,
)
else:
dH_dz = grad_z_periodic(H, kz=kz, ny_full=ny_full)
axis_m = -4
pad = [(0, 0)] * H.ndim
pad[axis_m] = (1, 1)
H_pad = jnp.pad(dH_dz, pad)
slc_plus = [slice(None)] * H.ndim
slc_minus = [slice(None)] * H.ndim
slc_plus[axis_m] = slice(2, None)
slc_minus[axis_m] = slice(0, -2)
H_plus = H_pad[tuple(slc_plus)]
H_minus = H_pad[tuple(slc_minus)]
return vth * (sqrt_p * H_plus + sqrt_m * H_minus)