Allow descending ordered coordinates in target grid (#29)

* Added sortby to target grid

* Output grid matches input target grid order

* Added coord order tests

* Linting fixes

* Make coord order tests more compact

* Use reindex_like for aligning coords, make most_common test more compact

---------

Co-authored-by: Bart Schilperoort <[email protected]>

Author: kjdoore

src/xarray_regrid/methods/conservative.py

@@ -58,17 +58,22 @@ def conservative_regrid(
     dim_order = list(target_ds.dims)
 
     coord_names = set(target_ds.coords).intersection(set(data.coords))
-    coords = {name: target_ds[name] for name in coord_names}
+    target_ds_sorted = target_ds.sortby(list(coord_names))
+    coords = {name: target_ds_sorted[name] for name in coord_names}
     data = data.sortby(list(coord_names))
 
     if isinstance(data, xr.Dataset):
-        return conservative_regrid_dataset(data, coords, latitude_coord).transpose(
-            *dim_order, ...
-        )
+        regridded_data = conservative_regrid_dataset(
+            data, coords, latitude_coord
+        ).transpose(*dim_order, ...)
     else:
-        return conservative_regrid_dataarray(data, coords, latitude_coord).transpose(
-            *dim_order, ...
-        )
+        regridded_data = conservative_regrid_dataarray(  # type: ignore
+            data, coords, latitude_coord
+        ).transpose(*dim_order, ...)
+
+    regridded_data = regridded_data.reindex_like(target_ds, copy=False)
+
+    return regridded_data
 
 
 def conservative_regrid_dataset(

src/xarray_regrid/methods/most_common.py

@@ -61,15 +61,18 @@ def most_common_wrapper(
         data = data.to_dataset(name=da_name)
 
     coords = utils.common_coords(data, target_ds)
+    target_ds_sorted = target_ds.sortby(list(coords))
     coord_size = [data[coord].size for coord in coords]
     mem_usage = np.prod(coord_size) * np.zeros((1,), dtype=np.int64).itemsize
 
     if max_mem is not None and mem_usage > max_mem:
         result = split_combine_most_common(
-            data=data, target_ds=target_ds, time_dim=time_dim, max_mem=max_mem
+            data=data, target_ds=target_ds_sorted, time_dim=time_dim, max_mem=max_mem
         )
     else:
-        result = most_common(data=data, target_ds=target_ds, time_dim=time_dim)
+        result = most_common(data=data, target_ds=target_ds_sorted, time_dim=time_dim)
+
+    result = result.reindex_like(target_ds, copy=False)
 
     if da_name is not None:
         return result[da_name]

tests/test_most_common.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 import xarray as xr
+from numpy.testing import assert_array_equal
 
 from xarray_regrid import Grid, create_regridding_dataset
 
@@ -95,3 +96,21 @@ def test_attrs_dataset(dummy_lc_data, dummy_target_grid):
     assert ds_regrid.attrs != {}
     assert ds_regrid.attrs == dummy_lc_data.attrs
     assert ds_regrid["longitude"].attrs == dummy_lc_data["longitude"].attrs
+
+
+@pytest.mark.parametrize("dataarray", [True, False])
+def test_coord_order_original(dummy_lc_data, dummy_target_grid, dataarray):
+    input_data = dummy_lc_data["lc"] if dataarray else dummy_lc_data
+    ds_regrid = input_data.regrid.most_common(dummy_target_grid)
+    assert_array_equal(ds_regrid["latitude"], dummy_target_grid["latitude"])
+    assert_array_equal(ds_regrid["longitude"], dummy_target_grid["longitude"])
+
+
+@pytest.mark.parametrize("coord", ["latitude", "longitude"])
+@pytest.mark.parametrize("dataarray", [True, False])
+def test_coord_order_reversed(dummy_lc_data, dummy_target_grid, coord, dataarray):
+    input_data = dummy_lc_data["lc"] if dataarray else dummy_lc_data
+    dummy_target_grid[coord] = list(reversed(dummy_target_grid[coord]))
+    ds_regrid = input_data.regrid.most_common(dummy_target_grid)
+    assert_array_equal(ds_regrid["latitude"], dummy_target_grid["latitude"])
+    assert_array_equal(ds_regrid["longitude"], dummy_target_grid["longitude"])

tests/test_regrid.py

@@ -1,7 +1,9 @@
+from copy import deepcopy
 from pathlib import Path
 
 import pytest
 import xarray as xr
+from numpy.testing import assert_array_equal
 
 import xarray_regrid
 
@@ -14,9 +16,15 @@
 }
 
 
+@pytest.fixture(scope="session")
+def load_input_data() -> xr.Dataset:
+    ds = xr.open_dataset(DATA_PATH / "era5_2m_dewpoint_temperature_2000_monthly.nc")
+    return ds.compute()
+
+
 @pytest.fixture
-def sample_input_data() -> xr.Dataset:
-    return xr.open_dataset(DATA_PATH / "era5_2m_dewpoint_temperature_2000_monthly.nc")
+def sample_input_data(load_input_data) -> xr.Dataset:
+    return deepcopy(load_input_data)
 
 
 @pytest.fixture
@@ -63,9 +71,15 @@ def test_basic_regridders_da(sample_input_data, sample_grid_ds, method, cdo_file
     xr.testing.assert_allclose(da_regrid.compute(), ds_cdo["d2m"].compute())
 
 
+@pytest.fixture(scope="session")
+def load_conservative_input_data() -> xr.Dataset:
+    ds = xr.open_dataset(DATA_PATH / "era5_total_precipitation_2020_monthly.nc")
+    return ds.compute()
+
+
 @pytest.fixture
-def conservative_input_data() -> xr.Dataset:
-    return xr.open_dataset(DATA_PATH / "era5_total_precipitation_2020_monthly.nc")
+def conservative_input_data(load_conservative_input_data) -> xr.Dataset:
+    return deepcopy(load_conservative_input_data)
 
 
 @pytest.fixture
@@ -156,3 +170,49 @@ def test_attrs_dataset_conservative(sample_input_data, sample_grid_ds):
     assert ds_regrid.attrs == sample_input_data.attrs
     assert ds_regrid["d2m"].attrs == sample_input_data["d2m"].attrs
     assert ds_regrid["longitude"].attrs == sample_input_data["longitude"].attrs
+
+
+class TestCoordOrder:
+    @pytest.mark.parametrize("method", ["linear", "nearest", "cubic"])
+    @pytest.mark.parametrize("dataarray", [True, False])
+    def test_original(self, sample_input_data, sample_grid_ds, method, dataarray):
+        input_data = sample_input_data["d2m"] if dataarray else sample_input_data
+        regridder = getattr(input_data.regrid, method)
+        ds_regrid = regridder(sample_grid_ds)
+        assert_array_equal(ds_regrid["latitude"], sample_grid_ds["latitude"])
+        assert_array_equal(ds_regrid["longitude"], sample_grid_ds["longitude"])
+
+    @pytest.mark.parametrize("coord", ["latitude", "longitude"])
+    @pytest.mark.parametrize("method", ["linear", "nearest", "cubic"])
+    @pytest.mark.parametrize("dataarray", [True, False])
+    def test_reversed(
+        self, sample_input_data, sample_grid_ds, method, coord, dataarray
+    ):
+        input_data = sample_input_data["d2m"] if dataarray else sample_input_data
+        regridder = getattr(input_data.regrid, method)
+        sample_grid_ds[coord] = list(reversed(sample_grid_ds[coord]))
+        ds_regrid = regridder(sample_grid_ds)
+        assert_array_equal(ds_regrid["latitude"], sample_grid_ds["latitude"])
+        assert_array_equal(ds_regrid["longitude"], sample_grid_ds["longitude"])
+
+    @pytest.mark.parametrize("dataarray", [True, False])
+    def test_conservative_original(self, sample_input_data, sample_grid_ds, dataarray):
+        input_data = sample_input_data["d2m"] if dataarray else sample_input_data
+        ds_regrid = input_data.regrid.conservative(
+            sample_grid_ds, latitude_coord="latitude"
+        )
+        assert_array_equal(ds_regrid["latitude"], sample_grid_ds["latitude"])
+        assert_array_equal(ds_regrid["longitude"], sample_grid_ds["longitude"])
+
+    @pytest.mark.parametrize("coord", ["latitude", "longitude"])
+    @pytest.mark.parametrize("dataarray", [True, False])
+    def test_conservative_reversed(
+        self, sample_input_data, sample_grid_ds, coord, dataarray
+    ):
+        input_data = sample_input_data["d2m"] if dataarray else sample_input_data
+        sample_grid_ds[coord] = list(reversed(sample_grid_ds[coord]))
+        ds_regrid = input_data.regrid.conservative(
+            sample_grid_ds, latitude_coord="latitude"
+        )
+        assert_array_equal(ds_regrid["latitude"], sample_grid_ds["latitude"])
+        assert_array_equal(ds_regrid["longitude"], sample_grid_ds["longitude"])