Add rootd mapping for ACCESS-ESM1.6 (#343)

* Implement rootd calculation based on #335

* Add unit tests for calc_rootd

* Dynamically determine pseudo-level dimension in calc_rootd function

* Add calc_rootd to derivations module and custom functions

* Use positional vegetation tile selection in calc_rootd

* Cover calc_rootd missing pseudo_level branch

Author: rbeucher

src/access_moppy/derivations/__init__.py

@@ -18,6 +18,7 @@
     calc_mrsll,
     calc_mrsol,
     calc_nitrogen_pool_kg_m2,
+    calc_rootd,
     calc_topsoil,
     calc_tsl,
     extract_tilefrac,
@@ -93,6 +94,7 @@
     "calc_mrsfl": calc_mrsfl,
     "calc_mrsll": calc_mrsll,
     "calc_mrsol": calc_mrsol,
+    "calc_rootd": calc_rootd,
     "calc_tsl": calc_tsl,
     "calc_hfds": calc_hfds,
     "calc_msftbarot": calc_msftbarot,

src/access_moppy/derivations/calc_land.py

@@ -639,6 +639,67 @@ def calc_mrsol(var1):
     return result
 
 
+def calc_rootd(tilefrac):
+    """
+    Calculate maximum root depth (rootd) as a fixed (fx) field.
+
+    Since all plant functional types (PFTs, tiles 1-13) in ACCESS-ESM1.6 can
+    access the lowest (6th) soil layer, rootd equals the total soil column
+    depth of 4.6 m wherever any vegetated tile exists at any timestep.  Grid
+    cells with land but no vegetation (tiles 14-17 only) get 0 m, and
+    ocean/no-land cells get the missing value (NaN).
+
+    Agreed calculation from issue #335 (pseudocode by har917)::
+
+        rootd = 4.6  if any(tilefrac[t, i, lat, lon] > 0,  i=1..13)
+        rootd = 0    if all(tilefrac[t, i, lat, lon] == 0, i=1..13)
+                       and any(tilefrac[t, :, lat, lon] > 0)
+        rootd = NaN  if all(tilefrac[t, :, lat, lon] == 0)
+
+    Parameters
+    ----------
+    tilefrac : xarray.DataArray
+        Tile fraction variable (fractional, 0-1) with a pseudo-level
+        dimension as the second axis.  Corresponds to STASH ``fld_s03i317``.
+
+    Returns
+    -------
+    xarray.DataArray
+        Maximum root depth in metres (m) without a time dimension.
+    """
+    TOTAL_SOIL_DEPTH = 4.6  # metres — full depth of the 6-layer soil column
+
+    # Find the pseudo-level dimension dynamically (could be "pseudo_level_0" or "pseudo_level_1")
+    pseudo_level_dims = [d for d in tilefrac.dims if "pseudo_level" in d]
+    if not pseudo_level_dims:
+        raise ValueError(
+            f"No pseudo_level dimension found in tilefrac dims: {tilefrac.dims}"
+        )
+    pseudo_level = pseudo_level_dims[0]
+
+    # Vegetation occupies the first 13 slices of the tile axis. Use positional
+    # indexing so this works whether tile coordinates are 0-based, 1-based, or
+    # absent.
+    veg_tiles = tilefrac.isel({pseudo_level: slice(0, 13)})
+
+    # Lazy boolean reduction over tiles (and any remaining non-spatial dims)
+    has_veg = (veg_tiles > 0.0).any(dim=pseudo_level)
+    has_land = (tilefrac > 0.0).any(dim=pseudo_level)
+
+    # Reduce over all remaining non-spatial dimensions (e.g. time)
+    spatial_dims = {"lat", "lon", "latitude", "longitude"}
+    for dim in list(has_veg.dims):
+        if dim not in spatial_dims:
+            has_veg = has_veg.any(dim=dim)
+            has_land = has_land.any(dim=dim)
+
+    # Build rootd: 4.6 m where vegetated, 0 m where land-only, NaN where ocean
+    rootd = has_veg.astype(float) * TOTAL_SOIL_DEPTH
+    rootd = rootd.where(has_land)
+
+    return rootd
+
+
 def calc_tsl(var1):
     """
     This function takes soil temperature (fld_s08i225) and transforms the

src/access_moppy/mappings/ACCESS-ESM1.6_mappings.json

@@ -3404,6 +3404,28 @@
                 "fld_s03i395": "1"
             }
         },
+        "rootd": {
+            "dimensions": {
+                "lat": "lat",
+                "lon": "lon"
+            },
+            "units": "m",
+            "positive": null,
+            "model_variables": [
+                "fld_s03i317"
+            ],
+            "calculation": {
+                "type": "formula",
+                "operation": "calc_rootd",
+                "args": [
+                    "fld_s03i317"
+                ]
+            },
+            "CF standard Name": "root_depth",
+            "model_variable_units": {
+                "fld_s03i317": "1"
+            }
+        },
         "sftgif": {
             "dimensions": {
                 "lat": "lat",

tests/unit/test_derivations_calc_land.py

@@ -4,7 +4,7 @@
 import pytest
 import xarray as xr
 
-from access_moppy.derivations.calc_land import calc_snc
+from access_moppy.derivations.calc_land import calc_rootd, calc_snc
 
 # ---------------------------------------------------------------------------
 # Helpers
@@ -161,3 +161,146 @@ def test_mismatched_pseudo_level_name_handled(self):
         result = calc_snc(tf, sn, lf)
 
         np.testing.assert_allclose(result.values, 100.0)
+
+
+# ---------------------------------------------------------------------------
+# Helpers for calc_rootd
+# ---------------------------------------------------------------------------
+
+NT_R = 2  # time steps
+N_TILES = 17  # tiles 1-17 (pseudo_level_1 coordinate values)
+NJ_R = 3
+NI_R = 4
+
+
+def _make_tilefrac(data):
+    """Build a tilefrac DataArray with pseudo_level_1 coords 1-17."""
+    times = xr.date_range("2000-01-01", periods=NT_R, freq="ME")
+    return xr.DataArray(
+        data,
+        dims=["time", "pseudo_level_1", "lat", "lon"],
+        coords={
+            "time": times,
+            "pseudo_level_1": np.arange(1, N_TILES + 1),
+        },
+        attrs={"units": "1"},
+    )
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestCalcRootd:
+    """Tests for calc_rootd() — issue #335."""
+
+    def test_vegetated_returns_4p6(self):
+        """Grid cells with any vegetated tile (1-13) get 4.6 m."""
+        data = np.zeros((NT_R, N_TILES, NJ_R, NI_R))
+        data[:, 0, :, :] = 0.5  # tile 1 (index 0) is vegetated
+        tilefrac = _make_tilefrac(data)
+
+        result = calc_rootd(tilefrac)
+
+        assert result.dims == ("lat", "lon")
+        np.testing.assert_allclose(result.values, 4.6)
+
+    def test_non_veg_land_returns_zero(self):
+        """Land cells with only non-veg tiles (14-17) get 0 m."""
+        data = np.zeros((NT_R, N_TILES, NJ_R, NI_R))
+        data[:, 13, :, :] = 1.0  # tile 14 (index 13) — ice/urban/barren/lake
+        tilefrac = _make_tilefrac(data)
+
+        result = calc_rootd(tilefrac)
+
+        assert result.dims == ("lat", "lon")
+        np.testing.assert_allclose(result.values, 0.0)
+
+    def test_ocean_returns_nan(self):
+        """Cells with all-zero tile fractions (ocean) get NaN (missing)."""
+        data = np.zeros((NT_R, N_TILES, NJ_R, NI_R))
+        tilefrac = _make_tilefrac(data)
+
+        result = calc_rootd(tilefrac)
+
+        assert result.dims == ("lat", "lon")
+        assert np.all(np.isnan(result.values))
+
+    def test_mixed_grid(self):
+        """Spatial mix: veg / non-veg land / ocean all in the same array."""
+        data = np.zeros((NT_R, N_TILES, NJ_R, NI_R))
+        # lat 0: vegetated
+        data[:, 0, 0, :] = 0.8  # tile 1
+        # lat 1: non-veg land only (tile 14)
+        data[:, 13, 1, :] = 1.0
+        # lat 2: ocean — all zeros
+
+        tilefrac = _make_tilefrac(data)
+        result = calc_rootd(tilefrac)
+
+        np.testing.assert_allclose(result.values[0, :], 4.6)  # veg
+        np.testing.assert_allclose(result.values[1, :], 0.0)  # non-veg land
+        assert np.all(np.isnan(result.values[2, :]))  # ocean
+
+    def test_veg_appears_in_only_one_timestep(self):
+        """Vegetation present in any timestep is enough to report 4.6 m."""
+        data = np.zeros((NT_R, N_TILES, NJ_R, NI_R))
+        data[0, 2, :, :] = 0.3  # tile 3, timestep 0 only
+        tilefrac = _make_tilefrac(data)
+
+        result = calc_rootd(tilefrac)
+
+        np.testing.assert_allclose(result.values, 4.6)
+
+    def test_no_time_dimension_removed(self):
+        """Output has no time dimension regardless of input."""
+        data = np.zeros((NT_R, N_TILES, NJ_R, NI_R))
+        data[:, 0, :, :] = 0.5
+        tilefrac = _make_tilefrac(data)
+
+        result = calc_rootd(tilefrac)
+
+        assert "time" not in result.dims
+
+    def test_dask_compatible(self):
+        """Function returns a lazy dask-backed array without calling compute."""
+        pytest.importorskip("dask")
+
+        data = np.zeros((NT_R, N_TILES, NJ_R, NI_R))
+        data[:, 0, :, :] = 0.5
+        tilefrac = _make_tilefrac(data).chunk({"time": 1})
+
+        result = calc_rootd(tilefrac)
+
+        assert result.chunks is not None, "result should be dask-backed (lazy)"
+        np.testing.assert_allclose(result.compute().values, 4.6)
+
+    def test_pseudo_level_0_without_coord_is_supported(self):
+        """Regression: works when tile dimension is pseudo_level_0 with no coord."""
+        times = xr.date_range("2000-01-01", periods=NT_R, freq="ME")
+        data = np.zeros((NT_R, N_TILES, NJ_R, NI_R))
+        data[:, 0, :, :] = 0.5
+        tilefrac = xr.DataArray(
+            data,
+            dims=["time", "pseudo_level_0", "lat", "lon"],
+            coords={"time": times},
+        )
+
+        result = calc_rootd(tilefrac)
+
+        assert result.dims == ("lat", "lon")
+        np.testing.assert_allclose(result.values, 4.6)
+
+    def test_missing_pseudo_level_dimension_raises(self):
+        """Defensive path: raise when tilefrac has no pseudo_level dimension."""
+        times = xr.date_range("2000-01-01", periods=NT_R, freq="ME")
+        data = np.zeros((NT_R, NJ_R, NI_R))
+        tilefrac = xr.DataArray(
+            data,
+            dims=["time", "lat", "lon"],
+            coords={"time": times},
+        )
+
+        with pytest.raises(ValueError, match="No pseudo_level dimension found"):
+            calc_rootd(tilefrac)