Spatial subsetting (#155)

* WIP

* Added spatial subset for both messages and grib2io xarray.

* fixed an introduced bug in unrelated code and reversed some format changes.

* Will only work on regular grids and clarified documentation.

src/grib2io/_grib2io.py

@@ -1334,6 +1334,139 @@ def interpolate(self, method, grid_def_out, method_options=None, drtn=None,
         return msg
 
 
+    def subset(self, lats, lons):
+        """
+        Return a spatial subset.
+
+        Currently only supports regular grids of the following types:
+
+        | Grid Type                                                    | gdtn  |
+        | :---:                                                        | :---: |
+        | Latitude/Longitude, Equidistant Cylindrical, or Plate Carree | 0     |
+        | Rotated Latitude/Longitude                                   | 1     |
+        | Mercator                                                     | 10    |
+        | Polar Stereographic                                          | 20    |
+        | Lambert Conformal                                            | 30    |
+        | Albers Equal-Area                                            | 31    |
+        | Gaussian Latitude/Longitude                                  | 40    |
+        | Equatorial Azimuthal Equidistant Projection                  | 110   |
+
+        Parameters
+        ----------
+        lats
+            List or tuple of latitudes.  The minimum and maximum latitudes will
+            be used to define the southern and northern boundaries.
+
+            The order of the latitudes is not important.  The function will
+            determine which is the minimum and maximum.
+
+            The latitudes should be in decimal degrees with 0.0 at the equator,
+            positive values in the northern hemisphere increasing to 90, and
+            negative values in the southern hemisphere decreasing to -90.
+        lons
+            List or tuple of longitudes.  The minimum and maximum longitudes
+            will be used to define the western and eastern boundaries.
+
+            The order of the longitudes is not important.  The function will
+            determine which is the minimum and maximum.
+
+            GRIB2 longitudes should be in decimal degrees with 0.0 at the prime
+            meridian, positive values increasing eastward to 360.  There are no
+            negative GRIB2 longitudes.
+
+            The typical west longitudes that start at 0.0 at the prime meridian
+            and decrease to -180 westward, are converted to GRIB2 longitudes by
+            '360 - (absolute value of the west longitude)' where typical
+            eastern longitudes are unchanged as GRIB2 longitudes.
+
+        Returns
+        -------
+        subset
+            A spatial subset of a GRIB2 message.
+        """
+        if self.gdtn not in [0, 1, 10, 20, 30, 31, 40, 110]:
+            raise ValueError(
+                """
+
+Subset only works for
+    Latitude/Longitude, Equidistant Cylindrical, or Plate Carree (gdtn=0)
+    Rotated Latitude/Longitude (gdtn=1)
+    Mercator (gdtn=10)
+    Polar Stereographic (gdtn=20)
+    Lambert Conformal (gdtn=30)
+    Albers Equal-Area (gdtn=31)
+    Gaussian Latitude/Longitude (gdtn=40)
+    Equatorial Azimuthal Equidistant Projection (gdtn=110)
+
+"""
+            )
+
+        if self.nx == 0 or self.ny == 0:
+            raise ValueError(
+                """
+
+Subset only works for regular grids.
+
+"""
+            )
+
+        newmsg = Grib2Message(
+            np.copy(self.section0),
+            np.copy(self.section1),
+            np.copy(self.section2),
+            np.copy(self.section3),
+            np.copy(self.section4),
+            np.copy(self.section5),
+        )
+
+        msglats, msglons = self.grid()
+
+        la1 = np.max(lats)
+        lo1 = np.min(lons)
+        la2 = np.min(lats)
+        lo2 = np.max(lons)
+
+        # Find the indices of the first and last grid points to the nearest
+        # lat/lon values in the grid.
+        first_lat = np.abs(msglats - la1)
+        first_lon = np.abs(msglons - lo1)
+        max_idx = np.maximum(first_lat, first_lon)
+        first_j, first_i = np.where(max_idx == np.min(max_idx))
+
+        last_lat = np.abs(msglats - la2)
+        last_lon = np.abs(msglons - lo2)
+        max_idx = np.maximum(last_lat, last_lon)
+        last_j, last_i = np.where(max_idx == np.min(max_idx))
+
+        setattr(newmsg, "latitudeFirstGridpoint", msglats[first_j[0], first_i[0]])
+        setattr(newmsg, "longitudeFirstGridpoint", msglons[first_j[0], first_i[0]])
+        setattr(newmsg, "nx", np.abs(first_i[0] - last_i[0]))
+        setattr(newmsg, "ny", np.abs(first_j[0] - last_j[0]))
+
+        # Set *LastGridpoint attributes even if only used for gdtn=[0, 1, 40].
+        # This information is used to subset xarray datasets and even though
+        # unnecessary for some supported grid types, it won't affect a grib2io
+        # message to set them.
+        setattr(newmsg, "latitudeLastGridpoint", msglats[last_j[0], last_i[0]])
+        setattr(newmsg, "longitudeLastGridpoint", msglons[last_j[0], last_i[0]])
+
+        setattr(
+            newmsg,
+            "data",
+            self.data[
+                min(first_j[0], last_j[0]) : max(first_j[0], last_j[0]),
+                min(first_i[0], last_i[0]) : max(first_i[0], last_i[0]),
+            ].copy(),
+        )
+
+        # Need to set the newmsg._sha1_section3 to a blank string so the grid
+        # method ignores the cached lat/lon values.  This will force the grid
+        # method to recompute the lat/lon values for the subsetted grid.
+        newmsg._sha1_section3 = ""
+        newmsg.grid()
+
+        return newmsg
+
     def validate(self):
         """
         Validate a complete GRIB2 message.

src/grib2io/xarray_backend.py

@@ -749,6 +749,29 @@ def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
             da.grib2io.to_grib2(filename, mode=mode)
             mode = "a"
 
+    def subset(self, lats, lons) -> xr.Dataset:
+        """
+        Subset the DataSet to a region defined by latitudes and longitudes.
+
+        Parameters
+        ----------
+        lats
+            Latitude bounds of the region.
+        lons
+            Longitude bounds of the region.
+
+        Returns
+        -------
+        subset
+            DataSet subset to the region.
+        """
+        ds = self._obj
+
+        newds = xr.Dataset()
+        for shortName in ds:
+            newds[shortName] = ds[shortName].grib2io.subset(lats, lons).copy()
+
+        return newds
 
 @xr.register_dataarray_accessor("grib2io")
 class Grib2ioDataArray:
@@ -1014,3 +1037,44 @@ def to_grib2(self, filename, mode: typing.Literal["x", "w", "a"] = "x"):
             with grib2io.open(filename, mode=mode) as f:
                 f.write(newmsg)
             mode = "a"
+
+    def subset(self, lats, lons) -> xr.DataArray:
+        """
+        Subset the DataArray to a region defined by latitudes and longitudes.
+
+        Parameters
+        ----------
+        lats
+            Latitude bounds of the region.
+        lons
+            Longitude bounds of the region.
+
+        Returns
+        -------
+        subset
+            DataArray subset to the region.
+        """
+        da = self._obj.copy(deep=True)
+
+        newmsg = Grib2Message(
+            da.attrs["GRIB2IO_section0"],
+            da.attrs["GRIB2IO_section1"],
+            da.attrs["GRIB2IO_section2"],
+            da.attrs["GRIB2IO_section3"],
+            da.attrs["GRIB2IO_section4"],
+            da.attrs["GRIB2IO_section5"],
+        )
+        newmsg.data = np.copy(da.values)
+
+        newmsg = newmsg.subset(lats, lons)
+
+        da.attrs["GRIB2IO_section3"] = newmsg.section3
+
+        mask_lat = (da.latitude >= newmsg.latitudeLastGridpoint) & (
+            da.latitude <= newmsg.latitudeFirstGridpoint
+        )
+        mask_lon = (da.longitude >= newmsg.longitudeFirstGridpoint) & (
+            da.longitude <= newmsg.longitudeLastGridpoint
+        )
+
+        return da.where((mask_lon & mask_lat).compute(), drop=True)

tests/test_subset.py

@@ -0,0 +1,86 @@
+import pytest
+import xarray as xr
+from numpy.testing import assert_array_equal
+
+import grib2io
+
+
+@pytest.fixture()
+def inp_ds(request):
+    datadir = request.config.rootdir / "tests" / "data" / "gfs_20221107"
+
+    filters = {
+        "typeOfFirstFixedSurface": 103,
+        "valueOfFirstFixedSurface": 2,
+        "productDefinitionTemplateNumber": 0,
+        "shortName": "TMP",
+    }
+
+    ids = xr.open_mfdataset(
+        [
+            datadir / "gfs.t00z.pgrb2.1p00.f009_subset",
+            datadir / "gfs.t00z.pgrb2.1p00.f012_subset",
+        ],
+        combine="nested",
+        concat_dim="leadTime",
+        engine="grib2io",
+        filters=filters,
+    )
+
+    yield ids
+
+
+@pytest.fixture()
+def inp_msgs(request):
+    datadir = request.config.rootdir / "tests" / "data" / "gfs_20221107"
+
+    with grib2io.open(datadir / "gfs.t00z.pgrb2.1p00.f012_subset") as imsgs:
+        yield imsgs
+
+
+@pytest.mark.parametrize(
+    "lats, lons, expected_section3",
+    [
+        pytest.param(
+            (43, 32.7),
+            (117, 79),
+            [
+                0,
+                380,
+                0,
+                0,
+                0,
+                6,
+                0,
+                0,
+                0,
+                0,
+                0,
+                0,
+                38,
+                10,
+                0,
+                -1,
+                43000000,
+                79000000,
+                48,
+                33000000,
+                117000000,
+                1000000,
+                1000000,
+                0,
+            ],
+            id="subset_1",
+        ),
+    ],
+)
+def test_message_subset(inp_msgs, inp_ds, lats, lons, expected_section3):
+    """Test subsetting a single DataArray to a single grib2 message."""
+    newmsg = inp_msgs[0].subset(lats=lats, lons=lons)
+    assert_array_equal(newmsg.section3, expected_section3)
+
+    newds = inp_ds["TMP"].grib2io.subset(lats=lats, lons=lons)
+    assert_array_equal(newds.attrs["GRIB2IO_section3"], expected_section3)
+
+    newds = inp_ds.grib2io.subset(lats=lats, lons=lons)
+    assert_array_equal(newds["TMP"].attrs["GRIB2IO_section3"], expected_section3)