Merge pull request #93 from csiro-coasts/get_all_geometry_names

Add `Convention.select_variables()`

Author: mx-moth

docs/releases/development.rst

@@ -13,3 +13,4 @@ Next release (in development)
   Add :attr:`.Convention.time_coordinate` and :attr:`.Convention.depth_coordinate`,
   deprecate :meth:`.Convention.get_times()` and :meth:`.Convention.get_depths()`
   (:pr:`92`).
+* Add :meth:`.Convention.select_variables` (:pr:`93`).

src/emsarray/conventions/_base.py

@@ -7,8 +7,8 @@
 import warnings
 from functools import cached_property
 from typing import (
-    TYPE_CHECKING, Any, Callable, Dict, FrozenSet, Generic, Hashable, List,
-    Optional, Tuple, TypeVar, Union, cast
+    TYPE_CHECKING, Any, Callable, Dict, FrozenSet, Generic, Hashable, Iterable,
+    List, Optional, Tuple, TypeVar, Union, cast
 )
 
 import numpy as np
@@ -1257,14 +1257,63 @@ def select_point(self, point: Point) -> xr.Dataset:
         return self.select_index(index.index)
 
     @abc.abstractmethod
+    def get_all_geometry_names(self) -> List[Hashable]:
+        """
+        Return a list of the names of all geometry variables used by this convention.
+
+        See Also
+        --------
+        drop_geometry
+        select_variables
+        """
+        pass
+
     def drop_geometry(self) -> xr.Dataset:
         """
         Return a new :class:`xarray.Dataset`
         with all geometry variables dropped.
         Useful when significantly transforming the dataset,
         such as :mod:`extracting point data <emsarray.operations.point_extraction>`.
+
+        See Also
+        --------
+        get_all_geometry_names
+        select_variables
         """
-        pass
+        return self.dataset.drop_vars(self.get_all_geometry_names())
+
+    def select_variables(self, variables: Iterable[Hashable]) -> xr.Dataset:
+        """Select only a subset of the variables in this dataset, dropping all others.
+
+        This will keep all coordinate variables and all geometry variables.
+
+        Parameters
+        ----------
+        variables : iterable of Hashable
+            The names of all data variables to select.
+
+        Returns
+        -------
+        xarray.DataArray
+            A new dataset with the same geometry and coordinates,
+            but only the selected data variables.
+
+        See also
+        --------
+        get_all_geometry_names
+        drop_geometry
+        """
+        all_vars = set(self.dataset.variables.keys())
+        keep_vars = {
+            *variables,
+            *self.get_all_geometry_names(),
+            *self.get_all_depth_names(),
+        }
+        try:
+            keep_vars.add(self.get_time_name())
+        except NoSuchCoordinateError:
+            pass
+        return self.dataset.drop_vars(all_vars - keep_vars)
 
     @abc.abstractmethod
     def make_clip_mask(

src/emsarray/conventions/arakawa_c.py

@@ -11,7 +11,7 @@
 import enum
 import logging
 from functools import cached_property
-from typing import Dict, Hashable, Optional, Tuple, cast
+from typing import Dict, Hashable, List, Optional, Tuple, cast
 
 import numpy as np
 import xarray as xr
@@ -307,8 +307,8 @@ def selector_for_index(self, index: ArakawaCIndex) -> Dict[Hashable, int]:
         topology = self._topology_for_grid_kind[kind]
         return {topology.j_dimension: j, topology.i_dimension: i}
 
-    def drop_geometry(self) -> xr.Dataset:
-        variables = [
+    def get_all_geometry_names(self) -> List[Hashable]:
+        return [
             self.face.longitude.name,
             self.face.latitude.name,
             self.node.longitude.name,
@@ -318,7 +318,6 @@ def drop_geometry(self) -> xr.Dataset:
             self.back.longitude.name,
             self.back.latitude.name,
         ]
-        return self.dataset.drop_vars(variables)
 
     def make_linear(self, data_array: xr.DataArray) -> xr.DataArray:
         kind, size = self.get_grid_kind_and_size(data_array)

src/emsarray/conventions/grid.py

@@ -11,7 +11,7 @@
 from contextlib import suppress
 from functools import cached_property
 from typing import (
-    Dict, Generic, Hashable, Optional, Tuple, Type, TypeVar, cast
+    Dict, Generic, Hashable, List, Optional, Tuple, Type, TypeVar, cast
 )
 
 import numpy as np
@@ -264,7 +264,10 @@ def unravel_index(
     def ravel_index(self, indices: CFGridIndex) -> int:
         return int(np.ravel_multi_index(indices, self.topology.shape))
 
-    def get_grid_kind_and_size(self, data_array: xr.DataArray) -> Tuple[CFGridKind, int]:
+    def get_grid_kind_and_size(
+        self,
+        data_array: xr.DataArray,
+    ) -> Tuple[CFGridKind, int]:
         expected = {self.topology.y_dimension, self.topology.x_dimension}
         dims = set(data_array.dims)
         if dims.issuperset(expected):
@@ -277,13 +280,27 @@ def selector_for_index(self, index: CFGridIndex) -> Dict[Hashable, int]:
         y, x = index
         return {self.topology.y_dimension: y, self.topology.x_dimension: x}
 
-    def drop_geometry(self) -> xr.Dataset:
-        dataset = self.dataset.drop_vars([
+    def get_all_geometry_names(self) -> List[Hashable]:
+        # Grid datasets contain latitude and longitude variables
+        # plus optional bounds variables.
+        names = [
             self.topology.longitude_name,
             self.topology.latitude_name,
-        ])
-        dataset.attrs.pop('Conventions', None)
+        ]
 
+        bounds_names: List[Optional[Hashable]] = [
+            self.topology.longitude.attrs.get('bounds', None),
+            self.topology.latitude.attrs.get('bounds', None),
+        ]
+        for bounds_name in bounds_names:
+            if bounds_name is not None and bounds_name in self.dataset.variables:
+                names.append(bounds_name)
+
+        return names
+
+    def drop_geometry(self) -> xr.Dataset:
+        dataset = super().drop_geometry()
+        dataset.attrs.pop('Conventions', None)
         return dataset
 
     def make_linear(self, data_array: xr.DataArray) -> xr.DataArray:

src/emsarray/conventions/ugrid.py

@@ -1315,34 +1315,34 @@ def integer_indices(data_array: xr.DataArray) -> np.ndarray:
         new_dataset = xr.open_mfdataset(mfdataset_paths, lock=False)
         return utils.dataset_like(dataset, new_dataset)
 
-    def drop_geometry(self) -> xr.Dataset:
-        dataset = self.dataset
+    def get_all_geometry_names(self) -> List[Hashable]:
         topology = self.topology
 
-        geometry_variables = [
+        names = [
             topology.mesh_variable.name,
             topology.face_node_connectivity.name,
             topology.node_x.name,
             topology.node_y.name,
         ]
         if topology.has_valid_face_edge_connectivity:
-            geometry_variables.append(topology.face_edge_connectivity.name)
+            names.append(topology.face_edge_connectivity.name)
         if topology.has_valid_face_face_connectivity:
-            geometry_variables.append(topology.face_face_connectivity.name)
+            names.append(topology.face_face_connectivity.name)
         if topology.has_valid_edge_node_connectivity:
-            geometry_variables.append(topology.edge_node_connectivity.name)
+            names.append(topology.edge_node_connectivity.name)
         if topology.has_valid_edge_face_connectivity:
-            geometry_variables.append(topology.edge_face_connectivity.name)
+            names.append(topology.edge_face_connectivity.name)
         if topology.edge_x is not None:
-            geometry_variables.append(topology.edge_x.name)
+            names.append(topology.edge_x.name)
         if topology.edge_y is not None:
-            geometry_variables.append(topology.edge_y.name)
+            names.append(topology.edge_y.name)
         if topology.face_x is not None:
-            geometry_variables.append(topology.face_x.name)
+            names.append(topology.face_x.name)
         if topology.face_y is not None:
-            geometry_variables.append(topology.face_y.name)
+            names.append(topology.face_y.name)
+        return names
 
-        dataset = dataset.drop_vars(geometry_variables)
+    def drop_geometry(self) -> xr.Dataset:
+        dataset = super().drop_geometry()
         dataset.attrs.pop('Conventions', None)
-
         return dataset

tests/conventions/test_base.py

@@ -4,7 +4,7 @@
 import enum
 import pathlib
 from functools import cached_property
-from typing import Dict, Hashable, Optional, Tuple
+from typing import Dict, Hashable, List, Optional, Tuple
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -54,6 +54,9 @@ def get_grid_kind_and_size(self, data_array: xr.DataArray) -> Tuple[SimpleGridKi
             return (SimpleGridKind.face, int(np.prod(self.shape)))
         raise ValueError("Invalid dimensions")
 
+    def get_all_geometry_names(self) -> List[Hashable]:
+        return ['x', 'y']
+
     def unravel_index(
         self,
         index: int,