tracmass.LazyFluxFields#

class tracmass.LazyFluxFields(_past, _present, dtm=None, fluxes_fn=None, zarr_chunk_xy=(1, 1), timestep=None, reverse=False, dtm_end=None, _computed=<factory>)[source]#

Bases: object

Two time-slots of flux data held as raw xarray DataArrays.

The DataArrays may be backed by Dask (for GLORYS from Zarr) or by in-memory arrays (for the theoretical adapter). In either case, no numpy or device memory is allocated until compute() is called.

This object is not a JAX pytree. Pass the result of compute() to the advection kernel instead.

The two time slots are stored as separate DataArrays without a trailing time dimension. The shape of each DataArray matches the spatial shape of the variable (e.g. (imt, jmt, km) for uflux). update_flux() advances the window by reassigning Python dict references — no Dask compute, no JAX, no GC.

Parameters:
_past#

Mapping {name: xr.DataArray} for the ‘past’ time slot.

Type:

dict

_present#

Mapping {name: xr.DataArray} for the ‘present’ time slot.

Type:

dict

dtm#

Timestamp of the ‘past’ (slot 0) time level. Set by new_flux() and update_flux(); None when constructed via lower-level helpers like swap_fields().

Type:

pd.Timestamp or None

fluxes_fn#

The GCM adapter callable used to create this object, e.g. gcm.glorys.fluxes. Stored by new_flux() and propagated by update_flux() so that time_dependent() can infer it automatically without the caller re-supplying it.

Type:

callable or None

timestep#

GCM output interval (always positive), copied from grid.timestep by new_flux(). update_flux() steps dtm by this amount.

Type:

pd.Timedelta or None

reverse#

True for a backward run (copied from grid.reverse): update_flux() then walks dtm backward and negates the velocity fluxes it loads.

Type:

bool

dtm_end#

The dtm_end passed to new_flux(), or None. Used by n_windows to size a run.

Type:

pd.Timestamp or None

Notes

compute() keeps the present slot’s arrays (already read and cleaned) in _computed, keyed by the slot’s DataArray. After update_flux() that DataArray is the new past slot, so the next compute() reuses the arrays instead of reading the same snapshot again: every snapshot is read once, not twice. _computed is shared by the copies update_flux() makes and holds one slot per variable.

__init__(_past, _present, dtm=None, fluxes_fn=None, zarr_chunk_xy=(1, 1), timestep=None, reverse=False, dtm_end=None, _computed=<factory>)#
Parameters:
Return type:

None

Methods

__init__(_past, _present[, dtm, fluxes_fn, ...])

compute([include_wflux])

Trigger Dask computation and transfer flux data to JAX device memory.

compute_patch(i0, i1, j0, j1, grid[, ...])

Load only the sub-patch [i0:i1, j0:j1] into device memory.

preload([include_wflux])

Return a copy with all Dask arrays preloaded to numpy.

Attributes

dtm

dtm_end

fluxes_fn

n_windows

Number of GCM windows left between dtm and dtm_end.

reverse

timestep

zarr_chunk_xy

property n_windows: int | None#

Number of GCM windows left between dtm and dtm_end.

ceil(|dtm_end - dtm| / timestep), so a dtm_end that falls inside a window still has that window integrated. Decreases by one with every update_flux() call. None when dtm_end (or dtm / timestep) is not set.

Returns:

int or None

compute_patch(i0, i1, j0, j1, grid, include_wflux=True)[source]#

Load only the sub-patch [i0:i1, j0:j1] into device memory.

Slices the lazy xarray DataArrays before calling .values, so Dask only computes (and reads from Zarr) the chunks that intersect the requested region. kmt and dxdy are sliced from grid and included in the returned LocalFluxFields as dynamic JAX arrays, so the advection kernel can be compiled once for a fixed patch shape and reused without recompilation as patch content changes.

Parameters:

  • i0 (int) – Zonal slice bounds (0-based, exclusive end).

  • i1 (int) – Zonal slice bounds (0-based, exclusive end).

  • j0 (int) – Meridional slice bounds (0-based, exclusive end).

  • j1 (int) – Meridional slice bounds (0-based, exclusive end).

  • grid (GridFields) – Full-grid static geometry; kmt and dxdy are sliced from here to match the patch.

  • include_wflux (bool, optional) – False skips reading wflux (returned as None), for runs without vertical motion. Default True.

Returns:

LocalFluxFields

Return type:

LocalFluxFields

preload(include_wflux=True)[source]#

Return a copy with all Dask arrays preloaded to numpy.

Issues a single dask.compute() call for every variable in _past and _present, replacing Dask-backed DataArrays with in-memory numpy-backed DataArrays. The returned object is otherwise identical and can be used anywhere a LazyFluxFields is expected.

Calling this once per GCM window in time_dependent_streaming() changes the I/O pattern from O(n_sub) small Dask computes (one per sub-cycle that has cell crossings) to O(1) per window. Subsequent extract_stencil() calls on the returned object use pure numpy fancy indexing — no further Zarr I/O regardless of how many sub-cycles or cell crossings occur.

Parameters:

include_wflux (bool, optional) – False leaves wflux out of the copy instead of reading it, for runs without vertical motion. Default True.

Returns:

LazyFluxFields – A copy with numpy-backed DataArrays and the same metadata (dtm, fluxes_fn, timestep, reverse, …).

Return type:

LazyFluxFields

compute(include_wflux=True)[source]#

Trigger Dask computation and transfer flux data to JAX device memory.

Variables are loaded one at a time (past then present) and stacked along a new trailing axis, so the peak extra allocation per variable is one numpy array of size t0, one of size t1, and the resulting JAX array of size t0 + t1.

NaN values in flux arrays (uflux, vflux, wflux) are replaced with 0.0 before stacking. Land cells in GLORYS/NEMO are masked as NaN in the Zarr files; the physical no-flow boundary condition is 0.0, so NaN → 0 is the correct substitution. Without this replacement, NaN fluxes propagate through _calc_pos_1d to produce a NaN particle position, which after astype(int32) conversion gives a garbage cell index that may point to a kmt=0 (land) cell, causing spurious REASON_LAND deactivations.

Parameters:

include_wflux (bool, optional) – False skips reading wflux (returned as None), for runs without vertical motion. Default True.

Returns:

FluxFields – Kernel-ready flux data with trailing slot dimension of size 2.

Return type:

FluxFields