Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added sortby to target grid #29

Merged
merged 6 commits into from
Feb 15, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 12 additions & 7 deletions src/xarray_regrid/methods/conservative.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
7 changes: 5 additions & 2 deletions src/xarray_regrid/methods/most_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
19 changes: 19 additions & 0 deletions tests/test_most_common.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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"])
68 changes: 64 additions & 4 deletions tests/test_regrid.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"])
Loading