From 69d86e096cf1e4a7dd6678429a58a6f52328ab6d Mon Sep 17 00:00:00 2001 From: Sam Levang Date: Fri, 20 Sep 2024 03:57:00 +0000 Subject: [PATCH 1/7] add weights sparsity --- pyproject.toml | 1 + src/xarray_regrid/methods/conservative.py | 31 ++++++++++++++++++++--- tests/test_regrid.py | 13 ++++++++-- 3 files changed, 40 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index aa0e178..5ea0b89 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "xarray", "flox", "scipy", + "sparse", ] [tool.hatch.build] diff --git a/src/xarray_regrid/methods/conservative.py b/src/xarray_regrid/methods/conservative.py index c620058..9c89b9d 100644 --- a/src/xarray_regrid/methods/conservative.py +++ b/src/xarray_regrid/methods/conservative.py @@ -5,6 +5,7 @@ import numpy as np import xarray as xr +from sparse import COO # type: ignore from xarray_regrid import utils @@ -125,9 +126,10 @@ def conservative_regrid_dataset( for array in data_vars.keys(): if coord in data_vars[array].dims: + var_weights = sparsify_weights(weights, data_vars[array]) data_vars[array], valid_fracs[array] = apply_weights( da=data_vars[array], - weights=weights, + weights=var_weights, coord=coord, valid_frac=valid_fracs[array], skipna=skipna, @@ -171,7 +173,9 @@ def apply_weights( # Renormalize the weights along this dim by the accumulated valid_frac # along previous dimensions if valid_frac.name != EMPTY_DA_NAME: - weights_norm = weights * valid_frac / valid_frac.mean(dim=[coord]) + weights_norm = weights * (valid_frac / valid_frac.mean(dim=[coord])).fillna( + 0 + ) da_reduced: xr.DataArray = xr.dot( da.fillna(0), weights_norm, dim=[coord], optimize=True @@ -180,7 +184,7 @@ def apply_weights( if skipna: weights_valid_sum: xr.DataArray = xr.dot( - weights_norm, notnull, dim=[coord], optimize=True + notnull, weights_norm, dim=[coord], optimize=True ) weights_valid_sum = weights_valid_sum.rename(coord_map) da_reduced /= weights_valid_sum.clip(1e-6, None) @@ -195,6 +199,13 @@ def apply_weights( valid_frac = valid_frac.rename(coord_map) valid_frac = valid_frac.clip(0, 1) + # In some cases, dot product of dask data and sparse weights fails + # to densify, which prevents future conversion to numpy + if da_reduced.chunks and isinstance(da_reduced.data._meta, COO): + da_reduced.data = da_reduced.data.map_blocks( + lambda x: x.todense(), dtype=da_reduced.dtype + ) + return da_reduced, valid_frac @@ -248,3 +259,17 @@ def lat_weight(latitude: np.ndarray, latitude_res: float) -> np.ndarray: lat = np.radians(latitude) h = np.sin(lat + dlat / 2) - np.sin(lat - dlat / 2) return h * dlat / (np.pi * 4) # type: ignore + + +def sparsify_weights(weights: xr.DataArray, da: xr.DataArray) -> xr.DataArray: + """Create a sparse version of the weights that matches the dtype and chunks + of the array to be regridded. Even though the weights can be constructed as + dense arrays, contraction is more efficient with sparse operations.""" + new_weights = weights.copy().astype(da.dtype) + if da.chunks: + chunks = {k: v for k, v in da.chunksizes.items() if k in weights.dims} + new_weights.data = new_weights.chunk(chunks).data.map_blocks(COO) + else: + new_weights.data = COO(weights.data) + + return new_weights diff --git a/tests/test_regrid.py b/tests/test_regrid.py index ccab1f4..ace71a1 100644 --- a/tests/test_regrid.py +++ b/tests/test_regrid.py @@ -92,10 +92,19 @@ def test_basic_regridders_da( xr.testing.assert_allclose(da_regrid, da_cdo, rtol=0.002, atol=2e-5) + +@pytest.mark.parametrize( + "chunks", + [None, {"time": 1}, {"longitude": 100, "latitude": 100}] +) def test_conservative_regridder( - conservative_input_data, conservative_sample_grid, cdo_comparison_data + conservative_input_data, + conservative_sample_grid, + cdo_comparison_data, + chunks, ): - ds_regrid = conservative_input_data.regrid.conservative( + input_data = conservative_input_data.chunk(chunks) + ds_regrid = input_data.regrid.conservative( conservative_sample_grid, latitude_coord="latitude" ) ds_cdo = cdo_comparison_data["conservative"] From 7497a129666713f5f364e51682e8b6fe1bb65b45 Mon Sep 17 00:00:00 2001 From: Sam Levang Date: Fri, 20 Sep 2024 14:28:33 +0000 Subject: [PATCH 2/7] add opt-einsum --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 5ea0b89..041db53 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ dependencies = [ "flox", "scipy", "sparse", + "opt-einsum", ] [tool.hatch.build] From ddc7a42b70a052d5ee82d3b2923b90caec27a578 Mon Sep 17 00:00:00 2001 From: Sam Levang Date: Fri, 20 Sep 2024 14:33:28 +0000 Subject: [PATCH 3/7] fix linting --- tests/test_regrid.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_regrid.py b/tests/test_regrid.py index ace71a1..b720a3a 100644 --- a/tests/test_regrid.py +++ b/tests/test_regrid.py @@ -92,10 +92,8 @@ def test_basic_regridders_da( xr.testing.assert_allclose(da_regrid, da_cdo, rtol=0.002, atol=2e-5) - @pytest.mark.parametrize( - "chunks", - [None, {"time": 1}, {"longitude": 100, "latitude": 100}] + "chunks", [{}, {"time": 1}, {"longitude": 100, "latitude": 100}] ) def test_conservative_regridder( conservative_input_data, From 28fead0008a0809fc54dec024da4d3b04c7e12fb Mon Sep 17 00:00:00 2001 From: Sam Levang Date: Sat, 21 Sep 2024 08:44:00 -0400 Subject: [PATCH 4/7] make sparse optional --- pyproject.toml | 9 ++++++--- src/xarray_regrid/methods/conservative.py | 24 +++++++++++++++++------ tests/test_regrid.py | 2 +- 3 files changed, 25 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 041db53..7ab0a8e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,8 +26,6 @@ dependencies = [ "xarray", "flox", "scipy", - "sparse", - "opt-einsum", ] [tool.hatch.build] @@ -41,10 +39,15 @@ Issues = "https://github.com/EXCITED-CO2/xarray-regrid/issues" Source = "https://github.com/EXCITED-CO2/xarray-regrid" [project.optional-dependencies] +accel = [ + "sparse", + "opt-einsum", +] benchmarking = [ "dask[distributed]", "matplotlib", "zarr", + "h5netcdf", "requests", "aiohttp", ] @@ -71,7 +74,7 @@ docs = [ # Required for ReadTheDocs path = "src/xarray_regrid/__init__.py" [tool.hatch.envs.default] -features = ["dev", "benchmarking"] +features = ["accel", "dev", "benchmarking"] [tool.hatch.envs.default.scripts] lint = [ diff --git a/src/xarray_regrid/methods/conservative.py b/src/xarray_regrid/methods/conservative.py index 30a96a4..8c384b6 100644 --- a/src/xarray_regrid/methods/conservative.py +++ b/src/xarray_regrid/methods/conservative.py @@ -5,7 +5,11 @@ import numpy as np import xarray as xr -from sparse import COO # type: ignore + +try: + import sparse # type: ignore +except ImportError: + sparse = None from xarray_regrid import utils @@ -126,7 +130,11 @@ def conservative_regrid_dataset( for array in data_vars.keys(): if coord in data_vars[array].dims: - var_weights = sparsify_weights(weights, data_vars[array]) + if sparse is not None: + var_weights = sparsify_weights(weights, data_vars[array]) + else: + var_weights = weights + data_vars[array], valid_fracs[array] = apply_weights( da=data_vars[array], weights=var_weights, @@ -200,8 +208,12 @@ def apply_weights( valid_frac = valid_frac.clip(0, 1) # In some cases, dot product of dask data and sparse weights fails - # to densify, which prevents future conversion to numpy - if da_reduced.chunks and isinstance(da_reduced.data._meta, COO): + # to automatically densify, which prevents future conversion to numpy + if ( + sparse is not None + and da_reduced.chunks + and isinstance(da_reduced.data._meta, sparse.COO) + ): da_reduced.data = da_reduced.data.map_blocks( lambda x: x.todense(), dtype=da_reduced.dtype ) @@ -268,8 +280,8 @@ def sparsify_weights(weights: xr.DataArray, da: xr.DataArray) -> xr.DataArray: new_weights = weights.copy().astype(da.dtype) if da.chunks: chunks = {k: v for k, v in da.chunksizes.items() if k in weights.dims} - new_weights.data = new_weights.chunk(chunks).data.map_blocks(COO) + new_weights.data = new_weights.chunk(chunks).data.map_blocks(sparse.COO) else: - new_weights.data = COO(weights.data) + new_weights.data = sparse.COO(weights.data) return new_weights diff --git a/tests/test_regrid.py b/tests/test_regrid.py index b720a3a..be9770a 100644 --- a/tests/test_regrid.py +++ b/tests/test_regrid.py @@ -208,7 +208,7 @@ def test_conservative_nan_thresholds_against_coarsen(nan_threshold): @pytest.mark.skipif(xesmf is None, reason="xesmf required") def test_conservative_nan_thresholds_against_xesmf(): - ds = xr.tutorial.open_dataset("ersstv5").sst.compute().isel(time=[0]) + ds = xr.tutorial.open_dataset("ersstv5").sst.isel(time=[0]).compute() ds = ds.rename(lon="longitude", lat="latitude") new_grid = xarray_regrid.Grid( north=90, From 319eb41e2503d6516e423eef044f989ed3c44843 Mon Sep 17 00:00:00 2001 From: Sam Levang Date: Sat, 21 Sep 2024 13:54:12 -0400 Subject: [PATCH 5/7] shuffle dependency groups --- environment.yml | 5 +---- pyproject.toml | 6 +++--- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/environment.yml b/environment.yml index 21ef4b8..65aa78a 100644 --- a/environment.yml +++ b/environment.yml @@ -7,7 +7,4 @@ dependencies: - xESMF - cdo - pip: - - "xarray-regrid" - - "cftime" - - "matplotlib" - - "dask[distributed]" \ No newline at end of file + - "xarray-regrid[accel,benchmarking,dev]" \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 7ab0a8e..dd97bdd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,14 +42,16 @@ Source = "https://github.com/EXCITED-CO2/xarray-regrid" accel = [ "sparse", "opt-einsum", + "dask[distributed]", ] benchmarking = [ - "dask[distributed]", "matplotlib", "zarr", "h5netcdf", "requests", "aiohttp", + "pooch", + "cftime", # required for decode time of test netCDF files ] dev = [ "hatch", @@ -57,9 +59,7 @@ dev = [ "mypy", "pytest", "pytest-cov", - "cftime", # required for decode time of test netCDF files "pandas-stubs", # Adds typing for pandas. - "cftime", ] docs = [ # Required for ReadTheDocs "myst_parser", From 8deb5e45b4afe681fd0f9aa8b1b81b1bbfa3c68c Mon Sep 17 00:00:00 2001 From: Sam Levang Date: Sat, 21 Sep 2024 14:19:41 -0400 Subject: [PATCH 6/7] add readme note, test other methods with varying chunks --- README.md | 13 +++++++++++++ tests/test_regrid.py | 16 +++++++++------- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 4b5b41e..f9e58b5 100644 --- a/README.md +++ b/README.md @@ -24,10 +24,23 @@ Regridding is a common operation in earth science and other fields. While xarray ## Installation +For a minimal install: ```console pip install xarray-regrid ``` +To improve performance in certain cases: +```console +pip install xarray-regrid[accel] +``` + +which includes optional extras such as: + - `dask`: parallelization over chunked data + - `sparse`: for performing conservative regridding using sparse weight matrices + - `opt-einsum`: optimized einsum routines used in conservative regridding + + Benchmarking varies across different hardware specifications, but the inclusion of these extras can often provide significant speedups. + ## Usage The xarray-regrid routines are accessed using the "regrid" accessor on an xarray Dataset: ```py diff --git a/tests/test_regrid.py b/tests/test_regrid.py index be9770a..d58fd14 100644 --- a/tests/test_regrid.py +++ b/tests/test_regrid.py @@ -20,6 +20,8 @@ "conservative": DATA_PATH / "cdo_conservative_64b.nc", } +CHUNK_SCHEMES = [{}, {"time": 1}, {"longitude": 100, "latitude": 100}] + @pytest.fixture(scope="session") def sample_input_data() -> xr.Dataset: @@ -71,30 +73,30 @@ def conservative_sample_grid(): @pytest.mark.parametrize("method", ["linear", "nearest"]) +@pytest.mark.parametrize("chunks", CHUNK_SCHEMES) def test_basic_regridders_ds( - sample_input_data, sample_grid_ds, cdo_comparison_data, method + sample_input_data, sample_grid_ds, cdo_comparison_data, method, chunks ): """Test the dataset regridders (except conservative).""" - regridder = getattr(sample_input_data.regrid, method) + regridder = getattr(sample_input_data.chunk(chunks).regrid, method) ds_regrid = regridder(sample_grid_ds) ds_cdo = cdo_comparison_data[method] xr.testing.assert_allclose(ds_regrid, ds_cdo, rtol=0.002, atol=2e-5) @pytest.mark.parametrize("method", ["linear", "nearest"]) +@pytest.mark.parametrize("chunks", CHUNK_SCHEMES) def test_basic_regridders_da( - sample_input_data, sample_grid_ds, cdo_comparison_data, method + sample_input_data, sample_grid_ds, cdo_comparison_data, method, chunks ): """Test the dataarray regridders (except conservative).""" - regridder = getattr(sample_input_data["d2m"].regrid, method) + regridder = getattr(sample_input_data["d2m"].chunk(chunks).regrid, method) da_regrid = regridder(sample_grid_ds) da_cdo = cdo_comparison_data[method]["d2m"] xr.testing.assert_allclose(da_regrid, da_cdo, rtol=0.002, atol=2e-5) -@pytest.mark.parametrize( - "chunks", [{}, {"time": 1}, {"longitude": 100, "latitude": 100}] -) +@pytest.mark.parametrize("chunks", CHUNK_SCHEMES) def test_conservative_regridder( conservative_input_data, conservative_sample_grid, From 784eb9869b9b08eaa1edbaa8b43b133426a60545 Mon Sep 17 00:00:00 2001 From: Sam Levang Date: Tue, 24 Sep 2024 09:24:33 -0400 Subject: [PATCH 7/7] update changelog --- CHANGELOG.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 282e910..c0f174d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/). ## Unreleased -# 0.3.0 (2024-09-05) +Added: + - If latitude/longitude coordinates are detected and the domain is global, apply automatic padding at the boundaries, which gives behavior more consistent with common tools like ESMF and CDO ([#45](https://github.com/xarray-contrib/xarray-regrid/pull/45)). + - Conservative regridding weights are converted to sparse matrices if the optional [sparse](https://github.com/pydata/sparse) package is installed, which improves compute and memory performance in most cases ([#49](https://github.com/xarray-contrib/xarray-regrid/pull/49)). + + +## 0.3.0 (2024-09-05) New contributors: - [@slevang](https://github.com/slevang)