# 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":

```python
"""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:

```python
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](../components/grid.md#index-conventions)) |
| `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:

```python
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](https://www.tracmass.org/docs/tutorials/setup_grid.html)
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:

```python
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](../tools/coordinates.md)).
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.