"""Local-stencil advection kernels for TRACMASS-JAX.
Implements two kernel factories for memory-efficient integration:
- :func:`piecewise_local` — patch-based advection for
:func:`~tracmass.integrate.time_dependent_patched`. The kernel closes
over only the four scalar grid dimensions; ``kmt`` and ``dxdy`` are read
from the ``LocalFluxFields`` at runtime so that the same compiled kernel
can be reused across time windows.
- :func:`piecewise_stencil` — per-particle stencil advection for
:func:`~tracmass.integrate.time_dependent_streaming`. Reads flux values
from a pre-extracted :class:`~tracmass.gcm.FluxStencil` instead of
indexing into a global flux array.
Both share the piecewise crossing-time and position helpers from
:mod:`~tracmass.advection_scheme.full_grid`.
"""
import functools
import types
import jax
import jax.numpy as jnp
from tracmass.particle import ParticleState, REASON_LAND, REASON_BOUNDS_H, REASON_BOUNDS_V, REASON_DXYZ
from tracmass.advection_scheme.full_grid import (
UNDEF,
EPS,
_cross_time_1d,
_calc_pos_1d,
_subcycle_one_particle,
)
# ===========================================================================
# Patch-based (local) piecewise kernel
# ===========================================================================
[docs]
@functools.lru_cache(maxsize=64)
def piecewise_local(imt: int, jmt: int, km: int, iperio: int,
no_wflux: bool = False):
"""Return a JIT-compiled step function for patch-based (local stencil) advection.
Unlike :func:`~tracmass.advection_scheme.full_grid.piecewise`, this factory
closes over only the four scalar grid dimensions. ``kmt`` and ``dxdy`` are
read from the *fluxes* argument at runtime (a
:class:`~tracmass.gcm.LocalFluxFields`), so the same compiled kernel can be
reused across windows as the patch content changes — without recompilation.
Parameters
----------
imt, jmt, km : int
Patch dimensions (fixed for the lifetime of the run).
iperio : int
East-west periodicity flag (1 = cyclic, 0 = closed).
no_wflux : bool
When True, vertical fluxes are ignored.
"""
@jax.jit
def step_all(fluxes, state: ParticleState, intrpg: float,
dsmin: float) -> ParticleState:
ib, jb, kb = state.ia, state.ja, state.ka
i0 = jnp.clip(ib - 1, 0, imt - 1)
j0 = jnp.clip(jb - 1, 0, jmt - 1)
k0 = jnp.clip(kb - 1, 0, km - 1)
# kmt and dxdy come from the fluxes object (dynamic, not XLA constants)
dxyz_raw = fluxes.dxdy[i0, j0] * (
intrpg * fluxes.dzt[i0, j0, k0, 1] * fluxes.zstot[i0, j0, 1]
+ (1.0 - intrpg) * fluxes.dzt[i0, j0, k0, 0] * fluxes.zstot[i0, j0, 0]
)
zero_volume = state.active & (dxyz_raw <= 0.0)
dxyz = jnp.maximum(dxyz_raw, EPS)
dsmin_arr = dsmin / dxyz
# Pass a lightweight namespace to _subcycle_one_particle: scalars
# (imt, jmt, km, iperio) are closed-over Python ints → XLA constants;
# kmt comes from the LocalFluxFields → dynamic JAX array.
_meta = types.SimpleNamespace(
imt=imt, jmt=jmt, km=km, iperio=iperio,
kmt=fluxes.kmt,
)
x1, y1, z1, ib_new, jb_new, kb_new, ds, still_ok = jax.vmap(
lambda x, y, z, ia, ja, ka, dsm, dv: _subcycle_one_particle(
x, y, z, ia, ja, ka, intrpg, dsm, dv, _meta, fluxes, no_wflux
)
)(state.x, state.y, state.z, state.ia, state.ja, state.ka,
dsmin_arr, dxyz)
dt_particles = ds * dxyz
new_tt = state.tt + dt_particles
i0_new = jnp.clip(ib_new - 1, 0, imt - 1)
j0_new = jnp.clip(jb_new - 1, 0, jmt - 1)
on_land = fluxes.kmt[i0_new, j0_new] == 0
horiz_oob = (jb_new < 1) | (jb_new > jmt)
if fluxes.edge is not None: # domain edge inside the patch
horiz_oob = horiz_oob | fluxes.edge[i0_new, j0_new]
vert_oob = (kb_new < 1) | (kb_new > km)
new_reason = jnp.where(
zero_volume,
REASON_DXYZ,
jnp.where(
state.active & on_land,
REASON_LAND,
jnp.where(
state.active & ~on_land & horiz_oob,
REASON_BOUNDS_H,
jnp.where(state.active & ~on_land & vert_oob, REASON_BOUNDS_V, state.reason),
),
),
)
new_active = state.active & still_ok & ~zero_volume & ~horiz_oob
return ParticleState(
x=jnp.where(state.active, x1, state.x),
y=jnp.where(state.active, y1, state.y),
z=jnp.where(state.active, z1, state.z),
ia=jnp.where(state.active, ib_new, state.ia),
ja=jnp.where(state.active, jb_new, state.ja),
ka=jnp.where(state.active, kb_new, state.ka),
tt=jnp.where(state.active, new_tt, state.tt),
ts=jnp.where(state.active,
state.ts + dt_particles / state.subvol, state.ts),
t0=state.t0,
subvol=state.subvol,
niter=jnp.where(state.active, state.niter + 1, state.niter),
active=new_active,
reason=new_reason,
)
return step_all
# ===========================================================================
# Streaming (per-particle stencil) kernel
# ===========================================================================
def _subcycle_one_particle_stencil(
x0, y0, z0, ia, ja, ka, intrpg, dsmin_s, grid, stencil_p, no_wflux=False
):
"""Advance one particle by one sub-cycle using pre-extracted stencil values.
Like :func:`~tracmass.advection_scheme.full_grid._subcycle_one_particle`
but reads flux values directly from a per-particle
:class:`~tracmass.gcm.FluxStencil` instead of indexing into a global
or patch flux array. This eliminates all large-array indexing inside the
vmapped kernel.
Parameters
----------
x0, y0, z0 : float
Current fractional cell position (1-based dimensionless).
ia, ja, ka : int
Current cell indices (1-based).
intrpg : float
Time interpolation weight (0 = past slot, 1 = present slot).
dsmin_s : float
Step cap in seconds: time left until the particle's next sub-step
boundary.
grid : GridFields
Closed-over grid geometry. Used only for the ``kmt`` land check
on the new cell after crossing and for scalar dimensions.
stencil_p : FluxStencil
Per-particle stencil. After :func:`jax.vmap` peels the N dimension,
each field is a scalar (``kmt``, ``dxdy``) or a length-2 array
(flux face values and geometry).
no_wflux : bool
Compile-time flag; when ``True``, vertical fluxes are ignored.
Returns
-------
tuple
``(x1, y1, z1, ib, jb, kb, dt, still_active)`` where *dt* is the
physical time elapsed in seconds (``ds * dxyz``).
"""
imt, jmt, km, iperio = grid.imt, grid.jmt, grid.km, grid.iperio
intrpr = 1.0 - intrpg
# dxyz from per-particle geometry (no global array indexing)
dxyz = stencil_p.dxdy * (
intrpg * stencil_p.dzt[1] * stencil_p.zstot[1]
+ intrpr * stencil_p.dzt[0] * stencil_p.zstot[0]
)
dxyz = jnp.maximum(dxyz, EPS)
dsmin = dsmin_s / dxyz # normalised minimum time step
iam = jnp.where((ia - 1 == 0) & (iperio == 1), imt, ia - 1)
# Flux interpolation from stencil fields (shape (2,) after vmap)
uu_x = intrpg * stencil_p.u_east[1] + intrpr * stencil_p.u_east[0]
um_x_raw = intrpg * stencil_p.u_west[1] + intrpr * stencil_p.u_west[0]
um_x = jnp.where((ia - 1 == 0) & (iperio == 0), 0.0, um_x_raw)
uu_y = intrpg * stencil_p.v_north[1] + intrpr * stencil_p.v_north[0]
um_y = intrpg * stencil_p.v_south[1] + intrpr * stencil_p.v_south[0]
# Crossing times
dse, dsw = _cross_time_1d(uu_x, um_x, ia * 1.0, x0)
dsn, dss = _cross_time_1d(uu_y, um_y, ja * 1.0, y0)
if no_wflux:
dsu = UNDEF
dsd = UNDEF
else:
uu_z = intrpg * stencil_p.w_top[1] + intrpr * stencil_p.w_top[0]
um_z = intrpg * stencil_p.w_bot[1] + intrpr * stencil_p.w_bot[0]
dsu, dsd = _cross_time_1d(uu_z, um_z, ka * 1.0, z0)
ds = jnp.minimum(
jnp.minimum(jnp.minimum(dse, dsw), jnp.minimum(dsn, dss)),
jnp.minimum(jnp.minimum(dsu, dsd), dsmin),
)
# Which wall was hit?
hit_e = ds == dse; hit_w = ds == dsw
hit_n = ds == dsn; hit_s = ds == dss
if no_wflux:
hit_u = jnp.zeros_like(hit_e)
hit_d = jnp.zeros_like(hit_e)
else:
hit_u = ds == dsu; hit_d = ds == dsd
hit_none = ~(hit_e | hit_w | hit_n | hit_s | hit_u | hit_d)
# New positions
# EPS_WALL: nudge continuous positions off exact cell-wall values.
# When ds is tiny (float32 ULP of tt), _calc_pos_1d can return exactly
# the wall coordinate, causing a bounce to fire on every subsequent step
# and trapping particles in a corner oscillation cycle.
_EPS_WALL = jnp.float32(1e-5)
x1c = _calc_pos_1d(uu_x, um_x, ia * 1.0, x0, ds, iperio, imt * 1.0)
x1c = jnp.where(x1c == iam.astype(jnp.float32), x1c + _EPS_WALL, x1c)
x1c = jnp.where(x1c == ia.astype(jnp.float32), x1c - _EPS_WALL, x1c)
x1 = jnp.where(hit_e, ia * 1.0, jnp.where(hit_w, iam * 1.0, x1c))
y1c = _calc_pos_1d(uu_y, um_y, ja * 1.0, y0, ds, 0, 0.0)
y1c = jnp.where(y1c == (ja - 1).astype(jnp.float32), y1c + _EPS_WALL, y1c)
y1c = jnp.where(y1c == ja.astype(jnp.float32), y1c - _EPS_WALL, y1c)
y1 = jnp.where(hit_n, ja * 1.0, jnp.where(hit_s, (ja - 1) * 1.0, y1c))
if no_wflux:
z1 = z0
else:
z1c = _calc_pos_1d(uu_z, um_z, ka * 1.0, z0, ds, 0, 0.0)
z1c = jnp.where(z1c == (ka - 1).astype(jnp.float32), z1c + _EPS_WALL, z1c)
z1c = jnp.where(z1c == ka.astype(jnp.float32), z1c - _EPS_WALL, z1c)
z1 = jnp.where(hit_u, ka * 1.0, jnp.where(hit_d, (ka - 1) * 1.0, z1c))
z1 = jnp.where(z1 > km, km * 1.0, z1)
# New cell indices
ib, jb, kb = ia, ja, ka
ib = jnp.where(hit_e & (uu_x > 0.0), ia + 1, ib)
ib = jnp.where((ib > imt) & (iperio == 1), ib - imt, ib)
ib = jnp.where(hit_w & (um_x < 0.0), iam, ib)
jb = jnp.where(hit_n & (uu_y > 0.0), ja + 1, jb)
jb = jnp.where(hit_s & (um_y < 0.0), ja - 1, jb)
if not no_wflux:
kb = jnp.where(hit_u & (uu_z > 0.0), ka + 1, kb)
kb = jnp.where(kb > km, km, kb)
kb = jnp.where(hit_d & (um_z < 0.0), ka - 1, kb)
# Corners
ib = jnp.where(hit_n & (ds == dse), ia + 1, ib)
ib = jnp.where(hit_n & (ds == dsw), iam, ib)
jb = jnp.where(hit_e & (ds == dsn), ja + 1, jb)
jb = jnp.where(hit_e & (ds == dss), ja - 1, jb)
if not no_wflux:
kb = jnp.where(hit_e & (ds == dsu), ka + 1, kb)
kb = jnp.where(hit_e & (ds == dsd), ka - 1, kb)
all_undef = (
(dse >= UNDEF) & (dsw >= UNDEF)
& (dsn >= UNDEF) & (dss >= UNDEF)
& (dsu >= UNDEF) & (dsd >= UNDEF)
)
ib = jnp.where(hit_none & all_undef, ia, ib)
jb = jnp.where(hit_none & all_undef, ja, jb)
kb = jnp.where(hit_none & all_undef, ka, kb)
# Cyclic wrap
x1 = jnp.where((iperio != 0) & (x1 <= 0.0), x1 + imt, x1)
x1 = jnp.where((iperio != 0) & (x1 >= imt), x1 - imt, x1)
x1f = jnp.floor(x1); y1f = jnp.floor(y1); z1f = jnp.floor(z1)
ib = jnp.where(x1 != x1f, x1f.astype(jnp.int32) + 1, ib)
jb = jnp.where(y1 != y1f, y1f.astype(jnp.int32) + 1, jb)
kb = jnp.where(z1 != z1f, z1f.astype(jnp.int32) + 1, kb)
# Land/boundary check on the new cell (uses global kmt closed over in factory)
on_land = grid.kmt[
jnp.clip(ib - 1, 0, imt - 1),
jnp.clip(jb - 1, 0, jmt - 1),
] == 0
out_of_bounds = (jb < 1) | (jb > jmt) | (kb < 1) | (kb > km)
still_active = ~on_land & ~out_of_bounds
dt = ds * dxyz
return x1, y1, z1, ib, jb, kb, dt, still_active
[docs]
def piecewise_stencil(grid, no_wflux: bool = False):
"""Return a JIT-compiled step function for streaming (per-particle stencil) advection.
Unlike :func:`~tracmass.advection_scheme.full_grid.piecewise`, this kernel
receives pre-extracted :class:`~tracmass.gcm.FluxStencil` objects rather
than indexing into a global or patch flux array. This enables
memory-efficient integration of widely scattered particles: each Python call
to the returned function uses only the flux values at the particles' current
cell positions, which are loaded on-demand by
:func:`~tracmass.gcm.extract_stencil`.
*grid* is closed over so that ``kmt`` is available for land checking
on the new cell after each crossing. Scalar dimensions (``imt``, ``jmt``,
``km``, ``iperio``) become XLA constants in the compiled graph.
Parameters
----------
grid : GridFields
Static grid geometry.
no_wflux : bool, optional
When ``True``, vertical fluxes are ignored.
Returns
-------
callable
``step_all(stencil, state, intrpg, dsmin_s) -> ParticleState`` where
*dsmin_s* is the step cap in seconds — a scalar, or per particle
``(N,)`` (the integrators pass the time left until each particle's
next sub-step boundary).
"""
from tracmass.gcm.flux_grid import edge_mask
# Outer ring of cells on non-periodic axes: entering it leaves the domain.
edge = jnp.asarray(edge_mask(grid)) # (imt, jmt) bool
@jax.jit
def step_all(stencil, state: ParticleState, intrpg: float,
dsmin_s: float) -> ParticleState:
dsmin_arr = jnp.broadcast_to(dsmin_s, state.x.shape)
x1, y1, z1, ib_new, jb_new, kb_new, dt, still_ok = jax.vmap(
lambda x, y, z, ia, ja, ka, s, dsm: _subcycle_one_particle_stencil(
x, y, z, ia, ja, ka, intrpg, dsm, grid, s, no_wflux
)
)(state.x, state.y, state.z, state.ia, state.ja, state.ka, stencil,
dsmin_arr)
new_tt = state.tt + dt
i0_new = jnp.clip(ib_new - 1, 0, grid.imt - 1)
j0_new = jnp.clip(jb_new - 1, 0, grid.jmt - 1)
on_land = grid.kmt[i0_new, j0_new] == 0
horiz_oob = (jb_new < 1) | (jb_new > grid.jmt) | edge[i0_new, j0_new]
vert_oob = (kb_new < 1) | (kb_new > grid.km)
new_reason = jnp.where(
state.active & on_land,
REASON_LAND,
jnp.where(
state.active & ~on_land & horiz_oob,
REASON_BOUNDS_H,
jnp.where(state.active & ~on_land & vert_oob, REASON_BOUNDS_V, state.reason),
),
)
new_active = state.active & still_ok & ~horiz_oob
return ParticleState(
x=jnp.where(state.active, x1, state.x),
y=jnp.where(state.active, y1, state.y),
z=jnp.where(state.active, z1, state.z),
ia=jnp.where(state.active, ib_new, state.ia),
ja=jnp.where(state.active, jb_new, state.ja),
ka=jnp.where(state.active, kb_new, state.ka),
tt=jnp.where(state.active, new_tt, state.tt),
ts=state.ts, # updated by the caller (ts = tt / tseas)
t0=state.t0,
subvol=state.subvol,
niter=jnp.where(state.active, state.niter + 1, state.niter),
active=new_active,
reason=new_reason,
)
return step_all
