fixed test grid nc, added 1x padding for psi grid subsets

tests/test_grids/test_sgrid.py

@@ -2,13 +2,12 @@
 import fsspec
 import numpy as np
 import xarray as xr
+import os
 
 import xarray_subset_grid.accessor  # noqa: F401
-from xarray_subset_grid.utils import ray_tracing_numpy
+from ..test_utils import ray_tracing_numpy, get_test_file_dir
 # open dataset as zarr object using fsspec reference file system and xarray
 
-from xarray_subset_grid.utils import get_test_file_dir
-
 from gridded import Variable, VectorVariable
 
 test_dir = get_test_file_dir()
@@ -58,19 +57,29 @@ def test_polygon_subset():
     #Check that the subset rho/psi/u/v positional relationsip makes sense aka psi point is 'between' it's neighbor rho points
     #Note that this needs to be better generalized; it's not trivial to write a test that works in all potential cases.
     assert ds_subset['lon_rho'][0,0] < ds_subset['lon_psi'][0,0] and ds_subset['lon_rho'][0,1] > ds_subset['lon_psi'][0,0]
+    breakpoint()
 
     #ds_subset.temp_sur.isel(ocean_time=0).plot(x="lon_rho", y="lat_rho")
 
 def test_polygon_subset_2(): 
-    ds = xr.open_dataset(sample_sgrid_file)
-    polygon = np.array(
+    ds = xr.open_dataset(sample_sgrid_file, decode_times=False)
+    polygon = np.array([
         [6.5, 37.5],
         [6.5, 39.5],
         [9.5, 40.5],
         [8.5, 37.5],
         [6.5, 37.5]
-    )
+    ])
     ds_subset = ds.xsg.subset_polygon(polygon)
 
-    breakpoint()
+    #Check that the subset dataset has the correct dimensions given the original padding
+    assert ds_subset.sizes['eta_rho'] == ds_subset.sizes['eta_psi'] + 1
+    assert ds_subset.sizes['eta_u'] == ds_subset.sizes['eta_psi'] + 1
+    assert ds_subset.sizes['eta_v'] == ds_subset.sizes['eta_psi']
+    assert ds_subset.sizes['xi_rho'] == ds_subset.sizes['xi_psi'] + 1
+    assert ds_subset.sizes['xi_u'] == ds_subset.sizes['xi_psi']
+    assert ds_subset.sizes['xi_v'] == ds_subset.sizes['xi_psi'] + 1
+
+    assert ds_subset.lon_psi.min() <= 6.5 and ds_subset.lon_psi.max() >= 9.5
+    assert ds_subset.lat_psi.min() <= 37.5 and ds_subset.lat_psi.max() >= 40.5
 

tests/test_utils.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pytest
+import os
 
 from xarray_subset_grid.utils import (
     normalize_bbox_x_coords,
@@ -9,6 +10,13 @@
 
 # normalize_polygon_x_coords tests.
 
+def get_test_file_dir():
+    """
+    returns the test file dir path
+    """
+    test_file_dir = os.path.join(os.path.dirname(__file__), 'test_data')
+    return test_file_dir
+
 poly1_180 = np.array(
     [
         [-73, 41],

xarray_subset_grid/grids/sgrid.py

@@ -123,9 +123,16 @@ def compute_polygon_subset_selector(
 
         node_mask = compute_2d_subset_mask(lat=node_lat, lon=node_lon, polygon=polygon)
         msk = np.where(node_mask)
-        subset_masks.append(([node_coords[0], node_coords[1]], node_mask))
+        index_bounding_box = np.array([[msk[0].min(), msk[0].max()+1], [msk[1].min(), msk[1].max()+1]])
+        # quick and dirty add a 1 node pad around the psi mask. This is to ensure the entire polygon is covered.
+        index_bounding_box[0,0] = max(0, index_bounding_box[0,0] - 1)
+        index_bounding_box[0,1] = min(node_lon.shape[0], index_bounding_box[0,1] + 1)
+        index_bounding_box[1,0] = max(0, index_bounding_box[1,0] - 1)
+        index_bounding_box[1,1] = min(node_lon.shape[1], index_bounding_box[1,1] + 1)
+        node_mask[index_bounding_box[0][0]:index_bounding_box[0][1],
+                  index_bounding_box[1][0]:index_bounding_box[1][1]] = True
 
-        index_bounding_box = [[msk[0].min(), msk[0].max()], [msk[1].min(), msk[1].max()]]
+        subset_masks.append(([node_coords[0], node_coords[1]], node_mask))
         for s in ('face', 'edge1', 'edge2'):
             dims = grid_topology.attrs.get(f"{s}_dimensions", None)
             coords = grid_topology.attrs.get(f"{s}_coordinates", None).split()
@@ -143,8 +150,8 @@ def compute_polygon_subset_selector(
             arranged_padding = [padding[d] for d in lon.dims]
             arranged_padding = [0 if p == 'none' or p == 'low' else 1 for p in arranged_padding]
             mask = np.zeros(lon.shape, dtype=bool)
-            mask[index_bounding_box[0][0]:index_bounding_box[0][1] + arranged_padding[0] + 1,
-                 index_bounding_box[1][0]:index_bounding_box[1][1] + arranged_padding[1] + 1] = True
+            mask[index_bounding_box[0][0]:index_bounding_box[0][1] + arranged_padding[0],
+                 index_bounding_box[1][0]:index_bounding_box[1][1] + arranged_padding[1]] = True
             xr_mask = xr.DataArray(mask, dims=lon.dims)
             subset_masks.append(([coords[0], coords[1]], xr_mask))