Skip to content

Improve interp performance #4069

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.

Sign up for GitHub

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

Merged
merged 6 commits into from
May 25, 2020
Merged

Improve interp performance #4069

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b81a624
Fixes 2223
fujiisoup May 16, 2020
7c1919f
more tests
fujiisoup May 17, 2020
4a4e295
add @requires_scipy to test
fujiisoup May 17, 2020
238a08d
fix tests
fujiisoup May 17, 2020
200dacc
black
fujiisoup May 17, 2020
1a7d738
update whatsnew. Added a test for nearest
fujiisoup May 24, 2020
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
7 changes: 7 additions & 0 deletions doc/whats-new.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,13 @@ Breaking changes
(:pull:`3274`)
By `Elliott Sales de Andrade <https://github.com/QuLogic>`_

Enhancements
~~~~~~~~~~~~
- Performance improvement of :py:meth:`DataArray.interp` and :py:func:`Dataset.interp`
For orthogonal linear- and nearest-neighbor interpolation, we do 1d-interpolation sequentially
rather than interpolating in multidimensional space. (:issue:`2223`)
By `Keisuke Fujii <https://github.com/fujiisoup>`_.

New Features
~~~~~~~~~~~~
- Added :py:meth:`DataArray.polyfit` and :py:func:`xarray.polyval` for fitting polynomials. (:issue:`3349`)
Expand Down
15 changes: 14 additions & 1 deletion xarray/core/missing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -619,6 +619,19 @@ def interp(var, indexes_coords, method, **kwargs):
# default behavior
kwargs["bounds_error"] = kwargs.get("bounds_error", False)

# check if the interpolation can be done in orthogonal manner
if (
len(indexes_coords) > 1
and method in ["linear", "nearest"]
and all(dest[1].ndim == 1 for dest in indexes_coords.values())
and len(set([d[1].dims[0] for d in indexes_coords.values()]))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This condition is confusing me. This will not speed up this case:

da = xr.DataArray(
        np.arange(6).reshape(3, 2),
        dims=["x", "y"],
        coords={"x": [0, 1, 2], "y": [-0.1, -0.3]},
)

x_new = xr.DataArray([0.5, 1.5, 2.5], dims=["x1"]) 
y_new = xr.DataArray([-0.15, -0.25, -0.35], dims=["x1"])
da.interp(x=x_new, y=y_new)

Correct? Is that intentional?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @dcherian
It was intentional, but you are right. This case can be also improved.

da.interp(x=x_new, y=y_new)

should be equivalent to

da.interp(x=x_new).interp(y=y_new)

But then I'm a bit confused about which case should be improved.
For example if len(x_new) = len(y_new) = 1000000, then the original interpretation may be faster, although this is a rare use case.
Maybe we can use some heuristics, such as

len(x_new) < len(x)

?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm now thinking that the simpler behavior is better; for an orthogonal interpolation we interpolate sequentially and otherwise we use interpn.
Further improvement may be done in upstream.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

da = xr.DataArray(
        np.arange(6).reshape(3, 2),
        dims=["x", "y"],
        coords={"x": [0, 1, 2], "y": [-0.1, -0.3]},
)

x_new = xr.DataArray([0.5, 1.5, 2.5], dims=["x1"]) 
y_new = xr.DataArray([-0.15, -0.25, -0.35], dims=["x1"])
da.interp(x=x_new, y=y_new)

It looks that this case is not slow even with our current code.
The problem is when the final destination is a regular grid, where interpn will compute many times.
So, probably this PR should work good enough for this case.

== len(indexes_coords)
):
# interpolate sequentially
for dim, dest in indexes_coords.items():
var = interp(var, {dim: dest}, method, **kwargs)
return var

# target dimensions
dims = list(indexes_coords)
x, new_x = zip(*[indexes_coords[d] for d in dims])
Expand Down Expand Up @@ -659,7 +672,7 @@ def interp_func(var, x, new_x, method, kwargs):
New coordinates. Should not contain NaN.
method: string
{'linear', 'nearest', 'zero', 'slinear', 'quadratic', 'cubic'} for
1-dimensional itnterpolation.
1-dimensional interpolation.
{'linear', 'nearest'} for multidimensional interpolation
**kwargs:
Optional keyword arguments to be passed to scipy.interpolator
Expand Down
7 changes: 1 addition & 6 deletions xarray/testing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,7 @@
from xarray.core.indexes import default_indexes
from xarray.core.variable import IndexVariable, Variable

__all__ = (
"assert_allclose",
"assert_chunks_equal",
"assert_equal",
"assert_identical",
)
__all__ = ("assert_allclose", "assert_chunks_equal", "assert_equal", "assert_identical")


def _decode_string_data(data):
Expand Down
18 changes: 18 additions & 0 deletions xarray/tests/test_interp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -699,3 +699,21 @@ def test_3641():
times = xr.cftime_range("0001", periods=3, freq="500Y")
da = xr.DataArray(range(3), dims=["time"], coords=[times])
da.interp(time=["0002-05-01"])


@requires_scipy
@pytest.mark.parametrize("method", ["nearest", "linear"])
def test_decompose(method):
da = xr.DataArray(
np.arange(6).reshape(3, 2),
dims=["x", "y"],
coords={"x": [0, 1, 2], "y": [-0.1, -0.3]},
)
x_new = xr.DataArray([0.5, 1.5, 2.5], dims=["x1"])
y_new = xr.DataArray([-0.15, -0.25], dims=["y1"])
x_broadcast, y_broadcast = xr.broadcast(x_new, y_new)
assert x_broadcast.ndim == 2

actual = da.interp(x=x_new, y=y_new, method=method).drop(("x", "y"))
expected = da.interp(x=x_broadcast, y=y_broadcast, method=method).drop(("x", "y"))
assert_allclose(actual, expected)