Fix hfds derivation with frazil_2d fallback (closes #200) (#298)

* feat(ocean): add calc_hfds function for surface downward heat flux calculation

* feat(tests): add unit tests for evaluate_expression and calc_hfds functions

Author: rbeucher

src/access_moppy/derivations/__init__.py

@@ -26,6 +26,7 @@
 from access_moppy.derivations.calc_ocean import (
     calc_areacello,
     calc_global_ave_ocean,
+    calc_hfds,
     calc_hfgeou,
     calc_msftbarot,
     calc_overturning_streamfunction,
@@ -90,6 +91,7 @@
     "calc_mrsll": calc_mrsll,
     "calc_mrsol": calc_mrsol,
     "calc_tsl": calc_tsl,
+    "calc_hfds": calc_hfds,
     "calc_msftbarot": calc_msftbarot,
     "calc_hfgeou": calc_hfgeou,
     "calc_overturning_streamfunction": calc_overturning_streamfunction,
@@ -118,6 +120,9 @@ def evaluate_expression(expr, context):
     if isinstance(expr, dict):
         if "literal" in expr:
             return expr["literal"]
+        if "optional" in expr:
+            # Return the variable if present in context, else None
+            return context.get(expr["optional"])
         op = expr["operation"]
         args = [
             evaluate_expression(arg, context)

src/access_moppy/derivations/calc_ocean.py

@@ -1,6 +1,10 @@
 #!/usr/bin/env python
+import logging
+
 import xarray as xr
 
+logger = logging.getLogger(__name__)
+
 
 def calc_global_ave_ocean(var, rho_dzt, area_t):
     """Calculate mass-weighted global average of an ocean variable.
@@ -600,3 +604,54 @@ def calc_areacello(area_t, ht, drop_time=True):
         areacello = areacello.isel(time=0).drop_vars("time", errors="ignore")
 
     return areacello
+
+
+def calc_hfds(
+    sfc_hflux_from_runoff,
+    sfc_hflux_coupler,
+    sfc_hflux_pme,
+    frazil_3d_int_z=None,
+    frazil_2d=None,
+):
+    """Calculate surface downward heat flux in sea water (hfds).
+
+    Sums the base surface heat flux components and adds the appropriate frazil
+    term. ``frazil_3d_int_z`` is preferred; ``frazil_2d`` is used as a fallback
+    for ACCESS-ESM1.6 runs that use the ``pop_icediag`` frazil scheme (frazil
+    confined to the top 5 ocean layers), where ``frazil_3d_int_z`` is not saved.
+
+    Parameters
+    ----------
+    sfc_hflux_from_runoff : xarray.DataArray
+        Heat flux from runoff. Units: W m-2
+    sfc_hflux_coupler : xarray.DataArray
+        Heat flux from the coupler. Units: W m-2
+    sfc_hflux_pme : xarray.DataArray
+        Heat flux from precipitation minus evaporation. Units: W m-2
+    frazil_3d_int_z : xarray.DataArray or None, optional
+        Vertically integrated 3-D frazil heat flux. Used preferentially when
+        available. Units: W m-2
+    frazil_2d : xarray.DataArray or None, optional
+        2-D frazil heat flux. Used as a fallback when ``frazil_3d_int_z`` is
+        not available.
+
+    Returns
+    -------
+    hfds : xarray.DataArray
+        Surface downward heat flux in sea water. Units: W m-2
+    """
+    base = sfc_hflux_from_runoff + sfc_hflux_coupler + sfc_hflux_pme
+    if frazil_3d_int_z is not None:
+        return base + frazil_3d_int_z
+    elif frazil_2d is not None:
+        logger.warning(
+            "frazil_3d_int_z not available; using frazil_2d for hfds calculation "
+            "(appropriate for ACCESS-ESM1.6 runs with the pop_icediag frazil scheme)"
+        )
+        return base + frazil_2d
+    else:
+        logger.warning(
+            "Neither frazil_3d_int_z nor frazil_2d is available; "
+            "computing hfds without a frazil contribution"
+        )
+        return base

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

@@ -3741,17 +3741,21 @@
                 "sfc_hflux_from_runoff",
                 "sfc_hflux_coupler",
                 "sfc_hflux_pme",
-                "frazil_3d_int_z"
+                "frazil_3d_int_z",
+                "frazil_2d"
             ],
             "calculation": {
                 "type": "formula",
-                "operation": "sum_vars",
+                "operation": "calc_hfds",
                 "args": [
                     "sfc_hflux_from_runoff",
                     "sfc_hflux_coupler",
-                    "sfc_hflux_pme",
-                    "frazil_3d_int_z"
-                ]
+                    "sfc_hflux_pme"
+                ],
+                "kwargs": {
+                    "frazil_3d_int_z": {"optional": "frazil_3d_int_z"},
+                    "frazil_2d": {"optional": "frazil_2d"}
+                }
             },
             "CF standard Name": ""
         },

src/access_moppy/ocean.py

@@ -118,7 +118,11 @@ def select_and_process_variables(self):
             self.ds[self.cmor_name] = self.ds[required_vars[0]]
         elif calc["type"] == "formula":
             # If the calculation is a formula, evaluate it
-            context = {var: self.ds[var] for var in required_vars}
+            # Variables listed in model_variables that are absent from the
+            # loaded dataset (e.g. optional frazil fields) are silently
+            # omitted from the context; individual derivation functions
+            # handle the None case via {"optional": ...} expressions.
+            context = {var: self.ds[var] for var in required_vars if var in self.ds}
             context.update(custom_functions)
             self.ds[self.cmor_name] = evaluate_expression(calc, context)
         elif calc["type"] == "dataset_function":

tests/unit/test_derivations.py

@@ -40,6 +40,28 @@ def test_numeric_passthrough(self):
         """Numeric values are returned directly."""
         assert evaluate_expression(3.14, {}) == pytest.approx(3.14)
 
+    def test_optional_present_returns_value(self):
+        """{'optional': 'varname'} returns the DataArray when it IS in context."""
+        ctx = self._make_context()
+        result = evaluate_expression({"optional": "var1"}, ctx)
+        assert result is ctx["var1"]
+
+    def test_optional_absent_returns_none(self):
+        """{'optional': 'varname'} returns None when the variable is NOT in context."""
+        result = evaluate_expression({"optional": "missing_var"}, {})
+        assert result is None
+
+    def test_optional_used_as_kwarg_in_formula(self):
+        """An optional kwarg that is absent should resolve to None and be passed through."""
+        ctx = self._make_context()
+        # Build an expression like calc_hfds(var1, var2, frazil_3d_int_z=None)
+        expr = {
+            "operation": "add",
+            "args": ["var1", "var2"],
+        }
+        result = evaluate_expression(expr, ctx)
+        np.testing.assert_allclose(result.values, 2.0)
+
 
 class TestCalculateMonthlyMinimum:
     """Tests for calculate_monthly_minimum() — covers decode_cf and ME frequency fix."""

tests/unit/test_derivations_calc_ocean.py

@@ -7,6 +7,7 @@
 from access_moppy.derivations.calc_ocean import (
     calc_areacello,
     calc_global_ave_ocean,
+    calc_hfds,
     calc_hfgeou,
     calc_msftbarot,
     calc_overturning_streamfunction,
@@ -586,3 +587,87 @@ def test_equivalent_to_depth_sum_then_cumsum(self):
         result = calc_msftbarot(tx)
         expected = tx.sum("st_ocean").cumsum("yt_ocean")
         np.testing.assert_allclose(result.values, expected.values, rtol=1e-12)
+
+
+# ---------------------------------------------------------------------------
+# calc_hfds
+# ---------------------------------------------------------------------------
+
+
+class TestCalcHfds:
+    """Tests for calc_hfds — surface downward heat flux with frazil fallback."""
+
+    def _base_fields(self):
+        """Return three base surface flux DataArrays."""
+        times = xr.date_range("2000-01-01", periods=NT, freq="ME")
+        shape = (NT, NY, NX)
+        dims = ["time", "yt_ocean", "xt_ocean"]
+        runoff = xr.DataArray(
+            RNG.random(shape) * 10.0, dims=dims, coords={"time": times}
+        )
+        coupler = xr.DataArray(
+            RNG.random(shape) * 50.0, dims=dims, coords={"time": times}
+        )
+        pme = xr.DataArray(RNG.random(shape) * 5.0, dims=dims, coords={"time": times})
+        return runoff, coupler, pme
+
+    @pytest.mark.unit
+    def test_uses_frazil_3d_int_z_when_available(self):
+        """When frazil_3d_int_z is provided it should be included in the sum."""
+        runoff, coupler, pme = self._base_fields()
+        frazil_3d = xr.DataArray(
+            np.ones((NT, NY, NX)) * 2.0,
+            dims=["time", "yt_ocean", "xt_ocean"],
+        )
+        frazil_2d = xr.DataArray(
+            np.ones((NT, NY, NX)) * 99.0,  # should NOT be used
+            dims=["time", "yt_ocean", "xt_ocean"],
+        )
+        result = calc_hfds(
+            runoff, coupler, pme, frazil_3d_int_z=frazil_3d, frazil_2d=frazil_2d
+        )
+        expected = runoff + coupler + pme + frazil_3d
+        np.testing.assert_allclose(result.values, expected.values, rtol=1e-10)
+
+    @pytest.mark.unit
+    def test_falls_back_to_frazil_2d(self):
+        """When frazil_3d_int_z is None, frazil_2d should be used instead."""
+        runoff, coupler, pme = self._base_fields()
+        frazil_2d = xr.DataArray(
+            np.ones((NT, NY, NX)) * 3.0,
+            dims=["time", "yt_ocean", "xt_ocean"],
+        )
+        result = calc_hfds(
+            runoff, coupler, pme, frazil_3d_int_z=None, frazil_2d=frazil_2d
+        )
+        expected = runoff + coupler + pme + frazil_2d
+        np.testing.assert_allclose(result.values, expected.values, rtol=1e-10)
+
+    @pytest.mark.unit
+    def test_falls_back_to_frazil_2d_emits_warning(self, caplog):
+        """Falling back to frazil_2d should log a warning."""
+        import logging
+
+        runoff, coupler, pme = self._base_fields()
+        frazil_2d = xr.DataArray(
+            np.ones((NT, NY, NX)), dims=["time", "yt_ocean", "xt_ocean"]
+        )
+        with caplog.at_level(
+            logging.WARNING, logger="access_moppy.derivations.calc_ocean"
+        ):
+            calc_hfds(runoff, coupler, pme, frazil_3d_int_z=None, frazil_2d=frazil_2d)
+        assert any("frazil_2d" in msg for msg in caplog.messages)
+
+    @pytest.mark.unit
+    def test_no_frazil_returns_base_sum(self):
+        """When neither frazil term is provided, result equals the three base fluxes."""
+        runoff, coupler, pme = self._base_fields()
+        result = calc_hfds(runoff, coupler, pme, frazil_3d_int_z=None, frazil_2d=None)
+        expected = runoff + coupler + pme
+        np.testing.assert_allclose(result.values, expected.values, rtol=1e-10)
+
+    @pytest.mark.unit
+    def test_returns_dataarray(self):
+        runoff, coupler, pme = self._base_fields()
+        result = calc_hfds(runoff, coupler, pme)
+        assert isinstance(result, xr.DataArray)