Add reusable subsetting API (#55)

* Miggrate to selctor API

* Add reusable subset examples

* Add an exmaple for saving and loading selector with pickle

* Add docstrings

xarray_subset_grid/grid.py

@@ -4,6 +4,8 @@
 import numpy as np
 import xarray as xr
 
+from xarray_subset_grid.selector import Selector
+
 FLOAT_MAX = np.finfo(np.float32).max
 FLOAT_MIN = np.finfo(np.float32).min
 
@@ -130,29 +132,71 @@ def subset_vertical_levels(
         return ds.sel(selection, method=method)
 
     @abstractmethod
+    def compute_polygon_subset_selector(
+        self, ds: xr.Dataset, polygon: list[tuple[float, float]]
+    ) -> Selector:
+        """Compute the subset selector for the polygon
+
+        This method will return a Selector that can be used to subset the
+        dataset to the polygon. The selector will contain all the logic needed to subset
+        a dataset with the same grid type to the polygon. Once returned, the user
+        can call the select method on the selector to subset the dataset as many times
+        as needed without recomputing the selector.
+        """
+        raise NotImplementedError()
+
+    def compute_bbox_subset_selector(
+        self,
+        ds: xr.Dataset,
+        bbox: tuple[float, float, float, float],
+    ) -> Selector:
+        """Compute the subset selector for the bounding box
+
+        This method will return a Selector that can be used to subset the
+        dataset to the bounding box. The selector will contain all the logic needed to subset
+        a dataset with the same grid type to the bounding box. Once returned, the user
+        can call the select method on the selector to subset the dataset as many times
+        as needed without recomputing the selector.
+        """
+        polygon = np.array(
+            [
+                [bbox[0], bbox[3]],
+                [bbox[0], bbox[1]],
+                [bbox[2], bbox[1]],
+                [bbox[2], bbox[3]],
+                [bbox[0], bbox[3]],
+            ]
+        )
+        return self.compute_polygon_subset_selector(ds, polygon)
+
     def subset_polygon(
         self, ds: xr.Dataset, polygon: list[tuple[float, float]] | np.ndarray
     ) -> xr.Dataset:
         """Subset the dataset to the grid
+
+        This is a conveinence method that will compute the subset selector
+        for the polygon and then apply it to the dataset. This is useful for
+        one off subsetting operations where the user does not want to keep
+        the selector around for later use.
+
         :param ds: The dataset to subset
         :param polygon: The polygon to subset to
         :return: The subsetted dataset
         """
-        return ds
+        selector = self.compute_polygon_subset_selector(ds, polygon)
+        return selector.select(ds)
 
     def subset_bbox(self, ds: xr.Dataset, bbox: tuple[float, float, float, float]) -> xr.Dataset:
         """Subset the dataset to the bounding box
+
+        This is a conveinence method that will compute the subset selector
+        for the polygon and then apply it to the dataset. This is useful for
+        one off subsetting operations where the user does not want to keep
+        the selector around for later use.
+
         :param ds: The dataset to subset
         :param bbox: The bounding box to subset to
         :return: The subsetted dataset
         """
-        polygon = np.array(
-            [
-                [bbox[0], bbox[3]],
-                [bbox[0], bbox[1]],
-                [bbox[2], bbox[1]],
-                [bbox[2], bbox[3]],
-                [bbox[0], bbox[3]],
-            ]
-        )
-        return self.subset_polygon(ds, polygon)
+        selector = self.compute_bbox_subset_selector(ds, bbox)
+        return selector.select(ds)

xarray_subset_grid/grids/regular_grid.py

@@ -95,14 +95,9 @@ def data_vars(self, ds: xr.Dataset) -> set[str]:
             and "longitude" in var.cf.coordinates
         }
 
-    def subset_polygon(
-        self, ds: xr.Dataset, polygon: list[tuple[float, float]] | np.ndarray
-    ) -> xr.Dataset:
-        """Subset the dataset to the grid
-        :param ds: The dataset to subset
-        :param polygon: The polygon to subset to
-        :return: The subsetted dataset
-        """
+    def compute_polygon_subset_selector(
+        self, ds: xr.Dataset, polygon: list[tuple[float, float]]
+    ) -> Selector:
         lat = ds.cf["latitude"]
         lon = ds.cf["longitude"]
 
@@ -111,14 +106,13 @@ def subset_polygon(
         polygon_mask = ray_tracing_numpy(x, lat.flat, polygon).reshape(lon.shape)
 
         selector = RegularGridPolygonSelector(polygon, polygon_mask)
-        return selector.select(ds)
+        return selector
 
-    def subset_bbox(self, ds: xr.Dataset, bbox: tuple[float, float, float, float]) -> xr.Dataset:
-        """Subset the dataset to the bounding box
-        :param ds: The dataset to subset
-        :param bbox: The bounding box to subset to
-        :return: The subsetted dataset
-        """
+    def compute_bbox_subset_selector(
+        self,
+        ds: xr.Dataset,
+        bbox: tuple[float, float, float, float],
+    ) -> Selector:
         bbox = normalize_bbox_x_coords(ds.cf["longitude"].values, bbox)
         selector = RegularGridBBoxSelector(bbox)
-        return selector.select(ds)
+        return selector

xarray_subset_grid/grids/regular_grid_2d.py

@@ -77,17 +77,11 @@ def data_vars(self, ds: xr.Dataset) -> set[str]:
             and "longitude" in var.cf.coordinates
         }
 
-    def subset_polygon(
-        self, ds: xr.Dataset, polygon: list[tuple[float, float]] | np.ndarray
-    ) -> xr.Dataset:
-        """Subset the dataset to the grid
-        :param ds: The dataset to subset
-        :param polygon: The polygon to subset to
-        :return: The subsetted dataset
-        """
+    def compute_polygon_subset_selector(
+        self, ds: xr.Dataset, polygon: list[tuple[float, float]]
+    ) -> Selector:
         lat = ds.cf["latitude"]
         lon = ds.cf["longitude"]
         subset_mask = compute_2d_subset_mask(lat=lat, lon=lon, polygon=polygon)
 
-        selector = RegularGrid2dSelector(polygon=polygon, subset_mask=subset_mask)
-        return selector.select(ds)
+        return RegularGrid2dSelector(polygon=polygon, subset_mask=subset_mask)

xarray_subset_grid/grids/sgrid.py

@@ -16,10 +16,14 @@ class SGridSelector(Selector):
     def __init__(
         self,
         polygon: list[tuple[float, float]] | np.ndarray,
+        grid_topology_key: str,
+        grid_topology: xr.DataArray,
         subset_masks: list[tuple[list[str], xr.DataArray]],
     ):
         super().__init__()
         self.polygon = polygon
+        self._grid_topology_key = grid_topology_key
+        self._grid_topology = grid_topology
         self._subset_masks = subset_masks
 
     def select(self, ds: xr.Dataset) -> xr.Dataset:
@@ -33,7 +37,6 @@ def select(self, ds: xr.Dataset) -> xr.Dataset:
 
             # Now we can use the mask to subset the data
             ds_subset = ds_subset[mask[0]].where(ds_subset.subset_mask, drop=True).drop_encoding()
-            ds_subset = ds_subset.drop_vars("subset_mask")
 
             # Add the subsetted dataset to the list for merging
             ds_out.append(ds_subset)
@@ -94,14 +97,9 @@ def data_vars(self, ds: xr.Dataset) -> set[str]:
 
         return {var for var in ds.data_vars if not set(ds[var].dims).isdisjoint(dims)}
 
-    def subset_polygon(
-        self, ds: xr.Dataset, polygon: list[tuple[float, float]] | np.ndarray
-    ) -> xr.Dataset:
-        """Subset the dataset to the grid
-        :param ds: The dataset to subset
-        :param polygon: The polygon to subset to
-        :return: The subsetted dataset
-        """
+    def compute_polygon_subset_selector(
+        self, ds: xr.Dataset, polygon: list[tuple[float, float]]
+    ) -> Selector:
         grid_topology_key = ds.cf.cf_roles["grid_topology"][0]
         grid_topology = ds[grid_topology_key]
         dims = _get_sgrid_dim_coord_names(grid_topology)
@@ -133,8 +131,12 @@ def subset_polygon(
 
             subset_masks.append((vars, subset_mask))
 
-        selector = SGridSelector(polygon=polygon, subset_masks=subset_masks)
-        return selector.select(ds)
+        return SGridSelector(
+            polygon=polygon,
+            grid_topology_key=grid_topology_key,
+            grid_topology=grid_topology,
+            subset_masks=subset_masks,
+        )
 
 
 def _get_sgrid_dim_coord_names(

xarray_subset_grid/grids/ugrid.py

@@ -164,14 +164,9 @@ def data_vars(self, ds: xr.Dataset) -> set[str]:
 
         return data_vars
 
-    def subset_polygon(
-        self, ds: xr.Dataset, polygon: list[tuple[float, float]] | np.ndarray
-    ) -> xr.Dataset:
-        """Subset the dataset to the grid
-        :param ds: The dataset to subset
-        :param polygon: The polygon to subset to
-        :return: The subsetted dataset
-        """
+    def compute_polygon_subset_selector(
+        self, ds: xr.Dataset, polygon: list[tuple[float, float]]
+    ) -> Selector:
         # For this grid type, we find all nodes that are connected to elements that are inside
         # the polygon. To do this, we first find all nodes that are inside the polygon and then
         # find all elements that are connected to those nodes.
@@ -259,7 +254,7 @@ def subset_polygon(
             if transpose_face_face_connectivity:
                 face_face_new = face_face_new.T
 
-        selector = UGridSelector(
+        return UGridSelector(
             polygon=polygon,
             node_dimension=node_dimension,
             selected_nodes=selected_nodes,
@@ -273,8 +268,6 @@ def subset_polygon(
             face_face_connectivity=face_face_new if has_face_face_connectivity else None,
         )
 
-        return selector.select(ds)
-
 
 def assign_ugrid_topology(
     ds: xr.Dataset,

xarray_subset_grid/selector.py

@@ -1,7 +1,29 @@
+from abc import abstractmethod
+
 import xarray as xr
 
 
 class Selector:
+    """Selector class to subset datasets
+
+    This is a base method that should be implemented by a subclass
+    to perform selection on a given xarray dataset with whatever
+    context or logic is desired by the implementation.
+
+    select should return a new xarray dataset that is a subset of the input dataset
+    and must be implemented by the subclass.
+    """
+
+    @abstractmethod
     def select(self, ds: xr.Dataset) -> xr.Dataset:
-        """Perform the selection on the dataset"""
-        return ds
+        """Perform the selection on the dataset
+
+        For example, a selector could hold predifined masks to apply to the dataset,
+        and the select method here would apply those masks to the dataset and return the result.
+        This workflow is useful because computing the masks can be expensive, and
+        we want to avoid recomputing them for every dataset that needs to be subsetted. It also
+        allows datasets that are non standard to be subset using information from manually or
+        otherwise standardized datasets..
+
+        """
+        raise NotImplementedError()