Add matmul to DataFrame, Series

Author: bnaul

doc/source/whatsnew/v0.23.0.txt

@@ -142,6 +142,7 @@ Other Enhancements
 - ``Categorical.rename_categories``, ``CategoricalIndex.rename_categories`` and :attr:`Series.cat.rename_categories`
   can now take a callable as their argument (:issue:`18862`)
 - :class:`Interval` and :class:`IntervalIndex` have gained a ``length`` attribute (:issue:`18789`)
+- :class:`DataFrame` and :class:`Series` now support matrix multiplication (```@```) operator (:issue:`10259`)
 
 .. _whatsnew_0230.api_breaking:
 

pandas/core/frame.py

@@ -866,6 +866,12 @@ def dot(self, other):
         else:  # pragma: no cover
             raise TypeError('unsupported type: %s' % type(other))
 
+    def __matmul__(self, other):
+        try:
+            return self.dot(other)
+        except TypeError:
+            return NotImplemented
+
     # ----------------------------------------------------------------------
     # IO methods (to / from other formats)
 

pandas/core/series.py

@@ -1625,6 +1625,12 @@ def dot(self, other):
         else:  # pragma: no cover
             raise TypeError('unsupported type: %s' % type(other))
 
+    def __matmul__(self, other):
+        try:
+            return self.dot(other)
+        except TypeError:
+            return NotImplemented
+
     @Substitution(klass='Series')
     @Appender(base._shared_docs['searchsorted'])
     @deprecate_kwarg(old_arg_name='v', new_arg_name='value')

pandas/tests/frame/test_analytics.py

@@ -2014,41 +2014,42 @@ def test_clip_with_na_args(self):
                               self.frame)
 
     # Matrix-like
-
+    @pytest.mark.parametrize('dot_fn', [DataFrame.dot, DataFrame.__matmul__])
     def test_dot(self):
+        # __matmul__ test is for GH #10259
         a = DataFrame(np.random.randn(3, 4), index=['a', 'b', 'c'],
                       columns=['p', 'q', 'r', 's'])
         b = DataFrame(np.random.randn(4, 2), index=['p', 'q', 'r', 's'],
                       columns=['one', 'two'])
 
-        result = a.dot(b)
+        result = dot_fn(a, b)
         expected = DataFrame(np.dot(a.values, b.values),
                              index=['a', 'b', 'c'],
                              columns=['one', 'two'])
         # Check alignment
         b1 = b.reindex(index=reversed(b.index))
-        result = a.dot(b)
+        result = dot_fn(a, b)
         tm.assert_frame_equal(result, expected)
 
         # Check series argument
-        result = a.dot(b['one'])
+        result = dot_fn(a, b['one'])
         tm.assert_series_equal(result, expected['one'], check_names=False)
         assert result.name is None
 
-        result = a.dot(b1['one'])
+        result = dot_fn(a, b1['one'])
         tm.assert_series_equal(result, expected['one'], check_names=False)
         assert result.name is None
 
         # can pass correct-length arrays
         row = a.iloc[0].values
 
-        result = a.dot(row)
-        exp = a.dot(a.iloc[0])
+        result = dot_fn(a, row)
+        exp = dot_fn(a, a.iloc[0])
         tm.assert_series_equal(result, exp)
 
         with tm.assert_raises_regex(ValueError,
                                     'Dot product shape mismatch'):
-            a.dot(row[:-1])
+            dot_fn(a, row[:-1])
 
         a = np.random.rand(1, 5)
         b = np.random.rand(5, 1)
@@ -2058,14 +2059,14 @@ def test_dot(self):
         B = DataFrame(b)  # noqa
 
         # it works
-        result = A.dot(b)
+        result = dot_fn(A, b)
 
         # unaligned
         df = DataFrame(randn(3, 4), index=[1, 2, 3], columns=lrange(4))
         df2 = DataFrame(randn(5, 3), index=lrange(5), columns=[1, 2, 3])
 
         with tm.assert_raises_regex(ValueError, 'aligned'):
-            df.dot(df2)
+            dot_fn(df, df2)
 
 
 @pytest.fixture

pandas/tests/series/test_analytics.py

@@ -811,28 +811,30 @@ def test_count(self):
         ts.iloc[[0, 3, 5]] = nan
         assert_series_equal(ts.count(level=1), right - 1)
 
-    def test_dot(self):
+    @pytest.mark.parametrize('dot_fn', [Series.dot, Series.__matmul__])
+    def test_dot(self, dot_fn):
+        # __matmul__ test is for GH #10259
         a = Series(np.random.randn(4), index=['p', 'q', 'r', 's'])
         b = DataFrame(np.random.randn(3, 4), index=['1', '2', '3'],
                       columns=['p', 'q', 'r', 's']).T
 
-        result = a.dot(b)
+        result = dot_fn(a, b)
         expected = Series(np.dot(a.values, b.values), index=['1', '2', '3'])
         assert_series_equal(result, expected)
 
         # Check index alignment
         b2 = b.reindex(index=reversed(b.index))
-        result = a.dot(b)
+        result = dot_fn(a, b)
         assert_series_equal(result, expected)
 
         # Check ndarray argument
-        result = a.dot(b.values)
+        result = dot_fn(a, b.values)
         assert np.all(result == expected.values)
-        assert_almost_equal(a.dot(b['2'].values), expected['2'])
+        assert_almost_equal(dot_fn(a, b['2'].values), expected['2'])
 
         # Check series argument
-        assert_almost_equal(a.dot(b['1']), expected['1'])
-        assert_almost_equal(a.dot(b2['1']), expected['1'])
+        assert_almost_equal(dot_fn(a, b['1']), expected['1'])
+        assert_almost_equal(dot_fn(a, b2['1']), expected['1'])
 
         pytest.raises(Exception, a.dot, a.values[:3])
         pytest.raises(ValueError, a.dot, b.T)