Regridding to larger grid results in NaNs outside of data range (#33)

* Added mask to replace zeros outside of original data grid with NaNs

* Masked regridded regions outside of data range with NaN

* Linting fixes

* Most common NaN test

* Updated CHANGELOG

Author: kjdoore

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 Fixed:
  - Ensured all attributes are kept upon regridding (dataset, variable and coordinate attrs).
+ - Regridding to larger grid now result in NaNs at locations outside of starting data grid.
 
 Changed:
  - Moved to the Ruff formatter, instead of black.

src/xarray_regrid/methods/conservative.py

@@ -90,7 +90,13 @@ def conservative_regrid_dataset(
     da_attrs = [da.attrs for da in dataarrays]
     coord_attrs = [data[coord].attrs for coord in data_coords]
 
+    # track which target coordinate values are not covered by the source grid
+    uncovered_target_grid = {}
     for coord in coords:
+        uncovered_target_grid[coord] = (coords[coord] <= data[coord].max()) & (
+            coords[coord] >= data[coord].min()
+        )
+
         target_coords = coords[coord].to_numpy()
         source_coords = data[coord].to_numpy()
         weights = get_weights(source_coords, target_coords)
@@ -112,6 +118,10 @@ def conservative_regrid_dataset(
         da.attrs = attr
     regridded = xr.merge(dataarrays)
 
+    # Replace zeros outside of original data grid with NaNs
+    for coord in coords:
+        regridded = regridded.where(uncovered_target_grid[coord])
+
     regridded.attrs = attrs
 
     new_coords = [regridded[coord] for coord in data_coords]
@@ -133,6 +143,10 @@ def conservative_regrid_dataarray(
     coord_attrs = [data[coord].attrs for coord in data_coords]
 
     for coord in coords:
+        uncovered_target_grid = (coords[coord] <= data[coord].max()) & (
+            coords[coord] >= data[coord].min()
+        )
+
         if coord in data.coords:
             target_coords = coords[coord].to_numpy()
             source_coords = data[coord].to_numpy()
@@ -150,6 +164,9 @@ def conservative_regrid_dataarray(
             data = data.transpose(coord, ...)
             data = apply_weights(data, weights, coord, target_coords)
 
+            # Replace zeros outside of original data grid with NaNs
+            data = data.where(uncovered_target_grid)
+
     new_coords = [data[coord] for coord in data_coords]
     for coord, attr in zip(new_coords, coord_attrs, strict=True):
         coord.attrs = attr

src/xarray_regrid/methods/most_common.py

@@ -190,6 +190,13 @@ def most_common(data: xr.Dataset, target_ds: xr.Dataset, time_dim: str) -> xr.Da
     ds_regrid = ds_regrid.rename({f"{coord}_bins": coord for coord in coords})
     for coord in coords:
         ds_regrid[coord] = target_ds[coord]
+
+        # Replace zeros outside of original data grid with NaNs
+        uncovered_target_grid = (target_ds[coord] <= data[coord].max()) & (
+            target_ds[coord] >= data[coord].min()
+        )
+        ds_regrid = ds_regrid.where(uncovered_target_grid)
+
         ds_regrid[coord].attrs = coord_attrs[coord]
 
     return ds_regrid.transpose(*dim_order)

tests/test_most_common.py

@@ -51,6 +51,19 @@ def dummy_target_grid():
     return create_regridding_dataset(new_grid)
 
 
+@pytest.fixture
+def oversized_dummy_target_grid():
+    new_grid = Grid(
+        north=48,
+        east=48,
+        south=-8,
+        west=-8,
+        resolution_lat=8,
+        resolution_lon=8,
+    )
+    return create_regridding_dataset(new_grid)
+
+
 def test_most_common(dummy_lc_data, dummy_target_grid):
     expected_data = np.array(
         [
@@ -81,6 +94,38 @@ def test_most_common(dummy_lc_data, dummy_target_grid):
     )
 
 
+def test_oversized_most_common(dummy_lc_data, oversized_dummy_target_grid):
+    expected_data = np.array(
+        [
+            [np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN],
+            [np.NaN, 2, 2, 0, 0, 0, 0, np.NaN],
+            [np.NaN, 0, 0, 0, 0, 0, 0, np.NaN],
+            [np.NaN, 0, 0, 0, 0, 0, 0, np.NaN],
+            [np.NaN, 0, 0, 0, 0, 0, 0, np.NaN],
+            [np.NaN, 0, 0, 0, 0, 0, 0, np.NaN],
+            [np.NaN, 3, 3, 0, 0, 0, 1, np.NaN],
+            [np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN, np.NaN],
+        ]
+    )
+
+    lat_coords = np.linspace(-8, 48, num=8)
+    lon_coords = np.linspace(-8, 48, num=8)
+
+    expected = xr.Dataset(
+        data_vars={
+            "lc": (["longitude", "latitude"], expected_data),
+        },
+        coords={
+            "longitude": (["longitude"], lon_coords),
+            "latitude": (["latitude"], lat_coords),
+        },
+    )
+    xr.testing.assert_equal(
+        dummy_lc_data.regrid.most_common(oversized_dummy_target_grid)["lc"],
+        expected["lc"],
+    )
+
+
 def test_attrs_dataarray(dummy_lc_data, dummy_target_grid):
     dummy_lc_data["lc"].attrs = {"test": "testing"}
     da_regrid = dummy_lc_data["lc"].regrid.most_common(dummy_target_grid)