Adding a GCM#

In Fortran TRACMASS, supporting a new model means writing a project folder with setup_grid.F90 (grid metrics, kmt) and read_field.F90 (reading velocities and turning them into fluxes), then recompiling. In TRACMASS-JAX it means writing one Python function, fluxes(dtm), which returns one snapshot as an xarray.Dataset. This tutorial writes a complete adapter for an analytical eddy, then shows what changes for real model output.

1. A minimal adapter#

Save this as eddy.py. It describes solid-body rotation on a closed 60 × 60 km box with 1 km cells, one turn every five days, and 6-hourly “snapshots”:

"""A solid-body-rotation eddy on a closed 60 x 60 x 1 box."""

import numpy as np
import pandas as pd
import xarray as xr

IMT, JMT, KM = 60, 60, 1
DX = DY = 1000.0                     # m
DZ = 10.0                            # m
OMEGA = 2 * np.pi / (5 * 86400.0)    # one turn every five days
TIMESTEP = pd.Timedelta(6, "h")
DTM0 = pd.Timestamp("2020-01-01")


def fluxes(dtm=None) -> xr.Dataset:
    """Return one snapshot of volume fluxes at time *dtm*."""
    dtm = DTM0 if dtm is None else pd.Timestamp(dtm)
    # Cell-centre coordinates in metres from the eddy centre.
    xc = (np.arange(IMT) + 0.5 - IMT / 2) * DX
    yc = (np.arange(JMT) + 0.5 - JMT / 2) * DY
    # Solid-body rotation, u = -omega*y on east faces, v = omega*x on north faces
    u = -OMEGA * yc[None, :] * np.ones((IMT, 1))
    v = OMEGA * xc[:, None] * np.ones((1, JMT))
    # Volume fluxes (m3/s), with a trailing size-1 time axis: (imt, jmt, km, 1)
    uflux = (u * DY * DZ)[:, :, None, None]
    vflux = (v * DX * DZ)[:, :, None, None]
    uflux[-1] = 0.0                  # closed east wall
    vflux[:, -1] = 0.0               # closed north wall
    return xr.Dataset(
        {
            "uflux": (("x", "y", "z", "t"), uflux),
            "vflux": (("x", "y", "z", "t"), vflux),
            "dzt": (("x", "y", "z", "t"), np.full((IMT, JMT, KM, 1), DZ)),
            "zstot": (("x", "y", "t"), np.ones((IMT, JMT, 1))),
            "kmt": (("x", "y"), np.full((IMT, JMT), KM, dtype=np.int32)),
            "dxdy": (("x", "y"), np.full((IMT, JMT), DX * DY)),
            "imt": IMT, "jmt": JMT, "km": KM,
            "iperio": 0, "jperio": 0,
            "timestep": TIMESTEP,
            "dtm": dtm,
            "reverse": False,
        }
    )

Run one particle 10 km east of the centre for one turn:

import numpy as np
import tracmass
from tracmass import seed
import eddy

grid, lazy = tracmass.new_flux(eddy.fluxes, dtm_start="2020-01-01",
                               dtm_end="2020-01-06")   # 5 days = 20 windows
state = seed.ijk(40, 40, 30, 30, 1, 1, grid=grid)
result = tracmass.time_dependent(state, grid, lazy, iter_sub=4,
                                 recorder="eddy.parquet")

df = tracmass.load("eddy.parquet")
r = np.hypot(df.x - 30, df.y - 30)
print(r.min(), r.max())        # 9.513... 9.529...: the radius is conserved

The particle stays on its circle and comes back to where it started after five days.

2. The dataset contract#

new_flux reads these variables from the dataset:

Variable	Dims	Notes
`uflux`	`(x, y, z, t)`	Flux through the east face of each cell, m³ s⁻¹, positive eastward
`vflux`	`(x, y, z, t)`	Flux through the north face, positive northward
`wflux`	`(x, y, z, t)`	Optional. Flux through the top face of each level (level 0 at the surface), positive upward, as NEMO and MITgcm store it. `new_flux` converts it for the kernels (see Grid)
`dzt`	`(x, y, z, t)`	Layer thickness, m
`zstot`	`(x, y, t)`	Free-surface stretching factor; 1 for fixed levels
`kmt`	`(x, y)`	Ocean levels per column, 0 on land
`dxdy`	`(x, y)`	Cell area, m²
`imt`, `jmt`, `km`, `iperio`, `jperio`	scalars	Grid size and periodicity
`timestep`	scalar	`pd.Timedelta` between snapshots
`dtm`	scalar	`pd.Timestamp` of this snapshot
`reverse`	scalar	Always `False`

Rules to follow:

Axis order x, y, z (longitude first). Transpose with ds.transpose("x", "y", ...) if your source is (z, y, x).
Add a trailing size-1 t axis to the time-varying fields, as above. new_flux removes a trailing size-1 dimension, so without a t axis a one-level grid (km = 1) would lose its z axis.
Fluxes, not velocities: multiply by the face area.
No NaN on purpose: NaN and fill values in the fluxes are replaced by 0 (no flow) when the data is read, but kmt must be 0 on land.
Forward fluxes only: backward runs are handled by new_flux.
dtm=None should give a sensible default date, which is used when dtm_start is not given.
The x and y dimensions must have the same name in every variable, or xarray will try to align them. Use .data on geometry arrays from another file, as the GLORYS adapter does.

3. Real model output#

For a real model, the adapter opens the model files lazily and builds the fluxes from the velocities and grid metrics. The structure of tracmass.gcm.glorys.fluxes is a good template:

def fluxes(dtm=None, include_wflux=True):
    vels = seagrid.mymodel.open_dataset(dtm)            # lazy, Dask-backed
    sub = _subset_isel()                                 # settings["subset"]
    vels = vels.isel(**sub)
    grd = seagrid.mymodel.grid_definition().isel(**sub)  # dx, dy, dz, kmt
    ds = xr.Dataset({
        "uflux": vels.u * grd.dzt.data * grd.dyu.data,
        "vflux": vels.v * grd.dzt.data * grd.dxv.data,
        "kmt": grd.kmt, "dxdy": grd.dxdy,
        "dzt": grd.dzt.expand_dims("t", axis=-1),
        "timestep": pd.Timedelta(1, "D"), "dtm": pd.Timestamp(dtm),
        "reverse": False,
    })
    if include_wflux:
        ds["wflux"] = vels.w * grd.dxdy.data
    ds["km"], ds["jmt"], ds["imt"] = ds.uflux.shape[:3]
    ds["iperio"], ds["jperio"] = 1, 0
    ds["zstot"] = grd.dxdy * 0 + 1
    ds.attrs.update(subset_offset_attrs(settings["subset"]))
    return ds.transpose("x", "y", ...)

Things to get right:

Staggering. TRACMASS expects uflux[i] on the east face and vflux[j] on the north face of cell (i, j). NEMO already stores them there. MITgcm stores U on the west face and V on the south face, so the scrippsgog adapter shifts them by one face. A one-cell shift is easy to miss and makes coastal particles freeze. tests/test_glorys_alignment.py shows how to test for it: the face flux must be zero exactly where a neighbouring cell is land.

Halos. Remove halo or ghost columns from velocities and grid metrics consistently. In GLORYS the velocity Zarrs and the mesh files have different halos.

Other Arakawa grids. Interpolate A- or B-grid velocities onto the C-grid faces first. The Fortran grid setup tutorial gives the averaging formulas.

Laziness. Keep the fields as Dask arrays and avoid .values. Mask fill values after cutting the subset, so only the chunks you need are touched.

Subsets. Keep a module-level settings = {"subset": {...}} dict with the keys i1, i2, istep, j1, j2, jstep, k1, k2, kstep. new_flux fills it from [env.gcm] subset before the first call. Record the subset start with ds.attrs.update(subset_offset_attrs(...)) so saved positions are absolute. Set iperio = 0 when the subset no longer wraps around.

Chunk hints. If the data is in Zarr, set ds.attrs["_zarr_chunk_x"] and ["_zarr_chunk_y"] to the storage chunk sizes. time_dependent_patched then lines its patches up with them.

4. Optional extras#

Add these functions to the module to enable more features:

def lonlat_grid(grid=None):
    """Cell-centre (lon, lat), each (imt, jmt), for the grid's subset."""

def lonlat_to_xy(lon, lat):
    """Geographic -> 0-based fractional model positions (full grid)."""

def xy_to_lonlat(x, y):
    """Model positions -> geographic (full grid)."""

tracmass.trajs.lonlat_to_xy and xy_to_lonlat do the work for any grid with 2-D lon/lat arrays (see Coordinates). With these, seed.latlon_box and the other lat/lon seeders and tracmass.load(lonlat=True) work for the new model.

5. Making it available by name#

To use [env.gcm] name = "mymodel" and new_flux("mymodel"), put the module in src/tracmass/gcm/mymodel.py (with the data-access code in src/tracmass/seagrid/mymodel.py) and import it in src/tracmass/gcm/__init__.py. Outside the package you can always pass the function itself: new_flux(mymodel.fluxes, ...).

Add a test that builds the grid, checks the face alignment against kmt, and runs a few windows. tests/test_scrippsgog.py is a compact example.