implement interp_like (#2222)

* implement interp_like

* flake8

* interp along datetime

* Support datetime coordinate

* Using reindex for object coordinate.

* Update based on the comments

Author: fujiisoup

doc/api.rst

@@ -111,6 +111,7 @@ Indexing
    Dataset.sel
    Dataset.squeeze
    Dataset.interp
+   Dataset.interp_like
    Dataset.reindex
    Dataset.reindex_like
    Dataset.set_index
@@ -265,6 +266,7 @@ Indexing
    DataArray.sel
    DataArray.squeeze
    DataArray.interp
+   DataArray.interp_like
    DataArray.reindex
    DataArray.reindex_like
    DataArray.set_index

doc/indexing.rst

@@ -193,6 +193,14 @@ Indexing axes with monotonic decreasing labels also works, as long as the
    reversed_da.loc[3.1:0.9]
 
 
+.. note::
+
+  If you want to interpolate along coordinates rather than looking up the
+  nearest neighbors, use :py:meth:`~xarray.Dataset.interp` and
+  :py:meth:`~xarray.Dataset.interp_like`.
+  See :ref:`interpolation <interp>` for the details.
+
+
 Dataset indexing
 ----------------
 

doc/interpolation.rst

@@ -34,7 +34,7 @@ indexing of a :py:class:`~xarray.DataArray`,
     da.sel(time=3)
 
     # interpolation
-    da.interp(time=3.5)
+    da.interp(time=2.5)
 
 
 Similar to the indexing, :py:meth:`~xarray.DataArray.interp` also accepts an
@@ -82,6 +82,31 @@ Array-like coordinates are also accepted:
     da.interp(time=[1.5, 2.5], space=[0.15, 0.25])
 
 
+:py:meth:`~xarray.DataArray.interp_like` method is a useful shortcut. This
+method interpolates an xarray object onto the coordinates of another xarray
+object. For example, if we want to compute the difference between
+two :py:class:`~xarray.DataArray` s (``da`` and ``other``) staying on slightly
+different coordinates,
+
+.. ipython:: python
+
+  other = xr.DataArray(np.sin(0.4 * np.arange(9).reshape(3, 3)),
+                       [('time', [0.9, 1.9, 2.9]),
+                       ('space', [0.15, 0.25, 0.35])])
+
+it might be a good idea to first interpolate ``da`` so that it will stay on the
+same coordinates of ``other``, and then subtract it.
+:py:meth:`~xarray.DataArray.interp_like` can be used for such a case,
+
+.. ipython:: python
+
+  # interpolate da along other's coordinates
+  interpolated = da.interp_like(other)
+  interpolated
+
+It is now possible to safely compute the difference ``other - interpolated``.
+
+
 Interpolation methods
 ---------------------
 

doc/whats-new.rst

@@ -36,6 +36,12 @@ Documentation
 Enhancements
 ~~~~~~~~~~~~
 
+- :py:meth:`~xarray.DataArray.interp_like` and
+  :py:meth:`~xarray.Dataset.interp_like` methods are newly added.
+  (:issue:`2218`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
+
+
 Bug fixes
 ~~~~~~~~~
 

xarray/core/dataarray.py

@@ -951,6 +951,54 @@ def interp(self, coords=None, method='linear', assume_sorted=False,
             **coords_kwargs)
         return self._from_temp_dataset(ds)
 
+    def interp_like(self, other, method='linear', assume_sorted=False,
+                    kwargs={}):
+        """Interpolate this object onto the coordinates of another object,
+        filling out of range values with NaN.
+
+        Parameters
+        ----------
+        other : Dataset or DataArray
+            Object with an 'indexes' attribute giving a mapping from dimension
+            names to an 1d array-like, which provides coordinates upon
+            which to index the variables in this dataset.
+        method: string, optional.
+            {'linear', 'nearest'} for multidimensional array,
+            {'linear', 'nearest', 'zero', 'slinear', 'quadratic', 'cubic'}
+            for 1-dimensional array. 'linear' is used by default.
+        assume_sorted: boolean, optional
+            If False, values of coordinates that are interpolated over can be
+            in any order and they are sorted first. If True, interpolated
+            coordinates are assumed to be an array of monotonically increasing
+            values.
+        kwargs: dictionary, optional
+            Additional keyword passed to scipy's interpolator.
+
+        Returns
+        -------
+        interpolated: xr.DataArray
+            Another dataarray by interpolating this dataarray's data along the
+            coordinates of the other object.
+
+        Note
+        ----
+        scipy is required.
+        If the dataarray has object-type coordinates, reindex is used for these
+        coordinates instead of the interpolation.
+
+        See Also
+        --------
+        DataArray.interp
+        DataArray.reindex_like
+        """
+        if self.dtype.kind not in 'uifc':
+            raise TypeError('interp only works for a numeric type array. '
+                            'Given {}.'.format(self.dtype))
+
+        ds = self._to_temp_dataset().interp_like(
+            other, method=method, kwargs=kwargs, assume_sorted=assume_sorted)
+        return self._from_temp_dataset(ds)
+
     def rename(self, new_name_or_name_dict=None, **names):
         """Returns a new DataArray with renamed coordinates or a new name.
 

xarray/core/dataset.py

@@ -1893,6 +1893,63 @@ def maybe_variable(obj, k):
                        .union(coord_vars))
         return obj._replace_vars_and_dims(variables, coord_names=coord_names)
 
+    def interp_like(self, other, method='linear', assume_sorted=False,
+                    kwargs={}):
+        """Interpolate this object onto the coordinates of another object,
+        filling the out of range values with NaN.
+
+        Parameters
+        ----------
+        other : Dataset or DataArray
+            Object with an 'indexes' attribute giving a mapping from dimension
+            names to an 1d array-like, which provides coordinates upon
+            which to index the variables in this dataset.
+        method: string, optional.
+            {'linear', 'nearest'} for multidimensional array,
+            {'linear', 'nearest', 'zero', 'slinear', 'quadratic', 'cubic'}
+            for 1-dimensional array. 'linear' is used by default.
+        assume_sorted: boolean, optional
+            If False, values of coordinates that are interpolated over can be
+            in any order and they are sorted first. If True, interpolated
+            coordinates are assumed to be an array of monotonically increasing
+            values.
+        kwargs: dictionary, optional
+            Additional keyword passed to scipy's interpolator.
+
+        Returns
+        -------
+        interpolated: xr.Dataset
+            Another dataset by interpolating this dataset's data along the
+            coordinates of the other object.
+
+        Note
+        ----
+        scipy is required.
+        If the dataset has object-type coordinates, reindex is used for these
+        coordinates instead of the interpolation.
+
+        See Also
+        --------
+        Dataset.interp
+        Dataset.reindex_like
+        """
+        coords = alignment.reindex_like_indexers(self, other)
+
+        numeric_coords = OrderedDict()
+        object_coords = OrderedDict()
+        for k, v in coords.items():
+            if v.dtype.kind in 'uifcMm':
+                numeric_coords[k] = v
+            else:
+                object_coords[k] = v
+
+        ds = self
+        if object_coords:
+            # We do not support interpolation along object coordinate.
+            # reindex instead.
+            ds = self.reindex(object_coords)
+        return ds.interp(numeric_coords, method, assume_sorted, kwargs)
+
     def rename(self, name_dict=None, inplace=False, **names):
         """Returns a new object with renamed variables and dimensions.
 

xarray/core/missing.py

@@ -395,6 +395,26 @@ def _localize(var, indexes_coords):
     return var.isel(**indexes), indexes_coords
 
 
+def _floatize_x(x, new_x):
+    """ Make x and new_x float.
+    This is particulary useful for datetime dtype.
+    x, new_x: tuple of np.ndarray
+    """
+    x = list(x)
+    new_x = list(new_x)
+    for i in range(len(x)):
+        if x[i].dtype.kind in 'Mm':
+            # Scipy casts coordinates to np.float64, which is not accurate
+            # enough for datetime64 (uses 64bit integer).
+            # We assume that the most of the bits are used to represent the
+            # offset (min(x)) and the variation (x - min(x)) can be
+            # represented by float.
+            xmin = np.min(x[i])
+            x[i] = (x[i] - xmin).astype(np.float64)
+            new_x[i] = (new_x[i] - xmin).astype(np.float64)
+    return x, new_x
+
+
 def interp(var, indexes_coords, method, **kwargs):
     """ Make an interpolation of Variable
 
@@ -523,6 +543,8 @@ def _interp1d(var, x, new_x, func, kwargs):
 
 
 def _interpnd(var, x, new_x, func, kwargs):
+    x, new_x = _floatize_x(x, new_x)
+
     if len(x) == 1:
         return _interp1d(var, x, new_x, func, kwargs)
 

xarray/tests/test_interp.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import, division, print_function
 
 import numpy as np
+import pandas as pd
 import pytest
 
 import xarray as xr
@@ -430,3 +431,50 @@ def test_interpolate_dimorder(case):
         actual = da.interp(x=0.5, z=new_z).dims
         expected = da.sel(x=0.5, z=new_z, method='nearest').dims
         assert actual == expected
+
+
+@requires_scipy
+def test_interp_like():
+    ds = create_test_data()
+    ds.attrs['foo'] = 'var'
+    ds['var1'].attrs['buz'] = 'var2'
+
+    other = xr.DataArray(np.random.randn(3), dims=['dim2'],
+                         coords={'dim2': [0, 1, 2]})
+    interpolated = ds.interp_like(other)
+
+    assert_allclose(interpolated['var1'],
+                    ds['var1'].interp(dim2=other['dim2']))
+    assert_allclose(interpolated['var1'],
+                    ds['var1'].interp_like(other))
+    assert interpolated['var3'].equals(ds['var3'])
+
+    # attrs should be kept
+    assert interpolated.attrs['foo'] == 'var'
+    assert interpolated['var1'].attrs['buz'] == 'var2'
+
+    other = xr.DataArray(np.random.randn(3), dims=['dim3'],
+                         coords={'dim3': ['a', 'b', 'c']})
+
+    actual = ds.interp_like(other)
+    expected = ds.reindex_like(other)
+    assert_allclose(actual, expected)
+
+
+@requires_scipy
+def test_datetime():
+    da = xr.DataArray(np.random.randn(24), dims='time',
+                      coords={'time': pd.date_range('2000-01-01', periods=24)})
+
+    x_new = pd.date_range('2000-01-02', periods=3)
+    actual = da.interp(time=x_new)
+    expected = da.isel(time=[1, 2, 3])
+    assert_allclose(actual, expected)
+
+    x_new = np.array([np.datetime64('2000-01-01T12:00'),
+                      np.datetime64('2000-01-02T12:00')])
+    actual = da.interp(time=x_new)
+    assert_allclose(actual.isel(time=0).drop('time'),
+                    0.5 * (da.isel(time=0) + da.isel(time=1)))
+    assert_allclose(actual.isel(time=1).drop('time'),
+                    0.5 * (da.isel(time=1) + da.isel(time=2)))