"""Full-grid advection kernels for TRACMASS-JAX.
Implements the :func:`piecewise` and :func:`analytical` step-function
factories that operate on the entire flux grid. Both close over a
:class:`~tracmass.gcm.GridFields` so that ``kmt`` and ``dxdy`` become XLA
constants in the compiled graph.
The returned step functions have the signature::
step_fn(fluxes, state, intrpg, dsmin) -> ParticleState
where *fluxes* is a computed :class:`~tracmass.gcm.FluxFields` and
*intrpg* is a scalar float; *dsmin* (the step cap in seconds) is a scalar
or a per-particle ``(N,)`` array — the integrators pass the time left until
each particle's next sub-step boundary.
"""
import jax
import jax.numpy as jnp
import numpy as np
from jax import lax
from tracmass.particle import (
ParticleState,
REASON_LAND,
REASON_BOUNDS_H,
REASON_BOUNDS_V,
REASON_DXYZ,
)
UNDEF = 1.0e20
EPS = 1.0e-10
# Module-level constants for the analytical scheme — built once at import time
# so they are not recreated on every JIT recompilation.
_DAWSON_H = 0.4
_DAWSON_NMAX = 6
_DAWSON_C = jnp.array(
[jnp.exp(-((2.0*(i + 1) - 1.0)*_DAWSON_H)**2) for i in range(_DAWSON_NMAX)],
dtype=jnp.float32,
)
_ASYM_H_IRANGE = jnp.arange(1.0, 11.0, dtype=jnp.float32)
# ===========================================================================
# Shared helpers
# ===========================================================================
def _flux_interp(field, i, j, k, intrpg):
"""Linearly interpolate a 4-D flux field between time levels."""
i0, j0, k0 = i - 1, j - 1, k - 1
intrpr = 1.0 - intrpg
return intrpg * field[i0, j0, k0, 1] + intrpr * field[i0, j0, k0, 0]
def _cross_time_1d(uu, um, ii, r0):
"""Return (sp, sn) — crossing times for the positive and negative walls."""
def _sp_divergent(uu, um, ii, r0):
ba = (r0 + (-ii + 1.0)) * (uu - um) + um
return jnp.where(ba > 0.0, (jnp.log(ba) - jnp.log(uu)) / (um - uu), UNDEF)
sp = jnp.where(
(uu > 0.0) & (r0 != ii),
jnp.where(um != uu, _sp_divergent(uu, um, ii, r0), (ii - r0) / uu),
UNDEF,
)
sp = jnp.where(sp <= 0.0, UNDEF, sp)
def _sn_divergent(uu, um, ii, r0):
ba = -((r0 - ii) * (uu - um) + uu)
return jnp.where(ba > 0.0, (jnp.log(ba) - jnp.log(-um)) / (um - uu), UNDEF)
sn = jnp.where(
(um < 0.0) & (r0 != (ii - 1.0)),
jnp.where(um != uu, _sn_divergent(uu, um, ii, r0), (ii - 1.0 - r0) / uu),
UNDEF,
)
sn = jnp.where(sn <= 0.0, UNDEF, sn)
return sp, sn
# ===========================================================================
# Time-analytical scheme (mod_pos_tanalytical.F90)
# ===========================================================================
def _dawson(x):
"""Dawson's integral F(x) = exp(-x²) ∫₀ˣ exp(t²) dt.
Port of daw() from mod_pos_tanalytical.F90.
Uses series expansion for |x| < 0.2, Rybicki sampling theorem otherwise.
JIT-compatible: no Python control flow on traced values.
"""
H = _DAWSON_H
A1, A2, A3 = 2.0/3.0, 0.4, 2.0/7.0
c = _DAWSON_C
x2 = x * x
series = x * (1.0 - A1*x2*(1.0 - A2*x2*(1.0 - A3*x2)))
xx = jnp.abs(x)
n0 = 2.0 * jnp.round(0.5 * xx / H) # even integer as float
xp = xx - n0 * H
e1_0 = jnp.exp(2.0 * xp * H)
e2 = e1_0 * e1_0
def _step(carry, ci):
e1, dd1, dd2, s = carry
s = s + ci * (e1/dd1 + 1.0/(dd2*e1))
return (e2*e1, dd1 + 2.0, dd2 - 2.0, s), None
(_, _, _, summ), _ = lax.scan(
_step,
(e1_0, n0 + 1.0, n0 - 1.0, jnp.zeros_like(x)),
c,
)
sampling = jnp.sign(x) / jnp.sqrt(jnp.pi) * jnp.exp(-xp**2) * summ
return jnp.where(jnp.abs(x) < 0.2, series, sampling)
def _erfexp_val(const, xi0, xi):
"""const * exp(xi²) * (erf(xi₀) − erf(xi)), numerically stable.
Port of erfexp() from mod_pos_tanalytical.F90.
Falls back to asymptotic expansion when |xi| > 25.
"""
standard = const * jnp.exp(xi**2) * (
jax.scipy.special.erf(xi0) - jax.scipy.special.erf(xi)
)
def _asym_h(z):
def _step(carry, i):
h0, h = carry
h0 = -0.5 * h0 * (2.0*i - 1.0) / (z*z)
return (h0, h + h0), None
(_, h), _ = lax.scan(_step,
(jnp.ones_like(z), jnp.ones_like(z)),
_ASYM_H_IRANGE)
return h
h0 = _asym_h(xi0)
h = _asym_h(xi)
asymp = const / jnp.sqrt(jnp.pi) * (
h/xi - jnp.exp(xi**2 - xi0**2) * h0/xi0
)
return jnp.where(jnp.abs(xi) > 25.0, asymp, standard)
def _r_apos(r0, ga, xi0, xic, c, daw_xi0):
"""Trajectory position r(xic) for the alpha > 0 (Dawson) regime."""
ef = jnp.exp(xi0**2 - xic**2)
return (r0 + ga)*ef - ga + c*(_dawson(xic) - ef*daw_xi0)
def _r_aneg(r0, ga, xi0, xic, c):
"""Trajectory position r(xic) for the alpha < 0 (erf) regime."""
return (r0 + ga)*jnp.exp(xic**2 - xi0**2) + _erfexp_val(c, xi0, xic) - ga
# ---------------------------------------------------------------------------
# Wall-crossing time for alpha > 0
# ---------------------------------------------------------------------------
def _apos_crossing(uu, um, vv, vm, ii, iim, r0, sm, s0, dstep, f0, f1):
"""Normalized crossing time (rw, sw) for the Dawson regime (alpha > 0).
Returns (rw, sw) where rw is the wall coordinate (-99 = no crossing)
and sw is the absolute normalized time of crossing.
"""
_EPS = 1.0e-10
denom = vv - vm - uu + um # = −alfa·dstep (<0 when alfa>0)
safe_d = jnp.where(jnp.abs(denom) < _EPS, _EPS, denom)
# Guard sqrt: only real when alfa > 0, but we compute safely
alfa_g = jnp.maximum(-(denom/dstep), _EPS) # guarded alfa
sqrt2a = jnp.sqrt(2.0 * alfa_g)
beta = um - uu - alfa_g * sm
ga = -iim + (f1*vm - f0*um) / safe_d
c = jnp.sqrt(2.0/alfa_g) * (f0*um*(vv - vm) + f1*vm*(um - uu)) / safe_d
xi0 = (beta + alfa_g*s0) / sqrt2a
daw_xi0 = _dawson(xi0)
# ssii / ssiim: normalized flux-zero times
den_ii = uu - vv + um*(f0 - 1.0) + vm*(1.0 - f1)
num_ii = uu + um*(f0 - 1.0)
ssii = jnp.where(jnp.abs(den_ii) < _EPS, num_ii/_EPS, num_ii/den_ii)
den_im = f0*um - f1*vm
ssiim = jnp.where(jnp.abs(den_im) < _EPS, f0*um/_EPS, f0*um/den_im)
xin_end = (vm - vv) / sqrt2a # xi at window end
xin_z_ii = (beta + alfa_g*(ssii*dstep + sm)) / sqrt2a # xi at flux-zero for ii
xin_z_im = (beta + alfa_g*(ssiim*dstep + sm)) / sqrt2a # xi at flux-zero for iim
def rval(xic):
return _r_apos(r0, ga, xi0, xic, c, daw_xi0)
# iiselect: determine initial wall target
no_z_ii = (ssii <= 0.0) | (ssii >= 1.0)
cfg_13 = (no_z_ii & ((uu > 0.0) | (vv > 0.0))) | (~no_z_ii & (vv > 0.0))
cfg_4ii = no_z_ii & ~((uu > 0.0) | (vv > 0.0))
cfg_2e = ~no_z_ii & ~(vv > 0.0) & (s0 <= sm + ssii*dstep)
cfg_2w = ~no_z_ii & ~(vv > 0.0) & (s0 > sm + ssii*dstep)
try_east = cfg_13 | cfg_2e
go_west_d = cfg_4ii | cfg_2w
xin_e = jnp.where(cfg_2e, xin_z_ii, xin_end)
dr_e = rval(xin_e) - ii
# Secondary check: drds <= 0 and flux-zero at iim side
drds = jnp.sqrt(alfa_g/2.0)*c - sqrt2a*xi0*(r0 + ga)
sec = (drds <= 0.0) & (ssiim > 0.0) & (ssiim < 1.0) & (vm > 0.0)
dra_im = rval(xin_z_im) - iim
east_ok = try_east & (dr_e >= 0.0) & ~(sec & (dra_im < 0.0))
go_west = go_west_d | (try_east & ~east_ok)
# iimselect
no_z_im = (ssiim <= 0.0) | (ssiim >= 1.0)
cfg_3b = no_z_im & ~((um > 0.0) | (vm > 0.0))
cfg_3c = ~no_z_im & (vm > 0.0) & (s0 < sm + ssiim*dstep)
cfg_3d = ~no_z_im & ~(vm > 0.0)
west_case = cfg_3b | cfg_3c | cfg_3d
xin_w = jnp.where(cfg_3c, xin_z_im, xin_end)
dr_w = rval(xin_w) - iim
west_ok = go_west & west_case & (dr_w <= 0.0)
do_east = east_ok
do_west = west_ok & ~east_ok
rw_sel = jnp.where(do_east, ii*1.0, jnp.where(do_west, iim*1.0, -99.0))
xin_s = jnp.where(do_east, xin_e, jnp.where(do_west, xin_w, xi0))
# Bisection in xi-space (100 fixed iterations)
def _bisect(_, carry):
xa, xb, xc = carry
dra = rval(xa) - rw_sel
drc = rval(xc) - rw_sel
new_xa = jnp.where(dra*drc > 0.0, xc, xa)
new_xb = jnp.where(dra*drc > 0.0, xb, xc)
return (new_xa, new_xb, 0.5*(new_xa + new_xb))
_, _, xc = lax.fori_loop(0, 100, _bisect, (xi0, xin_s, 0.5*(xi0 + xin_s)))
sw = jnp.where(do_east | do_west, (sqrt2a*xc - beta)/alfa_g, UNDEF)
return rw_sel, sw
# ---------------------------------------------------------------------------
# Wall-crossing time for alpha < 0
# ---------------------------------------------------------------------------
def _aneg_crossing(uu, um, vv, vm, ii, iim, r0, sm, s0, dstep, f0, f1):
"""Normalized crossing time (rw, sw) for the erf regime (alpha < 0)."""
_EPS = 1.0e-10
denom = vv - vm - uu + um # >0 when alfa<0
safe_d = jnp.where(jnp.abs(denom) < _EPS, _EPS, denom)
alfa_g = jnp.minimum(-(denom/dstep), -_EPS) # guarded alfa (<0)
sqrt_n2 = jnp.sqrt(-2.0 * alfa_g)
beta = um - uu - alfa_g * sm
ga = -iim + (f1*vm - f0*um) / safe_d
c = jnp.sqrt(jnp.pi / (-2.0*alfa_g)) * (
f0*um*(vv - vm) + f1*vm*(um - uu)) / safe_d
xi0 = (beta + alfa_g*s0) / sqrt_n2
den_ii = uu - vv + um*(f0 - 1.0) + vm*(1.0 - f1)
num_ii = uu + um*(f0 - 1.0)
ssii = jnp.where(jnp.abs(den_ii) < _EPS, num_ii/_EPS, num_ii/den_ii)
den_im = f0*um - f1*vm
ssiim = jnp.where(jnp.abs(den_im) < _EPS, f0*um/_EPS, f0*um/den_im)
xin_end = (vm - vv) / sqrt_n2
xin_z_ii = (beta + alfa_g*(ssii*dstep + sm)) / sqrt_n2
xin_z_im = (beta + alfa_g*(ssiim*dstep + sm)) / sqrt_n2
def rval(xic):
return _r_aneg(r0, ga, xi0, xic, c)
no_z_ii = (ssii <= 0.0) | (ssii >= 1.0)
cfg_13 = (no_z_ii & ((uu > 0.0) | (vv > 0.0))) | (~no_z_ii & (vv > 0.0))
cfg_4ii = no_z_ii & ~((uu > 0.0) | (vv > 0.0))
cfg_2e = ~no_z_ii & ~(vv > 0.0) & (s0 <= sm + ssii*dstep)
cfg_2w = ~no_z_ii & ~(vv > 0.0) & (s0 > sm + ssii*dstep)
try_east = cfg_13 | cfg_2e
go_west_d = cfg_4ii | cfg_2w
xin_e = jnp.where(cfg_2e, xin_z_ii, xin_end)
dr_e = rval(xin_e) - ii
drds = c * jnp.sqrt(-2.0*alfa_g/jnp.pi) - sqrt_n2*xi0*(r0 + ga)
sec = (drds <= 0.0) & (ssiim > 0.0) & (ssiim < 1.0) & (vm > 0.0)
dra_im = rval(xin_z_im) - iim
east_ok = try_east & (dr_e >= 0.0) & ~(sec & (dra_im < 0.0))
go_west = go_west_d | (try_east & ~east_ok)
no_z_im = (ssiim <= 0.0) | (ssiim >= 1.0)
cfg_3b = no_z_im & ~((um > 0.0) | (vm > 0.0))
cfg_3c = ~no_z_im & (vm > 0.0) & (s0 < sm + ssiim*dstep)
cfg_3d = ~no_z_im & ~(vm > 0.0)
west_case = cfg_3b | cfg_3c | cfg_3d
xin_w = jnp.where(cfg_3c, xin_z_im, xin_end)
dr_w = rval(xin_w) - iim
west_ok = go_west & west_case & (dr_w <= 0.0)
do_east = east_ok
do_west = west_ok & ~east_ok
rw_sel = jnp.where(do_east, ii*1.0, jnp.where(do_west, iim*1.0, -99.0))
xin_s = jnp.where(do_east, xin_e, jnp.where(do_west, xin_w, xi0))
def _bisect(_, carry):
xa, xb, xc = carry
dra = rval(xa) - rw_sel
drc = rval(xc) - rw_sel
new_xa = jnp.where(dra*drc > 0.0, xc, xa)
new_xb = jnp.where(dra*drc > 0.0, xb, xc)
return (new_xa, new_xb, 0.5*(new_xa + new_xb))
_, _, xc = lax.fori_loop(0, 100, _bisect, (xi0, xin_s, 0.5*(xi0 + xin_s)))
sw = jnp.where(do_east | do_west, (sqrt_n2*xc - beta)/alfa_g, UNDEF)
return rw_sel, sw
# ---------------------------------------------------------------------------
# Wall-crossing time for alpha = 0
# ---------------------------------------------------------------------------
def _azer_crossing(uu, um, vv, vm, ii, iim, r0, sm, s0, dstep, f0, f1):
"""Normalized crossing time (rw, sw) for the alpha = 0 regime."""
_EPS = 1.0e-10
beta = um - uu
gamma = (f0*um - f1*vm) / dstep
# ------------------------------------------------------------------
# Case A: beta = 0, gamma = 0 (uniform constant flow)
# ------------------------------------------------------------------
delta = -f0 * um
rw_A = jnp.where(delta < 0.0, ii*1.0,
jnp.where(delta > 0.0, iim*1.0, -99.0))
sw_A = jnp.where(jnp.abs(delta) > _EPS,
s0 - (rw_A - r0)/delta, UNDEF)
# ------------------------------------------------------------------
# Case B: beta = 0, gamma ≠ 0 (uniform time-varying, quadratic)
# r(s) = r0 − ds·(0.5·gamma·(s + s0 − 2·sm) − f0·um)
# ------------------------------------------------------------------
g_safe = jnp.where(jnp.abs(gamma) > _EPS, gamma, _EPS)
del_ga = -(f0*um) / g_safe - sm # position of parabola vertex − sm
s0d = s0 + del_ga # s0 + del_ga
pos_reg = s0d >= 0.0 # which branch of Fortran logic
# Wall choice when s0+del_ga >= 0
rw_B1 = jnp.where(gamma > 0.0, iim*1.0, ii*1.0)
disc1 = s0d**2 - 2.0*(rw_B1 - r0) / g_safe
sw_B1 = -del_ga + jnp.sqrt(jnp.maximum(disc1, 0.0)) # pos region
# Wall choice when s0+del_ga < 0, first try
rw_B2a = jnp.where(gamma > 0.0, ii*1.0, iim*1.0)
disc2a = s0d**2 - 2.0*(rw_B2a - r0) / g_safe
sw_B2a = -del_ga - jnp.sqrt(jnp.maximum(disc2a, 0.0)) # neg region, first try
# Fallback if disc2a < 0: swap wall
rw_B2b = jnp.where(gamma > 0.0, iim*1.0, ii*1.0)
disc2b = s0d**2 - 2.0*(rw_B2b - r0) / g_safe
sw_B2b = -del_ga + jnp.sqrt(jnp.maximum(disc2b, 0.0))
has_disc2a = disc2a >= 0.0
rw_B = jnp.where(pos_reg, rw_B1,
jnp.where(has_disc2a, rw_B2a, rw_B2b))
sw_B = jnp.where(pos_reg, sw_B1,
jnp.where(has_disc2a, sw_B2a, sw_B2b))
disc_B = jnp.where(pos_reg, disc1,
jnp.where(has_disc2a, disc2a, disc2b))
sw_B = jnp.where(disc_B >= 0.0, sw_B, UNDEF)
# ------------------------------------------------------------------
# Case C: beta ≠ 0, gamma = 0 (exponential, space-varying)
# ------------------------------------------------------------------
b_safe = jnp.where(jnp.abs(beta) > _EPS, beta, _EPS)
denom_e = uu - (1.0 - f0)*um
de_safe = jnp.where(jnp.abs(denom_e) > _EPS, denom_e, _EPS)
xi_e = 1.0 + (ii - r0)*beta / de_safe
sw_eC = jnp.where((uu > 0.0) & (xi_e > 0.0),
s0 + jnp.log(jnp.maximum(xi_e, _EPS))/b_safe,
UNDEF)
fu_safe = jnp.where(jnp.abs(f0*um) > _EPS, f0*um, _EPS)
xi_w = 1.0 - (r0 - iim)*beta / fu_safe
sw_wC = jnp.where((um < 0.0) & (xi_w > 0.0),
s0 + jnp.log(jnp.maximum(xi_w, _EPS))/b_safe,
UNDEF)
rw_C = jnp.where(sw_eC <= sw_wC,
jnp.where(sw_eC < UNDEF/2, ii*1.0, -99.0),
jnp.where(sw_wC < UNDEF/2, iim*1.0, -99.0))
sw_C = jnp.minimum(sw_eC, sw_wC)
# ------------------------------------------------------------------
# Case D: beta ≠ 0, gamma ≠ 0 (bisection in s-space)
# ------------------------------------------------------------------
ssii_D = (f0*uu - beta) / g_safe
ssiim_D = f0*um / g_safe
exp_term = (r0 - iim) + (gamma*(s0 - sm - 1.0/b_safe) - f0*um)/b_safe
drds_D = f0*um - gamma*(s0 - sm) - beta*(r0 - iim)
def r_s(s):
ds = s - s0
return r0 + exp_term*(jnp.exp(-beta*ds) - 1.0) - gamma*ds/b_safe
no_z_iiD = (ssii_D <= 0.0) | (ssii_D >= dstep)
cfg1D = (no_z_iiD & ((uu > 0.0) | (vv > 0.0))) | \
(~no_z_iiD & (vv > 0.0))
cfg4D = no_z_iiD & ~((uu > 0.0) | (vv > 0.0))
cfg2eD = ~no_z_iiD & ~(vv > 0.0) & (s0 <= sm + ssii_D)
cfg2wD = ~no_z_iiD & ~(vv > 0.0) & (s0 > sm + ssii_D)
try_e_D = cfg1D | cfg2eD
gw_d_D = cfg4D | cfg2wD
si00_end = sm + dstep - s0
si00_z_ii = sm + ssii_D - s0
si00_z_im = sm + ssiim_D - s0
si00_e_D = jnp.where(cfg2eD, si00_z_ii, si00_end)
dr_e_D = r_s(s0 + si00_e_D) - ii
sec_D = (drds_D <= 0.0) & (ssiim_D > 0.0) & \
(ssiim_D < dstep) & (vm > 0.0)
dra_D = r_s(s0 + si00_z_im) - iim
east_D = try_e_D & (dr_e_D >= 0.0) & ~(sec_D & (dra_D < 0.0))
gw_D = gw_d_D | (try_e_D & ~east_D)
no_z_imD = (ssiim_D <= 0.0) | (ssiim_D >= dstep)
cfg3bD = no_z_imD & ~((um > 0.0) | (vm > 0.0))
cfg3cD = ~no_z_imD & (vm > 0.0) & (s0 < sm + ssiim_D)
cfg3dD = ~no_z_imD & ~(vm > 0.0)
wcase_D = cfg3bD | cfg3cD | cfg3dD
si00_w_D = jnp.where(cfg3cD, si00_z_im, si00_end)
dr_w_D = r_s(s0 + si00_w_D) - iim
west_D = gw_D & wcase_D & (dr_w_D <= 0.0)
de_D = east_D
dw_D = west_D & ~east_D
rw_D = jnp.where(de_D, ii*1.0, jnp.where(dw_D, iim*1.0, -99.0))
si00_D = jnp.where(de_D, si00_e_D, jnp.where(dw_D, si00_w_D, 0.0))
def _bisect_s(_, carry):
xa, xb, xc = carry
dra = r_s(xa) - rw_D
drc = r_s(xc) - rw_D
new_xa = jnp.where(dra*drc > 0.0, xc, xa)
new_xb = jnp.where(dra*drc > 0.0, xb, xc)
return (new_xa, new_xb, 0.5*(new_xa + new_xb))
_, _, xc_D = lax.fori_loop(
0, 100, _bisect_s, (s0, s0 + si00_D, s0 + 0.5*si00_D)
)
sw_D = jnp.where(de_D | dw_D, xc_D, UNDEF)
# ------------------------------------------------------------------
# Select among A / B / C / D
# ------------------------------------------------------------------
bz = jnp.abs(beta) < _EPS
gz = jnp.abs(gamma) < _EPS
rw = jnp.where(bz & gz, rw_A,
jnp.where(bz & ~gz, rw_B,
jnp.where(~bz & gz, rw_C, rw_D)))
sw = jnp.where(bz & gz, sw_A,
jnp.where(bz & ~gz, sw_B,
jnp.where(~bz & gz, sw_C, sw_D)))
return rw, sw
# ---------------------------------------------------------------------------
# 1-D analytical crossing-time dispatcher
# ---------------------------------------------------------------------------
def _cross_time_analytical_1d(uu, um, vv, vm, ii, iim, r0, sm, s0, dstep, f0, f1):
"""Return (sp, sn) crossing times for the time-analytical scheme.
sp = time to cross positive wall (ii), sn = time to cross negative wall (iim).
Both in normalized (dxyz-divided) units. UNDEF means no crossing.
"""
alfa_sign = -(vv - vm - uu + um)
rw_p, sw_p = _apos_crossing(uu, um, vv, vm, ii, iim, r0, sm, s0, dstep, f0, f1)
rw_n, sw_n = _aneg_crossing(uu, um, vv, vm, ii, iim, r0, sm, s0, dstep, f0, f1)
rw_z, sw_z = _azer_crossing(uu, um, vv, vm, ii, iim, r0, sm, s0, dstep, f0, f1)
rw = jnp.where(alfa_sign > 0.0, rw_p,
jnp.where(alfa_sign < 0.0, rw_n, rw_z))
sw = jnp.where(alfa_sign > 0.0, sw_p,
jnp.where(alfa_sign < 0.0, sw_n, sw_z))
sp = jnp.where(rw == ii*1.0, sw - s0, UNDEF)
sn = jnp.where(rw == iim*1.0, sw - s0, UNDEF)
sp = jnp.where(sp <= 0.0, UNDEF, sp)
sn = jnp.where(sn <= 0.0, UNDEF, sn)
return sp, sn
# ---------------------------------------------------------------------------
# 1-D analytical position update
# ---------------------------------------------------------------------------
def _calc_pos_analytical_1d(uu, um, vv, vm, ii, iim, r0, ds, sm, s0, dstep, f0, f1):
"""New position after normalized time ds (analytical scheme).
Fluxes uu/um/vv/vm should already be adjusted by f0/f1 for vertical dir.
"""
_EPS = 1.0e-10
alfa_sign = -(vv - vm - uu + um)
denom = vv - vm - uu + um
safe_d = jnp.where(jnp.abs(denom) < _EPS, _EPS, denom)
# --- alpha > 0 ---
alfa_p = jnp.maximum(-denom/dstep, _EPS)
sqrt2a_p = jnp.sqrt(2.0 * alfa_p)
beta_p = um - uu - alfa_p*sm
xi0_p = (beta_p + alfa_p*s0) / sqrt2a_p
xis_p = (beta_p + alfa_p*(s0 + ds)) / sqrt2a_p
ga_p = -iim + (f1*vm - f0*um) / safe_d
daw_t = jnp.sqrt(2.0/alfa_p) * (f0*um*(vv-vm) + f1*vm*(um-uu)) / safe_d
r1_pos = (r0 + ga_p)*jnp.exp(xi0_p**2 - xis_p**2) - ga_p + daw_t*(
_dawson(xis_p) - jnp.exp(xi0_p**2 - xis_p**2)*_dawson(xi0_p))
# --- alpha < 0 ---
alfa_n = jnp.minimum(-denom/dstep, -_EPS)
sqrt_n2 = jnp.sqrt(-2.0 * alfa_n)
beta_n = um - uu - alfa_n*sm
xi0_n = (beta_n + alfa_n*s0) / sqrt_n2
xis_n = (beta_n + alfa_n*(s0 + ds)) / sqrt_n2
ga_n = -iim + (f1*vm - f0*um) / safe_d
erf_t = jnp.sqrt(jnp.pi/(-2.0*alfa_n)) * (
f0*um*(vv-vm) + f1*vm*(um-uu)) / safe_d
r1_neg = (r0 + ga_n)*jnp.exp(xis_n**2 - xi0_n**2) + \
_erfexp_val(erf_t, xi0_n, xis_n) - ga_n
# --- alpha = 0, beta = 0 ---
gamma_z = (f0*um - f1*vm) / dstep
r1_zbz = r0 - ds*(0.5*gamma_z*(2.0*s0 + ds - 2.0*sm) - f0*um)
# --- alpha = 0, beta ≠ 0 ---
beta_z = um - uu
b_safe = jnp.where(jnp.abs(beta_z) > _EPS, beta_z, _EPS)
exp_term = (r0 - iim) + (gamma_z*(s0 - sm - 1.0/b_safe) - f0*um)/b_safe
r1_zbg = r0 + exp_term*(jnp.exp(-beta_z*ds) - 1.0) - gamma_z*ds/b_safe
bz = jnp.abs(beta_z) < _EPS
r1_z = jnp.where(bz, r1_zbz, r1_zbg)
r1 = jnp.where(alfa_sign > 0.0, r1_pos,
jnp.where(alfa_sign < 0.0, r1_neg, r1_z))
# Clamp small overshoots (< 1e-7) back inside
r1 = jnp.where((r1 - ii > 0.0) & (r1 - ii < 1.0e-7), ii - 1.0e-7, r1)
r1 = jnp.where((iim - r1 > 0.0) & (iim - r1 < 1.0e-7), iim + 1.0e-7, r1)
return r1
# ---------------------------------------------------------------------------
# Analytical-scheme particle kernel
# ---------------------------------------------------------------------------
def _subcycle_one_particle_analytical(
x0, y0, z0, ia, ja, ka, tt, ts, dxyz, tseas, grid, fluxes, no_wflux=False
):
"""Advance one particle by one sub-cycle using the time-analytical scheme.
Reads fluxes from both time slots of *fluxes* directly (no intrpg
interpolation). Static grid data (kmt, dimensions, periodicity) is
read from *grid*.
Returns (x1, y1, z1, ib, jb, kb, ds, still_active).
When *no_wflux* is True (a Python bool, resolved at JIT compile time) the
vertical flux is ignored: particles never cross depth-layer walls and their
z-coordinate and ka index are held fixed.
"""
imt, jmt, km, iperio = grid.imt, grid.jmt, grid.km, grid.iperio
dstep = tseas / dxyz # window duration (normalized)
sm = jnp.floor(ts) * dstep # lower window boundary
s0 = tt / dxyz # current normalized time
dsmin = dstep - (s0 - sm) # remaining window time
iam = jnp.where((ia - 1 == 0) & (iperio == 1), imt, ia - 1)
# --- Zonal (x) ---
uu_x = fluxes.uflux[ia-1, ja-1, ka-1, 0]
um_x_raw = fluxes.uflux[iam-1, ja-1, ka-1, 0]
um_x = jnp.where((ia - 1 == 0) & (iperio == 0), 0.0, um_x_raw)
vv_x = fluxes.uflux[ia-1, ja-1, ka-1, 1]
vm_x_raw = fluxes.uflux[iam-1, ja-1, ka-1, 1]
vm_x = jnp.where((ia - 1 == 0) & (iperio == 0), 0.0, vm_x_raw)
# Enforce no-flow BC: zero flux toward any land neighbor (catches residuals
# not removed by the fill-value threshold in flux_grid.compute()).
_e_land = grid.kmt[jnp.clip(ia, 0, imt-1), jnp.clip(ja-1, 0, jmt-1)] == 0
_w_land = grid.kmt[jnp.clip(iam-1, 0, imt-1), jnp.clip(ja-1, 0, jmt-1)] == 0
uu_x = jnp.where(_e_land, 0.0, uu_x)
um_x = jnp.where(_w_land, 0.0, um_x)
vv_x = jnp.where(_e_land, 0.0, vv_x)
vm_x = jnp.where(_w_land, 0.0, vm_x)
dse, dsw = _cross_time_analytical_1d(
uu_x, um_x, vv_x, vm_x,
ia*1.0, iam*1.0, x0, sm, s0, dstep, 1.0, 1.0,
)
# --- Meridional (y) ---
uu_y = fluxes.vflux[ia-1, ja-1, ka-1, 0]
um_y = fluxes.vflux[ia-1, ja-2, ka-1, 0]
vv_y = fluxes.vflux[ia-1, ja-1, ka-1, 1]
vm_y = fluxes.vflux[ia-1, ja-2, ka-1, 1]
_n_land = grid.kmt[jnp.clip(ia-1, 0, imt-1), jnp.clip(ja, 0, jmt-1)] == 0
_s_land = grid.kmt[jnp.clip(ia-1, 0, imt-1), jnp.clip(ja-2, 0, jmt-1)] == 0
uu_y = jnp.where(_n_land, 0.0, uu_y)
um_y = jnp.where(_s_land, 0.0, um_y)
vv_y = jnp.where(_n_land, 0.0, vv_y)
vm_y = jnp.where(_s_land, 0.0, vm_y)
dsn, dss = _cross_time_analytical_1d(
uu_y, um_y, vv_y, vm_y,
ja*1.0, (ja-1)*1.0, y0, sm, s0, dstep, 1.0, 1.0,
)
# --- Vertical (z), with thickness correction ---
if no_wflux:
dsu = UNDEF
dsd = UNDEF
else:
uu_z_r = fluxes.wflux[ia-1, ja-1, ka-1, 0]
km_idx = jnp.clip(ka - 2, 0, km) # 0-based lower-wall index
um_z_r = fluxes.wflux[ia-1, ja-1, km_idx, 0]
vv_z_r = fluxes.wflux[ia-1, ja-1, ka-1, 1]
vm_z_r = fluxes.wflux[ia-1, ja-1, km_idx, 1]
dzt1 = fluxes.dzt[ia-1, ja-1, ka-1, 0]
dzt2 = fluxes.dzt[ia-1, ja-1, ka-1, 1]
intrpg_z = jnp.mod(ts, 1.0)
dzs = (1.0 - intrpg_z)*dzt1 + intrpg_z*dzt2
f0_z = dzt1 / jnp.maximum(dzs, EPS)
f1_z = dzt2 / jnp.maximum(dzs, EPS)
uu_z = uu_z_r / f0_z; um_z = um_z_r / f0_z
vv_z = vv_z_r / f1_z; vm_z = vm_z_r / f1_z
# Sea surface (ka==1): no flux above the surface layer. km_idx clips
# ka-2=-1 to 0, making um_z == uu_z and giving a spurious finite dsd
# that pushes particles out through kb=0.
um_z = jnp.where(ka == 1, 0.0, um_z)
vm_z = jnp.where(ka == 1, 0.0, vm_z)
dsu, dsd = _cross_time_analytical_1d(
uu_z, um_z, vv_z, vm_z,
ka*1.0, (ka-1)*1.0, z0, sm, s0, dstep, f0_z, f1_z,
)
# Minimum crossing time
ds = jnp.minimum(
jnp.minimum(jnp.minimum(dse, dsw), jnp.minimum(dsn, dss)),
jnp.minimum(jnp.minimum(dsu, dsd), dsmin),
)
hit_e = (ds == dse); hit_w = (ds == dsw)
hit_n = (ds == dsn); hit_s = (ds == dss)
# When no_wflux, vertical walls are never the limiting step
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)
# Position updates
x1c = _calc_pos_analytical_1d(
uu_x, um_x, vv_x, vm_x, ia*1.0, iam*1.0, x0, ds, sm, s0, dstep, 1.0, 1.0)
x1 = jnp.where(hit_e, ia*1.0, jnp.where(hit_w, iam*1.0, x1c))
y1c = _calc_pos_analytical_1d(
uu_y, um_y, vv_y, vm_y, ja*1.0, (ja-1)*1.0, y0, ds, sm, s0, dstep, 1.0, 1.0)
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_analytical_1d(
uu_z, um_z, vv_z, vm_z, ka*1.0, (ka-1)*1.0, z0, ds, sm, s0, dstep, f0_z, f1_z)
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)
# Box 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 — only step to the diagonal cell if it is ocean; otherwise the
# cardinal crossing already recorded in ib/jb is kept.
_ne_ocean = grid.kmt[jnp.clip(ia, 0, imt-1), jnp.clip(ja, 0, jmt-1)] > 0
_nw_ocean = grid.kmt[jnp.clip(iam-1, 0, imt-1), jnp.clip(ja, 0, jmt-1)] > 0
_se_ocean = grid.kmt[jnp.clip(ia, 0, imt-1), jnp.clip(ja-2, 0, jmt-1)] > 0
ib = jnp.where(hit_n & (ds == dse) & _ne_ocean, ia + 1, ib)
ib = jnp.where(hit_n & (ds == dsw) & _nw_ocean, iam, ib)
jb = jnp.where(hit_e & (ds == dsn) & _ne_ocean, ja + 1, jb)
jb = jnp.where(hit_e & (ds == dss) & _se_ocean, 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 — strict x1 > imt matches Fortran mod_pos_tstep.F90.
# x1 == imt is left unwrapped intentionally: both a westward crossing into
# cell imt (ib=imt, x1=imt) and an eastward crossing out of cell imt
# (ib=1, x1=imt) leave x1=imt; _calc_pos_analytical_1d handles r0=imt
# with ii=1 via the r0_adj special case on the next sub-step.
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)
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
return x1, y1, z1, ib, jb, kb, ds, still_active
# ===========================================================================
# Vectorised step — analytical scheme
# ===========================================================================
[docs]
def analytical(grid, tseas: float, no_wflux: bool = False):
"""Return a JIT-compiled step function that advances all particles analytically.
*grid* (a ``GridFields``) is closed over as XLA constants.
The returned function takes a ``FluxFields`` as its first dynamic argument,
enabling lazy loading: flux data need not be computed until the step
function is actually called.
Parameters
----------
grid : GridFields
Static grid geometry. Closed over by the compiled kernel.
tseas : float
GCM time-window duration in seconds.
no_wflux : bool
When True, vertical (w) fluxes are ignored. Automatically True
when ``fluxes.wflux is None`` at call time, but must be set here
to select the compiled variant.
"""
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(fluxes, state: ParticleState, intrpg: float, dsmin: float) -> ParticleState:
ib, jb, kb = state.ia, state.ja, state.ka
i0 = jnp.clip(ib - 1, 0, grid.imt - 1)
j0 = jnp.clip(jb - 1, 0, grid.jmt - 1)
k0 = jnp.clip(kb - 1, 0, grid.km - 1)
dxyz_raw = grid.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)
x1, y1, z1, ib_new, jb_new, kb_new, ds, still_ok = jax.vmap(
lambda x, y, z, ia, ja, ka, dv, tt, ts:
_subcycle_one_particle_analytical(
x, y, z, ia, ja, ka, tt, ts, dv, tseas, grid, fluxes, no_wflux
)
)(state.x, state.y, state.z,
state.ia, state.ja, state.ka,
dxyz, state.tt, state.ts)
dt_particles = ds * dxyz
new_tt = state.tt + dt_particles
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(
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
# ===========================================================================
# Piecewise-linear scheme
# ===========================================================================
def _calc_pos_1d(uu, um, ii, r0, ds, iperio, imt):
"""New position after advection for time ds."""
ratio = jnp.where(jnp.abs(uu - um) > 1e-30, um / (uu - um), 1.0e15)
use_exp = jnp.abs(ratio) < 1.0e14
r0_adj = jnp.where(
(iperio == 1) & (r0 == imt) & (ii == 1), 0.0, r0,
)
r1_exp = (r0_adj + (-ii + 1.0 + ratio)) * jnp.exp((uu - um) * ds) + (ii - 1.0) - ratio
r1_lin = r0_adj + uu * ds
return jnp.where(use_exp, r1_exp, r1_lin)
def _subcycle_one_particle(
x0, y0, z0, ia, ja, ka, intrpg, dsmin, dxyz, grid, fluxes, no_wflux=False
):
"""Advance one particle by one sub-cycle. Returns (x1,y1,z1,ib,jb,kb,ds,active).
Static grid data (kmt, dimensions, periodicity) is read from *grid*.
Time-varying flux data is read from *fluxes*.
When *no_wflux* is True (a Python bool, resolved at JIT compile time) the
vertical flux is ignored: particles never cross depth-layer walls and their
z-coordinate and ka index are held fixed.
"""
imt, jmt, km, iperio = grid.imt, grid.jmt, grid.km, grid.iperio
intrpr = 1.0 - intrpg
iam = jnp.where((ia - 1 == 0) & (iperio == 1), imt, ia - 1)
# -----------------------------------------------------------------------
# update_bounce (Fortran mod_pos_tstep.F90): when a particle sits exactly
# on a cell wall and the time-interpolated flux pushes it further into the
# adjacent cell, advance the cell index before computing crossing times.
# This prevents indefinite oscillation (crossing time = 0) at walls.
#
# Phase 1 — read fluxes at the PRE-BOUNCE indices to determine direction.
# Phase 2 — after updating ia/iam/ja/ka, re-read fluxes at the
# POST-BOUNCE indices for the crossing-time calculation.
# The re-read is essential: using pre-bounce fluxes with post-bounce
# indices gives incorrect (often huge) crossing times that cause step-size
# errors and tt-overshoot past timax_seconds.
# -----------------------------------------------------------------------
# --- Phase 1: bounce direction check ---
# Precompute land-neighbor flags (kmt lookups on the static grid).
# These prevent the bounce from firing toward a land cell even when the
# stored velocity has a small residual value that is not a fill-value.
_e_land = grid.kmt[jnp.clip(ia, 0, imt-1), jnp.clip(ja-1, 0, jmt-1)] == 0
_w_land = grid.kmt[jnp.clip(iam-1, 0, imt-1), jnp.clip(ja-1, 0, jmt-1)] == 0
_n_land = grid.kmt[jnp.clip(ia-1, 0, imt-1), jnp.clip(ja, 0, jmt-1)] == 0
_s_land = grid.kmt[jnp.clip(ia-1, 0, imt-1), jnp.clip(ja-2, 0, jmt-1)] == 0
# X: east/west
_bx_uu = _flux_interp(fluxes.uflux, ia, ja, ka, intrpg)
_bx_um_raw = _flux_interp(fluxes.uflux, iam, ja, ka, intrpg)
_bx_um = jnp.where((ia - 1 == 0) & (iperio == 0), 0.0, _bx_um_raw)
_bx_e = (x0 == ia.astype(jnp.float32)) & (_bx_uu > 0.0) & ((ia < imt) | (iperio == 1)) & ~_e_land
_bx_w = (x0 == iam.astype(jnp.float32)) & (_bx_um < 0.0) & ~_w_land
# Y: north/south
_by_uu = _flux_interp(fluxes.vflux, ia, ja, ka, intrpg)
_km_y = jnp.maximum(ja - 2, 0) # safe 0-based south-wall index
_by_um = intrpg * fluxes.vflux[ia-1, _km_y, ka-1, 1] + intrpr * fluxes.vflux[ia-1, _km_y, ka-1, 0]
_by_um = jnp.where(ja == 1, 0.0, _by_um) # no south flux at southern boundary
_by_n = (y0 == ja.astype(jnp.float32)) & (_by_uu > 0.0) & (ja < jmt) & ~_n_land
_by_s = (y0 == (ja - 1).astype(jnp.float32)) & (_by_um < 0.0) & (ja > 1) & ~_s_land
if not no_wflux:
# Z: up/down
_bz_uu = _flux_interp(fluxes.wflux, ia, ja, ka, intrpg)
_km_idx = jnp.maximum(ka - 2, 0)
_bz_um = intrpg * fluxes.wflux[ia-1, ja-1, _km_idx, 1] + intrpr * fluxes.wflux[ia-1, ja-1, _km_idx, 0]
_bz_u = (z0 == ka.astype(jnp.float32)) & (_bz_uu > 0.0) & (ka < km)
_bz_d = (z0 == (ka - 1).astype(jnp.float32)) & (_bz_um < 0.0) & (ka > 1)
# --- Apply bounces: update ia/iam, ja, ka ---
_old_ia, _old_iam = ia, iam
ia = jnp.where(_bx_e, _old_ia + 1, jnp.where(_bx_w, _old_iam, _old_ia))
iam = jnp.where(_bx_e, _old_ia, jnp.where(_bx_w, _old_iam - 1, _old_iam))
ia = jnp.where((ia > imt) & (iperio == 1), ia - imt, ia)
iam = jnp.where((iam == 0) & (iperio == 1), imt, iam)
ja = jnp.where(_by_n, ja + 1, jnp.where(_by_s, ja - 1, ja))
if not no_wflux:
ka = jnp.where(_bz_u, ka + 1, jnp.where(_bz_d, ka - 1, ka))
ka = jnp.clip(ka, 1, km)
# --- Phase 2: re-read fluxes at post-bounce indices for crossing times ---
# X crossing times
uu_x = _flux_interp(fluxes.uflux, ia, ja, ka, intrpg)
um_x_raw = _flux_interp(fluxes.uflux, iam, ja, ka, intrpg)
um_x = jnp.where((ia - 1 == 0) & (iperio == 0), 0.0, um_x_raw)
# Enforce no-flow BC: zero flux toward any land neighbor (catches residuals
# not removed by the fill-value threshold in flux_grid.compute()).
_e_land2 = grid.kmt[jnp.clip(ia, 0, imt-1), jnp.clip(ja-1, 0, jmt-1)] == 0
_w_land2 = grid.kmt[jnp.clip(iam-1, 0, imt-1), jnp.clip(ja-1, 0, jmt-1)] == 0
uu_x = jnp.where(_e_land2, 0.0, uu_x)
um_x = jnp.where(_w_land2, 0.0, um_x)
dse, dsw = _cross_time_1d(uu_x, um_x, ia * 1.0, x0)
# Y crossing times
uu_y = _flux_interp(fluxes.vflux, ia, ja, ka, intrpg)
_km_y2 = jnp.maximum(ja - 2, 0)
um_y = intrpg * fluxes.vflux[ia-1, _km_y2, ka-1, 1] + intrpr * fluxes.vflux[ia-1, _km_y2, ka-1, 0]
um_y = jnp.where(ja == 1, 0.0, um_y) # no south flux at southern boundary
_n_land2 = grid.kmt[jnp.clip(ia-1, 0, imt-1), jnp.clip(ja, 0, jmt-1)] == 0
_s_land2 = grid.kmt[jnp.clip(ia-1, 0, imt-1), jnp.clip(ja-2, 0, jmt-1)] == 0
uu_y = jnp.where(_n_land2, 0.0, uu_y)
um_y = jnp.where(_s_land2, 0.0, um_y)
dsn, dss = _cross_time_1d(uu_y, um_y, ja * 1.0, y0)
# Z crossing times
if no_wflux:
dsu = UNDEF
dsd = UNDEF
else:
uu_z = _flux_interp(fluxes.wflux, ia, ja, ka, intrpg)
km_idx_pw = jnp.maximum(ka - 2, 0)
um_z = intrpg * fluxes.wflux[ia-1, ja-1, km_idx_pw, 1] + intrpr * fluxes.wflux[ia-1, ja-1, km_idx_pw, 0]
# Sea surface (ka==1): km_idx_pw clips to 0, making um_z == uu_z and
# giving a spurious finite dsd that exits particles through kb=0.
um_z = jnp.where(ka == 1, 0.0, um_z)
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?
hit_e = (ds == dse); hit_w = (ds == dsw)
hit_n = (ds == dsn); hit_s = (ds == dss)
# When no_wflux, vertical walls are never the limiting step
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)
# 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. Mirroring the
# clamping already present in _calc_pos_analytical_1d breaks the 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)
# Box 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 — only step to the diagonal cell if it is ocean; otherwise the
# cardinal crossing already recorded in ib/jb is kept.
_ne_ocean = grid.kmt[jnp.clip(ia, 0, imt-1), jnp.clip(ja, 0, jmt-1)] > 0
_nw_ocean = grid.kmt[jnp.clip(iam-1, 0, imt-1), jnp.clip(ja, 0, jmt-1)] > 0
_se_ocean = grid.kmt[jnp.clip(ia, 0, imt-1), jnp.clip(ja-2, 0, jmt-1)] > 0
ib = jnp.where(hit_n & (ds == dse) & _ne_ocean, ia+1, ib)
ib = jnp.where(hit_n & (ds == dsw) & _nw_ocean, iam, ib)
jb = jnp.where(hit_e & (ds == dsn) & _ne_ocean, ja+1, jb)
jb = jnp.where(hit_e & (ds == dss) & _se_ocean, 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 — strict x1 > imt matches Fortran mod_pos_tstep.F90.
# x1 == imt is left unwrapped intentionally: both a westward crossing into
# cell imt (ib=imt, x1=imt) and an eastward crossing out of cell imt
# (ib=1, x1=imt) leave x1=imt; _calc_pos_1d handles r0=imt with ii=1
# via the r0_adj special case on the next sub-step.
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)
# Termination checks
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
return x1, y1, z1, ib, jb, kb, ds, still_active
# ===========================================================================
# Vectorised step over all particles — piecewise-linear scheme
# ===========================================================================
[docs]
def piecewise(grid, no_wflux: bool = False):
"""Return a JIT-compiled step function that advances all particles by one sub-cycle.
*grid* (a ``GridFields``) is closed over as XLA constants.
The returned function takes a ``FluxFields`` as its first dynamic argument,
enabling lazy loading: flux data need not be computed until the step
function is actually called.
Parameters
----------
grid : GridFields
Static grid geometry. Closed over by the compiled kernel.
no_wflux : bool
When True, vertical (w) fluxes are ignored. Automatically True
when ``fluxes.wflux is None`` at call time, but must be set here
to select the compiled variant.
"""
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(fluxes, state: ParticleState, intrpg: float, dsmin: float) -> ParticleState:
ib, jb, kb = state.ia, state.ja, state.ka
i0 = jnp.clip(ib - 1, 0, grid.imt - 1)
j0 = jnp.clip(jb - 1, 0, grid.jmt - 1)
k0 = jnp.clip(kb - 1, 0, grid.km - 1)
# Use the global time-interpolation weight (mean ts across all active
# particles, computed in _integrate_window.body_fn and passed in as
# `intrpg`) for all flux reads inside the particle kernel.
#
# A per-particle weight (mod(ts_i, 1.0)) changes by Δintrpg = Δts
# between sub-cycles. In GLORYS (and similar reanalysis data), the
# monthly flux can swing by orders of magnitude between slot 0 and
# slot 1, so even Δintrpg ≈ 1e-4 (one fast cell crossing) is enough
# to flip the interpolated flux sign at a wall. This causes the
# update_bounce logic to fire in alternating directions on consecutive
# sub-cycles, producing a "bounce-chain" oscillation with thousands
# of sub-cycles per window.
#
# The Fortran TRACMASS reference uses global (mean) intrpg; reproducing
# that here keeps the bounce stable while accepting a slight physical
# approximation for particles spread across the window.
#
# Per-particle intrpg_pp is kept for `dxyz` (layer-thickness
# interpolation) only, because that affects the step-size cap
# (dsmin_arr) rather than the crossing direction.
intrpg_pp = jnp.mod(state.ts, 1.0) # per-particle, for dxyz only
dxyz_raw = grid.dxdy[i0, j0] * (
intrpg_pp * fluxes.dzt[i0, j0, k0, 1] * fluxes.zstot[i0, j0, 1]
+ (1.0 - intrpg_pp) * 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
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, grid, 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, 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(
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
