Add tensordot to dataarray class also add its test to test_dataarray

doc/api.rst

@@ -144,6 +144,7 @@ Computation
 :py:attr:`~Dataset.round`
 :py:attr:`~Dataset.real`
 :py:attr:`~Dataset.T`
+:py:attr:`~DataArray.dot`
 
 **Grouped operations**:
 :py:attr:`~core.groupby.DatasetGroupBy.assign`

doc/whats-new.rst

@@ -18,6 +18,8 @@ v0.7.2 (unreleased)
 
 Enhancements
 ~~~~~~~~~~~~
+-xarray version of np.dot :py:meth:`~DataArray.dot`. Performs dot product of 
+two DataArrays along their shared dims
 
 - Rolling window operations on DataArray objects are now supported via a new
   :py:meth:`xarray.DataArray.rolling` method.

xarray/core/dataarray.py

@@ -1374,5 +1374,64 @@ def real(self):
     def imag(self):
         return self._replace(self.variable.imag)
 
+    def dot(self, other):
+        """Perform dot product of two DataArrays along their shared dims.
+
+        Equivalent to taking taking tensordot over all shared dims.
+
+        Parameters
+        ----------
+        other : DataArray
+            The other array with which the dot product is performed.
+
+        Returns
+        -------
+        result : DataArray
+            Array resulting from the dot product over all shared dimensions.
+
+        See also
+        --------
+        np.tensordot(a, b, axes)
+
+        Examples
+        --------
+
+        >>> da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4))
+        >>> da = DataArray(da_vals, dims=['x', 'y', 'z'])    
+        >>> dm_vals = np.arange(4)
+        >>> dm = DataArray(dm_vals, dims=['z'])
+
+        >>> dm.dims
+        ('z')
+        >>> da.dims
+        ('x', 'y', 'z')
+
+        >>> dot_result = da.dot(dm)
+        >>> dot_result.dims
+        ('x', 'y')
+        """
+        if isinstance(other, Dataset):
+            raise NotImplementedError('dot products are not yet supported '
+                                      'with Dataset objects.')
+        if not isinstance(other, DataArray):
+            raise TypeError('dot only operates on DataArrays.')
+
+        # sum over the common dims
+        dims = set(self.dims) & set(other.dims)
+        if len(dims) == 0:
+            raise ValueError('DataArrays have no shared dimensions over which '
+                             'to perform dot.')
+
+        self, other = align(self, other, join='inner', copy=False)
+
+        axes = (self.get_axis_num(dims), other.get_axis_num(dims))
+        new_data = ops.tensordot(self.data, other.data, axes=axes)
+
+        new_coords = self.coords.merge(other.coords).drop(dims)
+        new_dims = ([d for d in self.dims if d not in dims] +
+                    [d for d in other.dims if d not in dims])
+
+        return type(self)(new_data, new_coords, new_dims)
+
 # priority most be higher than Variable to properly work with binary ufuncs
 ops.inject_all_ops_and_reduce_methods(DataArray, priority=60)

xarray/core/ops.py

@@ -97,6 +97,8 @@ def _fail_on_dask_array_input(values, msg=None, func_name=None):
 array_all = _dask_or_eager_func('all')
 array_any = _dask_or_eager_func('any')
 
+tensordot = _dask_or_eager_func('tensordot', n_array_args=2)
+
 
 def _interleaved_indices_required(indices):
     """With dask, we care about data locality and would rather avoid splitting

xarray/test/test_dask.py

@@ -300,3 +300,8 @@ def test_stack(self):
                              dims=['w', 'z'])
         assert stacked.data.chunks == expected.data.chunks
         self.assertLazyAndIdentical(expected, stacked)
+
+    def test_dot(self):
+        eager = self.eager_array.dot(self.eager_array[0])
+        lazy = self.lazy_array.dot(self.lazy_array[0])
+        self.assertLazyAndAllClose(eager, lazy)

xarray/test/test_dataarray.py

@@ -1792,3 +1792,41 @@ def test_full_like(self):
         actual = _full_like(DataArray([1, 2, 3]), fill_value=np.nan)
         self.assertEqual(actual.dtype, np.float)
         np.testing.assert_equal(actual.values, np.nan)
+
+    def test_dot(self):
+        x = np.linspace(-3, 3, 6)
+        y = np.linspace(-3, 3, 5)
+        z = range(4) 
+        da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4))
+        da = DataArray(da_vals, coords=[x, y, z], dims=['x', 'y', 'z'])
+
+        dm_vals = range(4)
+        dm = DataArray(dm_vals, coords=[z], dims=['z'])
+
+        # nd dot 1d
+        actual = da.dot(dm)
+        expected_vals = np.tensordot(da_vals, dm_vals, [2, 0])
+        expected = DataArray(expected_vals, coords=[x, y], dims=['x', 'y'])
+        self.assertDataArrayEqual(expected, actual)
+
+        # all shared dims
+        actual = da.dot(da)
+        expected_vals = np.tensordot(da_vals, da_vals, axes=([0, 1, 2], [0, 1, 2]))
+        expected = DataArray(expected_vals)
+        self.assertDataArrayEqual(expected, actual)
+
+        # multiple shared dims
+        dm_vals = np.arange(20 * 5 * 4).reshape((20, 5, 4))
+        j = np.linspace(-3, 3, 20)
+        dm = DataArray(dm_vals, coords=[j, y, z], dims=['j', 'y', 'z'])
+        actual = da.dot(dm)
+        expected_vals = np.tensordot(da_vals, dm_vals, axes=([1, 2], [1, 2]))
+        expected = DataArray(expected_vals, coords=[x, j], dims=['x', 'j'])
+        self.assertDataArrayEqual(expected, actual)
+
+        with self.assertRaises(NotImplementedError):
+            da.dot(dm.to_dataset(name='dm'))
+        with self.assertRaises(TypeError):
+            da.dot(dm.values)
+        with self.assertRaisesRegexp(ValueError, 'no shared dimensions'):
+            da.dot(DataArray(1))