# Grid and flux data ## The TRACMASS grid TRACMASS works on the staggered **Arakawa C-grid**. Tracers and cell volumes sit at the centre of each grid cell, and the velocity components sit at the centres of the cell faces: `u` on the east and west faces, `v` on the north and south faces, `w` on the faces between levels. TRACMASS does not use velocities directly. It uses **volume fluxes** through the faces (m³ s⁻¹), velocity times face area: $$ \begin{aligned} U_{i,j,k} &= u_{i,j,k}\,\Delta y_{i,j}\,\Delta z_{i,j,k}, \\ V_{i,j,k} &= v_{i,j,k}\,\Delta x_{i,j}\,\Delta z_{i,j,k}, \\ W_{i,j,k} &= w_{i,j,k}\,\Delta x_{i,j}\,\Delta y_{i,j}. \end{aligned} $$ A GCM adapter computes these fluxes from the model's velocities and grid metrics (see [GCM adapters](adapters.md)). Models on other Arakawa grids have to be moved onto the C-grid first. The Fortran documentation's [grid setup tutorial](https://www.tracmass.org/docs/tutorials/setup_grid.html) gives the A-to-C and B-to-C averaging formulas. The `scrippsgog` adapter shows how to move MITgcm's west/south face staggering onto TRACMASS's east/north faces. ```{image} ../_static/grid_cell.svg :class: diagram :alt: A C-grid cell with face fluxes and boxface numbers :width: 520px :align: center ``` ### Index conventions Positions are continuous coordinates in units of grid cells: - Cell `ia` (1-based, as in Fortran) covers `ia - 1 <= x <= ia`. Its west wall is at `x = ia - 1` and its east wall at `x = ia`, and likewise for `ja`/`y` and `ka`/`z`. - `uflux[i, j, k]` (0-based array index `i = ia - 1`) is the flux through the **east** face of cell `ia`, positive eastward. The west face of `ia` is `uflux[i - 1]`. With zonal periodicity (`iperio = 1`) the west face of the first column is the east face of the last. - `vflux[i, j, k]` is the flux through the **north** face of cell `ja`, positive northward. - **Vertical flux.** Adapters give `wflux` the way ocean models store it: `wflux[i, j, k]` is the flux through the **top** face of level `k`, positive **upward** (NEMO `depthw`, MITgcm `WVEL`). Inside `LazyFluxFields` and `FluxFields` it is converted to the kernel's convention: the face between levels `ka` and `ka + 1` (the **bottom** face of `ka`), positive towards increasing `k`, i.e. downward. That is `w_kernel[..., k] = -wflux[..., k + 1]`, with zero below the deepest level. `wflux` is `None` in runs without vertical motion. - `kmt[i, j]` is the number of ocean levels in each column. `kmt == 0` marks land. - Levels are numbered from the top: `k = 1` is the first model level, and a level `k > kmt` is below the sea floor. The "boxface" numbers in the figure are those of the Fortran code (`mod_pos_tstep.F90`): 1 east, 2 west, 3 north, 4 south, 5 and 6 for the two vertical faces. Saved trajectories use the same coordinates. Cell `i` of a 0-based array covers `i <= x < i + 1`, with its centre at `x = i + 0.5`. `gcm.glorys.xy_to_lonlat` and `tracmass.trajs.xy_to_lonlat` convert these positions to longitude and latitude (see [Coordinates](../tools/coordinates.md)). ## The three-layer split A trajectory run needs three kinds of data, which change at very different rates: - the **grid geometry** (`kmt`, cell areas, sizes): fixed for the whole run; - the **velocity snapshots**: change at every GCM output time and can be tens of gigabytes per snapshot for a global model; - the **particle state**: changes at every sub-step. TRACMASS-JAX keeps the first two in three objects, so the compiled kernel only has to be built once and flux data is only read when it is needed: ```{image} ../_static/architecture.svg :class: diagram :alt: GCM adapter feeding GridFields and LazyFluxFields; compute produces FluxFields :width: 720px :align: center ``` ### `GridFields` Static geometry, built once per run by `new_grid` or `new_flux`: | Field | Shape / type | Description | |---|---|---| | `kmt` | `(imt, jmt)` int | Number of ocean levels; 0 on land | | `dxdy` | `(imt, jmt)` | Horizontal cell area (m²) | | `imt`, `jmt`, `km` | int | Grid size | | `timestep` | `pd.Timedelta` | GCM output interval, always positive | | `reverse` | bool | `True` for a backward run | | `iperio`, `jperio` | int | Zonal periodicity, north fold (`jperio = 1` is not supported yet) | | `i_offset`, `j_offset` | int | Where the grid starts in the adapter's full grid (subset `i1`, `j1`) | `GridFields` is a JAX pytree with `kmt` and `dxdy` as data. The kernel factories (`advection_scheme.piecewise(grid)`, `analytical(grid, tseas)`) close over it, so `kmt`, `dxdy` and the sizes become constants in the compiled XLA program. ### `LazyFluxFields` The two-snapshot window of flux data, **not yet read**. It holds `_past` (slot 0) and `_present` (slot 1) dictionaries of `xarray.DataArray`s. For the reanalysis adapters these are Dask-backed, so no data has been read from disk. It also carries the run's time stepping: | Field | Description | |---|---| | `dtm` | Time of the past (slot 0) snapshot | | `timestep`, `reverse` | Copied from the grid | | `dtm_end` | End of the run, as given to `new_flux` | | `n_windows` | Windows left until `dtm_end`: `ceil(abs(dtm_end - dtm) / timestep)`, or `None` | | `fluxes_fn` | The adapter function, used by `update_flux` to open the next snapshot | Two operations act on it: - `lazy.compute()` reads both slots and returns a `FluxFields`. It keeps the present slot's arrays, so after the next `update_flux` the new past slot does not have to be read again. Each snapshot is read from disk once. - `update_flux(lazy)` moves the window one timestep: present becomes past (a dictionary reference swap, no I/O), and the adapter is called for the next snapshot, which again only opens it lazily. This is the counterpart of Fortran's `swap_time` in `mod_swap.F90`. For backward runs `update_flux` steps backward in time, and `uflux`, `vflux` and `wflux` are negated lazily, so particles move upstream. Both `new_flux` and `update_flux` also convert `wflux` to the kernel convention (see [Index conventions](#index-conventions)). The adapters always return forward fluxes. ### `FluxFields` The computed window, a JAX pytree passed to the kernel as a dynamic argument: | Field | Shape | Description | |---|---|---| | `uflux` | `(imt, jmt, km, 2)` | Zonal volume flux | | `vflux` | `(imt, jmt + 1, km, 2)` | Meridional volume flux | | `wflux` | `(imt, jmt, km, 2)` or `None` | Vertical volume flux, kernel convention (bottom face, positive down) | | `dzt` | `(imt, jmt, km, 2)` | Layer thickness | | `zstot` | `(imt, jmt, 2)` | Free-surface stretching factor (1 for fixed levels) | The trailing axis holds the two snapshots: index 0 is the past and index 1 the present. NaN and fill values in the fluxes are replaced by 0 when the data is read, which is the no-flow condition at land. The kernel accepts any `FluxFields` of the same shape and dtype without recompiling, so a new window costs only the read. ### The key invariant `new_flux` and `update_flux` never read flux data. Reading happens only in `compute()`, just before the kernel runs. After each window the integrator drops the computed arrays (`del concrete; gc.collect()`) so device memory is released before the next window is read. `compute()` warns if a window is larger than 8 GB, and suggests an `[env.gcm] subset`. ## Building a grid ```python import tracmass from tracmass import gcm # Grid and buffer in one call (slot 0 is fetched once and shared) grid, lazy = tracmass.new_flux(gcm.theoretical.fluxes, dtm_start="2007-06-01") # A GCM can be given by name; [env.gcm] subset is applied first grid, lazy = tracmass.new_flux("glorys", dtm_start="2010-06-01", dtm_end="2010-07-01") # Build the grid once and reuse it for several buffers grid = tracmass.new_grid("glorys", dtm_start="2010-06-01") lazy_june = tracmass.new_flux(grid, "glorys", dtm_start="2010-06-01") lazy_july = tracmass.new_flux(grid, "glorys", dtm_start="2010-07-01") # Advance one window by hand lazy = tracmass.update_flux(lazy) ``` With `dtm_end` earlier than `dtm_start` the run is backward: `grid.reverse` is `True` and the fluxes are negated. When you reuse a grid, `new_flux` checks that the `dtm_end` you pass agrees with `grid.reverse`.