Merge pull request #188 from csiro-coasts/159-ugrid-data-vars

Don't use Dataset.data_vars to access DataArrays in UGrid

Author: mx-moth

docs/releases/development.rst

@@ -12,3 +12,6 @@ Next release (in development)
 * Fixed :func:`emsarray.utils.datetime_from_np_time`
   when the system timezone is not UTC and a specific timezone is requested
   (:issue:`176`, :pr:`183`).
+* Fixed an issue with UGrid datasets when some of the mesh topology variables
+  are not present in Dataset.data_vars as they are detected as coordinate variables
+  (:issue:`159`, :pr:`188`).

src/emsarray/conventions/ugrid.py

@@ -378,10 +378,11 @@ def mesh_variable(self) -> xarray.DataArray:
         ``cf_role`` of ``"mesh_topology"``.
         """
         if self.topology_key is not None:
-            return self.dataset.data_vars[self.topology_key]
+            return self.dataset[self.topology_key]
         try:
+            data_arrays = (self.dataset[name] for name in self.dataset.variables.keys())
             return next(
-                data_array for data_array in self.dataset.data_vars.values()
+                data_array for data_array in data_arrays
                 if data_array.attrs.get('cf_role') == 'mesh_topology'
             )
         except StopIteration:
@@ -436,42 +437,42 @@ def _face_coordinates(self) -> tuple[Hashable, Hashable]:
     @property
     def node_x(self) -> xarray.DataArray:
         """Data array of node X / longitude coordinates."""
-        return self.dataset.data_vars[self._node_coordinates[0]]
+        return self.dataset[self._node_coordinates[0]]
 
     @property
     def node_y(self) -> xarray.DataArray:
         """Data array of node Y / latitude coordinates."""
-        return self.dataset.data_vars[self._node_coordinates[1]]
+        return self.dataset[self._node_coordinates[1]]
 
     @property
     def edge_x(self) -> xarray.DataArray | None:
         """Data array of characteristic edge X / longitude coordinates. Optional."""
         try:
-            return self.dataset.data_vars[self._edge_coordinates[0]]
+            return self.dataset[self._edge_coordinates[0]]
         except KeyError:
             return None
 
     @property
     def edge_y(self) -> xarray.DataArray | None:
         """Data array of characteristic edge y / latitude coordinates. Optional."""
         try:
-            return self.dataset.data_vars[self._edge_coordinates[1]]
+            return self.dataset[self._edge_coordinates[1]]
         except KeyError:
             return None
 
     @property
     def face_x(self) -> xarray.DataArray | None:
         """Data array of characteristic face x / longitude coordinates. Optional."""
         try:
-            return self.dataset.data_vars[self._face_coordinates[0]]
+            return self.dataset[self._face_coordinates[0]]
         except KeyError:
             return None
 
     @property
     def face_y(self) -> xarray.DataArray | None:
         """Data array of characteristic face y / latitude coordinates. Optional."""
         try:
-            return self.dataset.data_vars[self._face_coordinates[1]]
+            return self.dataset[self._face_coordinates[1]]
         except KeyError:
             return None
 
@@ -537,7 +538,7 @@ def has_valid_edge_node_connectivity(self) -> bool:
             return False
 
         try:
-            data_array = self.dataset.data_vars[self.mesh_attributes["edge_node_connectivity"]]
+            data_array = self.dataset[self.mesh_attributes["edge_node_connectivity"]]
         except KeyError:
             return False
 
@@ -570,7 +571,7 @@ def edge_node_connectivity(self) -> xarray.DataArray:
                 "No valid edge_node_connectivity defined")
 
         name = self.mesh_attributes['edge_node_connectivity']
-        return self.dataset.data_vars[name]
+        return self.dataset[name]
 
     @cached_property
     def edge_node_array(self) -> numpy.ndarray:
@@ -625,7 +626,7 @@ def has_valid_edge_face_connectivity(self) -> bool:
             return False
 
         try:
-            data_array = self.dataset.data_vars[self.mesh_attributes["edge_face_connectivity"]]
+            data_array = self.dataset[self.mesh_attributes["edge_face_connectivity"]]
         except KeyError:
             return False
 
@@ -650,7 +651,7 @@ def edge_face_connectivity(self) -> xarray.DataArray:
             raise NoConnectivityVariableException(
                 "No valid edge_face_connectivity defined")
         name = self.mesh_attributes['edge_face_connectivity']
-        return self.dataset.data_vars[name]
+        return self.dataset[name]
 
     @cached_property
     def edge_face_array(self) -> numpy.ndarray:
@@ -689,7 +690,7 @@ def has_valid_face_node_connectivity(self) -> bool:
         If this is invalid, the entire dataset is invalid.
         """
         try:
-            data_array = self.dataset.data_vars[self.mesh_attributes["face_node_connectivity"]]
+            data_array = self.dataset[self.mesh_attributes["face_node_connectivity"]]
         except KeyError:
             return False
 
@@ -713,7 +714,7 @@ def face_node_connectivity(self) -> xarray.DataArray:
         all the others can be derived from this if required.
         """
         name = self.mesh_attributes['face_node_connectivity']
-        return self.dataset.data_vars[name]
+        return self.dataset[name]
 
     @cached_property
     def face_node_array(self) -> numpy.ndarray:
@@ -734,7 +735,7 @@ def has_valid_face_edge_connectivity(self) -> bool:
         variable is incorrect.
         """
         try:
-            data_array = self.dataset.data_vars[self.mesh_attributes["face_edge_connectivity"]]
+            data_array = self.dataset[self.mesh_attributes["face_edge_connectivity"]]
         except KeyError:
             return False
 
@@ -773,7 +774,7 @@ def face_edge_connectivity(self) -> xarray.DataArray:
             raise NoConnectivityVariableException(
                 "No valid face_edge_connectivity defined")
         name = self.mesh_attributes['face_edge_connectivity']
-        return self.dataset.data_vars[name]
+        return self.dataset[name]
 
     @cached_property
     def face_edge_array(self) -> numpy.ndarray:
@@ -816,7 +817,7 @@ def has_valid_face_face_connectivity(self) -> bool:
         variable is incorrect.
         """
         try:
-            data_array = self.dataset.data_vars[self.mesh_attributes["face_face_connectivity"]]
+            data_array = self.dataset[self.mesh_attributes["face_face_connectivity"]]
         except KeyError:
             return False
 
@@ -841,7 +842,7 @@ def face_face_connectivity(self) -> xarray.DataArray:
             raise NoConnectivityVariableException(
                 "No valid face_face_connectivity defined")
         name = self.mesh_attributes['face_face_connectivity']
-        return self.dataset.data_vars[name]
+        return self.dataset[name]
 
     @cached_property
     def face_face_array(self) -> numpy.ndarray:
@@ -1201,7 +1202,7 @@ def apply_clip_mask(self, clip_mask: xarray.Dataset, work_dir: Pathish) -> xarra
         topology_variables: list[xarray.DataArray] = [topology.mesh_variable]
 
         # This is the fill value used in the mask.
-        new_fill_value = clip_mask.data_vars['new_node_index'].encoding['_FillValue']
+        new_fill_value = clip_mask['new_node_index'].encoding['_FillValue']
 
         def integer_indexes(data_array: xarray.DataArray) -> numpy.ndarray:
             masked_values = numpy.ma.masked_invalid(data_array.values)
@@ -1211,11 +1212,11 @@ def integer_indexes(data_array: xarray.DataArray) -> numpy.ndarray:
             masked_integers: numpy.ndarray = masked_values.astype(numpy.int_)
             return masked_integers
 
-        new_node_indexes = integer_indexes(clip_mask.data_vars['new_node_index'])
-        new_face_indexes = integer_indexes(clip_mask.data_vars['new_face_index'])
-        has_edges = 'new_edge_index' in clip_mask.data_vars
+        new_node_indexes = integer_indexes(clip_mask['new_node_index'])
+        new_face_indexes = integer_indexes(clip_mask['new_face_index'])
+        has_edges = 'new_edge_index' in clip_mask.variables
         if has_edges:
-            new_edge_indexes = integer_indexes(clip_mask.data_vars['new_edge_index'])
+            new_edge_indexes = integer_indexes(clip_mask['new_edge_index'])
 
         # Re-index the face_node_connectivity variable
         topology_variables.append(update_connectivity(
@@ -1273,7 +1274,8 @@ def integer_indexes(data_array: xarray.DataArray) -> numpy.ndarray:
             dimension_masks[topology.edge_dimension] = ~numpy.ma.getmask(new_edge_indexes)
         mesh_dimensions = set(dimension_masks.keys())
 
-        for name, data_array in dataset.data_vars.items():
+        for name in dataset.variables.keys():
+            data_array = dataset[name]
             data_array_path = work_path / (str(name) + '.nc')
             if name in topology_variable_names:
                 logger.debug("Skipping %r as it is a topology variable", name)