tracmass.advection_scheme#

Full-grid advection kernels for TRACMASS-JAX.

Implements the piecewise() and analytical() step-function factories that operate on the entire flux grid. Both close over a 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 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.

tracmass.advection_scheme.full_grid.analytical(grid, tseas, no_wflux=False)[source]#

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.

tracmass.advection_scheme.full_grid.piecewise(grid, no_wflux=False)[source]#

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.

Local-stencil advection kernels for TRACMASS-JAX.

Implements two kernel factories for memory-efficient integration:

  • piecewise_local() — patch-based advection for 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.

  • piecewise_stencil() — per-particle stencil advection for time_dependent_streaming(). Reads flux values from a pre-extracted FluxStencil instead of indexing into a global flux array.

Both share the piecewise crossing-time and position helpers from full_grid.

tracmass.advection_scheme.stencil.piecewise_local(imt, jmt, km, iperio, no_wflux=False)[source]#

Return a JIT-compiled step function for patch-based (local stencil) advection.

Unlike piecewise(), this factory closes over only the four scalar grid dimensions. kmt and dxdy are read from the fluxes argument at runtime (a LocalFluxFields), so the same compiled kernel can be reused across windows as the patch content changes — without recompilation.

Parameters:

  • imt (int) – Patch dimensions (fixed for the lifetime of the run).

  • jmt (int) – Patch dimensions (fixed for the lifetime of the run).

  • 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.

tracmass.advection_scheme.stencil.piecewise_stencil(grid, no_wflux=False)[source]#

Return a JIT-compiled step function for streaming (per-particle stencil) advection.

Unlike piecewise(), this kernel receives pre-extracted 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 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).