From 9290afd526f60728acb63eb373c63f0ffa1c6d76 Mon Sep 17 00:00:00 2001 From: David Huard Date: Thu, 6 Feb 2020 13:58:21 -0500 Subject: [PATCH 1/7] added test demonstrating interp bug for nd indexes sharing coordinate with array --- xarray/tests/test_interp.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index e3af8b5873a..5a1019c53c2 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -243,6 +243,28 @@ def test_interpolate_nd(case): actual = da.interp(y=ydest, x=xdest, method="linear") assert_allclose(actual.transpose("y", "z"), expected) +@pytest.mark.xfail +def test_interpolate_nd_nd(): + """Interpolate nd array with an nd indexer sharing coordinates.""" + # Create original array + a = [0, 2] + x = [0, 1, 2] + da = xr.DataArray(np.arange(2 * 3).reshape(2, 3), + dims=('a', 'x'), + coords={'a': a, 'x': x}) + + # Create indexer into `a` with dimensions (y, x) + y = [11, 12] + c = {'x': [0, ], 'y': y} + ia = xr.DataArray([[1], [2]], dims=('y', 'x'), coords=c) + + # This fails for now + out = da.interp(a=ia) + # But should be equivalent to + # da.sel(x=0).interp(a=is.sel(x=0)) + expected = xr.DataArray([[1.5], [3]], dims=('y', 'x'), coords=c) + assert_allclose(out, expected) + @pytest.mark.parametrize("method", ["linear"]) @pytest.mark.parametrize("case", [0, 1]) From 4acc7008748750b4511668e9c772d86e4c7efb70 Mon Sep 17 00:00:00 2001 From: David Huard Date: Thu, 6 Feb 2020 16:56:46 -0500 Subject: [PATCH 2/7] fix test so it works with sel --- xarray/core/missing.py | 3 +++ xarray/tests/test_interp.py | 21 ++++++++++----------- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 40f010b3514..29092f4ec92 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -616,6 +616,9 @@ def interp(var, indexes_coords, method, **kwargs): # target dimensions dims = list(indexes_coords) x, new_x = zip(*[indexes_coords[d] for d in dims]) + + # Shared dimensions between var and indexes coords + sdims = list(set(var.dims).intersection(*[set(nx.dims) for nx in new_x])) destination = broadcast_variables(*new_x) # transpose to make the interpolated axis to the last position diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index 5a1019c53c2..8209b1470ae 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -243,7 +243,8 @@ def test_interpolate_nd(case): actual = da.interp(y=ydest, x=xdest, method="linear") assert_allclose(actual.transpose("y", "z"), expected) -@pytest.mark.xfail + +#@pytest.mark.xfail def test_interpolate_nd_nd(): """Interpolate nd array with an nd indexer sharing coordinates.""" # Create original array @@ -254,16 +255,14 @@ def test_interpolate_nd_nd(): coords={'a': a, 'x': x}) # Create indexer into `a` with dimensions (y, x) - y = [11, 12] - c = {'x': [0, ], 'y': y} - ia = xr.DataArray([[1], [2]], dims=('y', 'x'), coords=c) - - # This fails for now - out = da.interp(a=ia) - # But should be equivalent to - # da.sel(x=0).interp(a=is.sel(x=0)) - expected = xr.DataArray([[1.5], [3]], dims=('y', 'x'), coords=c) - assert_allclose(out, expected) + # x coordinates should match da x coordinates. + y = [10, ] + c = {'x': x, 'y': y} + ia = xr.DataArray([[0, 2, 2]], dims=('y', 'x'), coords=c) + out = da.sel(a=ia) # This works + out = da.interp(a=ia) # This fails for now + expected = xr.DataArray([[0, 4, 5]], dims=('y', 'x'), coords=c) + xr.testing.assert_allclose(out.drop_vars('a'), expected) @pytest.mark.parametrize("method", ["linear"]) From c191150460f1017b7c81326bc633f50972a663de Mon Sep 17 00:00:00 2001 From: David Huard Date: Fri, 7 Feb 2020 10:06:54 -0500 Subject: [PATCH 3/7] support shared dimensions in interp --- xarray/core/dataset.py | 11 +++++++++++ xarray/core/missing.py | 3 --- xarray/tests/test_interp.py | 32 ++++++++++++++++++++------------ 3 files changed, 31 insertions(+), 15 deletions(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 07bea6dac19..562f789d912 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -2571,6 +2571,17 @@ def interp( coords = either_dict_or_kwargs(coords, coords_kwargs, "interp") indexers = dict(self._validate_interp_indexers(coords)) + if coords: + # This avoids broadcasting over coordinates that are both in + # the original array AND in the indexing array. It essentially + # forces interpolation along the shared coordinates. + sdims = ( + set(self.dims) + .intersection(*[set(nx.dims) for nx in indexers.values()]) + .difference(coords.keys()) + ) + indexers.update({d: self.variables[d] for d in sdims}) + obj = self if assume_sorted else self.sortby([k for k in coords]) def maybe_variable(obj, k): diff --git a/xarray/core/missing.py b/xarray/core/missing.py index 29092f4ec92..40f010b3514 100644 --- a/xarray/core/missing.py +++ b/xarray/core/missing.py @@ -616,9 +616,6 @@ def interp(var, indexes_coords, method, **kwargs): # target dimensions dims = list(indexes_coords) x, new_x = zip(*[indexes_coords[d] for d in dims]) - - # Shared dimensions between var and indexes coords - sdims = list(set(var.dims).intersection(*[set(nx.dims) for nx in new_x])) destination = broadcast_variables(*new_x) # transpose to make the interpolated axis to the last position diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index 8209b1470ae..13e4eb93b59 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -244,25 +244,33 @@ def test_interpolate_nd(case): assert_allclose(actual.transpose("y", "z"), expected) -#@pytest.mark.xfail def test_interpolate_nd_nd(): """Interpolate nd array with an nd indexer sharing coordinates.""" # Create original array a = [0, 2] x = [0, 1, 2] - da = xr.DataArray(np.arange(2 * 3).reshape(2, 3), - dims=('a', 'x'), - coords={'a': a, 'x': x}) + da = xr.DataArray( + np.arange(6).reshape(2, 3), dims=("a", "x"), coords={"a": a, "x": x} + ) # Create indexer into `a` with dimensions (y, x) - # x coordinates should match da x coordinates. - y = [10, ] - c = {'x': x, 'y': y} - ia = xr.DataArray([[0, 2, 2]], dims=('y', 'x'), coords=c) - out = da.sel(a=ia) # This works - out = da.interp(a=ia) # This fails for now - expected = xr.DataArray([[0, 4, 5]], dims=('y', 'x'), coords=c) - xr.testing.assert_allclose(out.drop_vars('a'), expected) + y = [10] + c = {"x": x, "y": y} + ia = xr.DataArray([[1, 2, 2]], dims=("y", "x"), coords=c) + out = da.interp(a=ia) + expected = xr.DataArray([[1.5, 4, 5]], dims=("y", "x"), coords=c) + xr.testing.assert_allclose(out.drop_vars("a"), expected) + + # If the *shared* indexing coordinates do not match, interp should fail. + with pytest.raises(ValueError): + c = {"x": [1], "y": y} + ia = xr.DataArray([[1]], dims=("y", "x"), coords=c) + da.interp(a=ia) + + with pytest.raises(ValueError): + c = {"x": [5, 6, 7], "y": y} + ia = xr.DataArray([[1]], dims=("y", "x"), coords=c) + da.interp(a=ia) @pytest.mark.parametrize("method", ["linear"]) From 5df6c9c0f99376dbc43f2f30567661ee49c00655 Mon Sep 17 00:00:00 2001 From: David Huard Date: Fri, 7 Feb 2020 10:09:34 -0500 Subject: [PATCH 4/7] isort fixes --- xarray/core/dataarray.py | 2 +- xarray/core/dataset.py | 4 ++-- xarray/tests/test_dask.py | 2 +- xarray/tests/test_interp.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 062cc6342df..551baf97a8a 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -22,7 +22,6 @@ import numpy as np import pandas as pd -from ..plot.plot import _PlotMethods from . import ( computation, dtypes, @@ -34,6 +33,7 @@ rolling, utils, ) +from ..plot.plot import _PlotMethods from .accessor_dt import CombinedDatetimelikeAccessor from .accessor_str import StringAccessor from .alignment import ( diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 562f789d912..fb5d44a892b 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -33,8 +33,6 @@ import xarray as xr -from ..coding.cftimeindex import _parse_array_of_cftime_strings -from ..plot.dataset_plot import _Dataset_PlotMethods from . import ( alignment, dtypes, @@ -47,6 +45,8 @@ rolling, utils, ) +from ..coding.cftimeindex import _parse_array_of_cftime_strings +from ..plot.dataset_plot import _Dataset_PlotMethods from .alignment import _broadcast_helper, _get_broadcast_dims_map_common_coords, align from .common import ( DataWithCoords, diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index cc554850839..361537e3569 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -16,7 +16,6 @@ from xarray.testing import assert_chunks_equal from xarray.tests import mock -from ..core.duck_array_ops import lazy_array_equiv from . import ( assert_allclose, assert_array_equal, @@ -26,6 +25,7 @@ raises_regex, requires_scipy_or_netCDF4, ) +from ..core.duck_array_ops import lazy_array_equiv from .test_backends import create_tmp_file dask = pytest.importorskip("dask") diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index 13e4eb93b59..b1bb2760d1a 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -5,8 +5,8 @@ import xarray as xr from xarray.tests import assert_allclose, assert_equal, requires_cftime, requires_scipy -from ..coding.cftimeindex import _parse_array_of_cftime_strings from . import has_dask, has_scipy +from ..coding.cftimeindex import _parse_array_of_cftime_strings from .test_dataset import create_test_data try: From 3974edb5b55015bf5c1ae358834e16ecb53adcce Mon Sep 17 00:00:00 2001 From: David Huard Date: Fri, 7 Feb 2020 10:14:19 -0500 Subject: [PATCH 5/7] update whats new --- doc/whats-new.rst | 3 +++ 1 file changed, 3 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index bf8e63eb926..94abb35d87d 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -27,6 +27,9 @@ New Features Bug fixes ~~~~~~~~~ +- Fix :py:meth:`Dataset.interp` when indexing array shares coordinates with the + indexed variable (:issue:`3252`). + By `David Huard `_. Documentation ~~~~~~~~~~~~~ From f3f9b19d2565c9634cae26fdf361ff5bff3c9975 Mon Sep 17 00:00:00 2001 From: David Huard Date: Fri, 7 Feb 2020 10:26:14 -0500 Subject: [PATCH 6/7] Revert "isort fixes" This reverts commit 5df6c9c0f99376dbc43f2f30567661ee49c00655. --- xarray/core/dataarray.py | 2 +- xarray/core/dataset.py | 4 ++-- xarray/tests/test_dask.py | 2 +- xarray/tests/test_interp.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 551baf97a8a..062cc6342df 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -22,6 +22,7 @@ import numpy as np import pandas as pd +from ..plot.plot import _PlotMethods from . import ( computation, dtypes, @@ -33,7 +34,6 @@ rolling, utils, ) -from ..plot.plot import _PlotMethods from .accessor_dt import CombinedDatetimelikeAccessor from .accessor_str import StringAccessor from .alignment import ( diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index fb5d44a892b..562f789d912 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -33,6 +33,8 @@ import xarray as xr +from ..coding.cftimeindex import _parse_array_of_cftime_strings +from ..plot.dataset_plot import _Dataset_PlotMethods from . import ( alignment, dtypes, @@ -45,8 +47,6 @@ rolling, utils, ) -from ..coding.cftimeindex import _parse_array_of_cftime_strings -from ..plot.dataset_plot import _Dataset_PlotMethods from .alignment import _broadcast_helper, _get_broadcast_dims_map_common_coords, align from .common import ( DataWithCoords, diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index 361537e3569..cc554850839 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -16,6 +16,7 @@ from xarray.testing import assert_chunks_equal from xarray.tests import mock +from ..core.duck_array_ops import lazy_array_equiv from . import ( assert_allclose, assert_array_equal, @@ -25,7 +26,6 @@ raises_regex, requires_scipy_or_netCDF4, ) -from ..core.duck_array_ops import lazy_array_equiv from .test_backends import create_tmp_file dask = pytest.importorskip("dask") diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index b1bb2760d1a..13e4eb93b59 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -5,8 +5,8 @@ import xarray as xr from xarray.tests import assert_allclose, assert_equal, requires_cftime, requires_scipy -from . import has_dask, has_scipy from ..coding.cftimeindex import _parse_array_of_cftime_strings +from . import has_dask, has_scipy from .test_dataset import create_test_data try: From 77824c20b9fac87c1cd6155818bce18a1c81402e Mon Sep 17 00:00:00 2001 From: David Huard Date: Fri, 7 Feb 2020 10:34:14 -0500 Subject: [PATCH 7/7] test requires scipy --- xarray/tests/test_interp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py index 13e4eb93b59..0502348160e 100644 --- a/xarray/tests/test_interp.py +++ b/xarray/tests/test_interp.py @@ -244,6 +244,7 @@ def test_interpolate_nd(case): assert_allclose(actual.transpose("y", "z"), expected) +@requires_scipy def test_interpolate_nd_nd(): """Interpolate nd array with an nd indexer sharing coordinates.""" # Create original array