Changes to support subsetting datasets by constraining coordinate values.

coast/__init__.py

@@ -24,6 +24,7 @@
 from ._utils.process_data import Process_data
 from .data.opendap import OpendapInfo
 from .data.copernicus import Copernicus, Product
+from ._utils.coordinates import Coordinates2D, Coordinates3D, Coordinates4D, Coordinates
 
 # Set default for logging level when coast is imported
 import logging

coast/_utils/coordinates.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from numpy import number
+from numbers import Number
+from typing import Union, Optional
+
+
+Numeric = Optional[Union[Number, number]]
+
+
+@dataclass
+class Coordinates2D:
+    """Represent a point in one-to-two-dimensional space with optional X and Y coordinates."""
+
+    x: Numeric
+    y: Numeric
+
+
+@dataclass
+class Coordinates3D(Coordinates2D):
+    """Represent a point in one-to-three-dimensional space with optional X, Y, and Z coordinates."""
+
+    z: Numeric
+
+
+@dataclass
+class Coordinates4D(Coordinates3D):
+    """Represent a point in one-to-four-dimensional spacetime with optional X, Y, Z, and T coordinates."""
+
+    t: Numeric  # TODO Should this be a datetime or is it likely to be something like a Unix timestamp?
+
+
+Coordinates = Union[Coordinates2D, Coordinates3D, Coordinates4D]

coast/data/coast.py

@@ -5,6 +5,7 @@
 from dask.distributed import Client
 import copy
 from .._utils.logging_util import get_slug, debug, info, warn, warning
+from .._utils.coordinates import Coordinates, Coordinates3D, Coordinates4D, Numeric
 from .opendap import OpendapInfo
 
 
@@ -467,5 +468,165 @@ def plot_cartopy(self, var: str, plot_var: array, params, time_counter: int = 0)
         info("Displaying plot!")
         plt.show()
 
+    def set_constraint(self, start: Coordinates, end: Coordinates, drop: bool = True) -> None:
+        """Constrain the underlying dataset to values within an arbitrarily sized orthotope.
+
+        Args:
+            start: The start coordinates of the shape to define.
+            end: The end coordinates of the shape to define.
+            drop: Whether values should be dropped from the constrained dataset (if False, they will be NaNed).
+        """
+        self.dataset = self.constrain(start, end, drop=drop)
+
+    def constrain(self, start: Coordinates, end: Coordinates, drop: bool = True) -> xr.Dataset:
+        """Return the underlying dataset with values constrained to an arbitrarily sized orthotope.
+
+        Args:
+            start: The start coordinates of the shape to define.
+            end: The end coordinates of the shape to define.
+            drop: Whether values should be dropped from the constrained dataset (if False, they will be NaNed).
+
+        Returns:
+            The underlying dataset with values constrained to within the defined selection.
+        """
+        return constrain(self.dataset, start, end, drop=drop)
+
+    @property
+    def x_dim(self) -> xr.DataArray:
+        """Return the X coordinate array of the underlying dataset."""
+        return x_dim(self.dataset)
+
+    @property
+    def y_dim(self) -> xr.DataArray:
+        """Return the Y coordinate array of the underlying dataset."""
+        return y_dim(self.dataset)
+
+    @property
+    def z_dim(self) -> xr.DataArray:
+        """Return the Z coordinate array of the underlying dataset."""
+        return z_dim(self.dataset)
+
+    @property
+    def t_dim(self) -> xr.DataArray:
+        """Return the T[ime] coordinate array of the underlying dataset."""
+        return t_dim(self.dataset)
+
+    def get_coord(self, dim: str) -> xr.DataArray:
+        """Get the coordinate array for a dimension from the underlying dataset.
+
+        Args:
+            dim: The name of the dimension (i.e. "x", "y", "z", or "t").
+
+        Returns:
+            The corresponding coordinate array from the underlying dataset.
+        """
+        return get_coord(self.dataset, dim)
+
     def plot_movie(self):
         raise NotImplementedError
+
+
+def create_constraint(start: Numeric, end: Numeric, dim: xr.DataArray) -> np.typing.NDArray[bool]:
+    """Create a mask to exclude coordinates that do not fall within a range of two arbitrary values.
+
+    Args:
+        start: The start of the range of values to constrain within.
+        end: The end of the range of values ot constrain within.
+        dim: The coordinate array to constrain values from.
+
+    Returns:
+        A mask that can be applied to dim to exclude unwanted values.
+    """
+    return np.logical_and(dim >= start, dim <= end)
+
+
+def get_coord(dataset: xr.Dataset, dim: str) -> xr.DataArray:
+    """Get the coordinate array for a dimension in a dataset.
+
+    Args:
+        dataset: The dataset to interrogate.
+        dim: The name of the dimension (i.e. "x", "y", "z", or "t").
+
+    Returns:
+        The corresponding coordinate array from the provided dataset.
+    """
+    # TODO Really not a fan of this, is there an easier way to get the mapping?
+    return dataset[list(dataset[f"{dim.lower()}_dim"].coords)[0]]
+
+
+def x_dim(dataset: xr.Dataset) -> xr.DataArray:
+    """Get the X coordinate array for a dimension in a dataset.
+
+    Args:
+        dataset: The dataset to interrogate.
+
+    Returns:
+        The corresponding coordinate array from the provided dataset.
+    """
+    return get_coord(dataset, "x")
+
+
+def y_dim(dataset: xr.Dataset) -> xr.DataArray:
+    """Get the Y coordinate array for a dimension in a dataset.
+
+    Args:
+        dataset: The dataset to interrogate.
+
+    Returns:
+        The corresponding coordinate array from the provided dataset.
+    """
+    return get_coord(dataset, "y")
+
+
+def z_dim(dataset: xr.Dataset) -> xr.DataArray:
+    """Get the Z coordinate array for a dimension in a dataset.
+
+    Args:
+        dataset: The dataset to interrogate.
+
+    Returns:
+        The corresponding coordinate array from the provided dataset.
+    """
+    return get_coord(dataset, "z")
+
+
+def t_dim(dataset: xr.Dataset) -> xr.DataArray:
+    """Get the T[ime] coordinate array for a dimension in a dataset.
+
+    Args:
+        dataset: The dataset to interrogate.
+
+    Returns:
+        The corresponding coordinate array from the provided dataset.
+    """
+    return get_coord(dataset, "t")
+
+
+def constrain(dataset: xr.Dataset, start: Coordinates, end: Coordinates, drop: bool = True) -> xr.Dataset:
+    """Constrain values within a dataset to an arbitrarily sized orthotope.
+
+    Args:
+        dataset: The dataset to constrain values from.
+        start: The start coordinates of the shape to define.
+        end: The end coordinates of the shape to define.
+        drop: Whether values should be dropped from the constrained dataset (if False, they will be NaNed).
+
+    Returns:
+        The provided dataset with values constrained to within the defined selection.
+    """
+    assert type(start) == type(end), "Coordinates must be of the same dimensionality!"
+
+    constrained = dataset
+    if (x_start := start.x is not None) and (x_end := end.x is not None):
+        assert x_start == x_end, "Tried to constrain on X with a missing paired value!"
+        constrained = constrained.where(create_constraint(start.x, end.x, x_dim(constrained)), drop=drop)
+    if (y_start := start.y is not None) and (y_end := end.y is not None):
+        assert y_start == y_end, "Tried to constrain on Y with a missing paired value!"
+        constrained = constrained.where(create_constraint(start.y, end.y, y_dim(constrained)), drop=drop)
+    if isinstance(start, Coordinates3D) and (z_start := start.z is not None) and (z_end := end.z is not None):
+        assert z_start == z_end, "Tried to constrain on Z with a missing paired value!"
+        constrained = constrained.where(create_constraint(start.z, end.y, z_dim(constrained)), drop=drop)
+    if isinstance(start, Coordinates4D) and (t_start := start.t is not None) and (t_end := end.t is not None):
+        assert t_start == t_end, "Tried to constrain on Z with a missing paired value!"
+        constrained = constrained.where(create_constraint(start.t, end.t, t_dim(constrained)), drop=drop)
+    return constrained