Refactor level height transformation to use cli_level_to_height function and update calculation type in ACCESS-ESM1.6 mappings

Author: rbeucher

src/access_moppy/atmosphere.py

@@ -152,18 +152,17 @@ def select_and_process_variables(self):
                 ],
                 errors="ignore",
             )
+        elif calc["type"] == "dataset_function":
+            # Function that operates on the full dataset
+            func_name = calc["function"]
+            self.ds = self.ds.rename({required_vars[0]: self.cmor_name})
+            self.ds = custom_functions[func_name](self.ds, **calc.get("kwargs", {}))
         else:
             raise ValueError(f"Unsupported calculation type: {calc['type']}")
 
-        # Handle level coordinate transformation
-        if "theta_level_height" in self.ds:
-            self.ds = (self.ds.assign_coords({"lev": self.ds["theta_level_height"]})
-                            .swap_dims({"model_theta_level_number": "lev"})
-                            .drop_vars(["theta_level_height", "model_theta_level_number"], errors="ignore"))
-            axes_rename_map.pop("theta_level_height", None)  
-
         # Rename axes and bounds variables
-        self.ds = self.ds.rename({**bounds_rename_map, **axes_rename_map})
+        rename_map = {k: v for k, v in {**bounds_rename_map, **axes_rename_map}.items() if k in self.ds}
+        self.ds = self.ds.rename(rename_map)
 
         # Transpose the data variable according to the CMOR dimensions
         cmor_dims = re.sub(r'\w*level', 'lev', self.vocab.variable["dimensions"]).split()

src/access_moppy/derivations/__init__.py

@@ -1,7 +1,7 @@
 import operator
 from functools import reduce
 
-from access_moppy.derivations.calc_atmos import level_to_height, swap_dimensions
+from access_moppy.derivations.calc_atmos import cli_level_to_height
 from access_moppy.derivations.calc_land import (
     average_tile,
     calc_landcover,
@@ -19,12 +19,11 @@
     "mean": lambda *args: sum(args) / len(args),
     "kelvin_to_celsius": lambda x: x - 273.15,
     "celsius_to_kelvin": lambda x: x + 273.15,
-    "level_to_height": level_to_height,
+    "cli_level_to_height": cli_level_to_height,
     "calc_topsoil": calc_topsoil,
     "calc_landcover": calc_landcover,
     "extract_tilefrac": extract_tilefrac,
     "average_tile": average_tile,
-    "swap_dimensions": swap_dimensions,
 }
 
 

src/access_moppy/derivations/calc_atmos.py

@@ -247,48 +247,11 @@
 # ----------------------------------------------------------------------
 
 
-def level_to_height(var, levs=None):
-    """Returns model level variable with level height instead of
-    number as dimension
+def cli_level_to_height(ds):
+    # Handle level coordinate transformation
+    if "theta_level_height" in ds:
+        ds = (ds.assign_coords({"lev": ds["theta_level_height"]})
+                        .swap_dims({"model_theta_level_number": "lev"})
+                        .drop_vars(["theta_level_height", "model_theta_level_number"], errors="ignore"))
+    return ds
 
-    Parameters
-    ----------
-    var : Xarray DataArray
-        Variable defined on model levels number
-    levs : tuple(str,str)
-        slice of levels to apply (optional, default is None)
-
-    Returns
-    -------
-    vout : Xarray DataArray
-        Same variable defined on model levels height
-
-    """
-    zdim = var.dims
-    zdim_height = zdim.replace("number", "height").replace("model_", "")
-    var = var.swap_dims({zdim: zdim_height})
-    if levs is not None:
-        var = var.isel({zdim_height: slice(int(levs[0]), int(levs[1]))})
-    return var
-
-
-def swap_dimensions(var, in_dim, out_dim):
-    """Swap dimensions of an xarray DataArray.
-
-    Parameters
-    ----------
-    var : xarray.DataArray
-        The input DataArray.
-    in_dim : str
-        The name of the dimension to be replaced.
-    out_dim : str
-        The name of the new dimension.
-
-    Returns
-    -------
-    xarray.DataArray
-        The DataArray with swapped dimensions.
-    """
-    var = var.swap_dims({in_dim: out_dim})
-
-    return var

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

@@ -115,8 +115,9 @@
                 "formula": "z = a + b*orog"
             },
             "calculation": {
-                "type": "direct"
-
+                "type": "dataset_function",
+                "function": "cli_level_to_height",
+                "kwargs": {}
             }
         },
         "clivi": {