Add GFs notebook, fix regular grid implementation for bounding boxes

examples/regular_grid_2d.ipynb

@@ -1344,13 +1344,6 @@
     "ds.xsg.grid.name"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "markdown",
    "metadata": {},

tests/test_utils.py

@@ -2,6 +2,7 @@
 import pytest
 
 from xarray_subset_grid.utils import (
+    normalize_bbox_x_coords,
     normalize_polygon_x_coords,
     ray_tracing_numpy,
 )
@@ -65,31 +66,28 @@ def test_normalize_x_coords(lons, poly, norm_poly):
 
     assert np.allclose(normalized_polygon, norm_poly)
 
-    # normalized_polygon = normalize_polygon_x_coords(x2, polygon)
-    # assert np.allclose(
-    #     normalized_polygon,
-    #     np.array(
-    #         [
-    #             [-126, 41],
-    #             [-126, 41],
-    #             [-110, 39],
-    #             [-70, 41],
-    #         ]
-    #     ),
-    # )
-
-    # normalized_polygon = normalize_polygon_x_coords(x3, polygon)
-    # assert np.allclose(
-    #     normalized_polygon,
-    #     np.array(
-    #         [
-    #             [234, 41],
-    #             [234, 41],
-    #             [250, 39],
-    #             [290, 41],
-    #         ]
-    #     ),
-    # )
+
+bbox1_180 = [-73, 39, -70, 41]
+bbox1_360 = [287, 39, 290, 41]
+bbox2_360 = [234, 39, 290, 41]
+bbox2_180 = [-126, 39, -70, 41]
+
+
+@pytest.mark.parametrize(
+    "lons, bbox, norm_bbox",
+    [
+        ([-85, -84, -83, 10], bbox1_180, bbox1_180),  # x1
+        ([60, 45, 85, 70], bbox1_180, bbox1_180),  # x2
+        ([190, 200, 220, 250, 260], bbox1_180, bbox1_360),  # x3
+        ([-85, -84, -83, 10], bbox2_360, bbox2_180),  # x1
+        ([60, 45, 85, 70], bbox2_360, bbox2_360),  # x2
+        ([190, 200, 220, 250, 260], bbox2_360, bbox2_360),  # x3
+    ],
+)
+def test_normalize_x_coords_bbox(lons, bbox, norm_bbox):
+    lons = np.array(lons)
+    normalized_polygon = normalize_bbox_x_coords(lons, bbox)
+    assert np.allclose(normalized_polygon, norm_bbox)
 
 
 def test_ray_tracing_numpy():

xarray_subset_grid/grids/regular_grid.py

@@ -2,7 +2,11 @@
 import xarray as xr
 
 from xarray_subset_grid.grid import Grid
-from xarray_subset_grid.utils import normalize_polygon_x_coords, ray_tracing_numpy
+from xarray_subset_grid.utils import (
+    normalize_bbox_x_coords,
+    normalize_polygon_x_coords,
+    ray_tracing_numpy,
+)
 
 
 class RegularGrid(Grid):
@@ -17,10 +21,9 @@ def recognize(ds: xr.Dataset) -> bool:
             return False
 
         # Make sure the coordinates are 1D and match
-        lat_dim = ds[lat[0]].dims
-        ndim = ds[lon[0]].ndim
-        lon_dim = ds[lon[0]].dims
-        return lat_dim == lon_dim and ndim == 1
+        lat_ndim = ds[lat[0]].ndim
+        lon_ndim = ds[lon[0]].ndim
+        return lat_ndim == lon_ndim and lon_ndim == 1
 
     @property
     def name(self) -> str:
@@ -81,4 +84,5 @@ def subset_bbox(self, ds: xr.Dataset, bbox: tuple[float, float, float, float]) -
         :param bbox: The bounding box to subset to
         :return: The subsetted dataset
         """
+        bbox = normalize_bbox_x_coords(ds.cf["longitude"].values, bbox)
         return ds.cf.sel(lon=slice(bbox[0], bbox[2]), lat=slice(bbox[1], bbox[3]))

xarray_subset_grid/utils.py

@@ -37,6 +37,28 @@ def normalize_polygon_x_coords(x, poly):
     return poly
 
 
+def normalize_bbox_x_coords(x, bbox):
+    """Normalize the bbox x coordinates (longitude) to the
+    same coord system as used by given x coordinates.
+
+    e.g. If the longitude values are between 0 and 360, we need to normalize
+    the bbox x coordinates to be between 0 and 360. Vice versa if the
+    longitude values are between -180 and 180.
+    """
+    x_min, x_max = x.min(), x.max()
+
+    bbox_x_min, bbox_x_max = bbox[0], bbox[2]
+
+    if x_max > 180 and bbox_x_max < 0:
+        bbox_x_min += 360
+        bbox_x_max += 360
+    elif x_min < 0 and bbox_x_max > 180:
+        bbox_x_min -= 360
+        bbox_x_max -= 360
+
+    return bbox_x_min, bbox[1], bbox_x_max, bbox[3]
+
+
 def ray_tracing_numpy(x, y, poly):
     """Find vertices inside of the given polygon