Add matmul to DataFrame, Series

Author: bnaul

doc/source/whatsnew/v0.23.0.txt

@@ -344,6 +344,7 @@ Other Enhancements
 - :meth:`DataFrame.to_sql` now performs a multivalue insert if the underlying connection supports itk rather than inserting row by row.
   ``SQLAlchemy`` dialects supporting multivalue inserts include: ``mysql``, ``postgresql``, ``sqlite`` and any dialect with ``supports_multivalues_insert``. (:issue:`14315`, :issue:`8953`)
 - :func:`read_html` now accepts a ``displayed_only`` keyword argument to controls whether or not hidden elements are parsed (``True`` by default) (:issue:`20027`)
+- :class:`DataFrame` and :class:`Series` now support matrix multiplication (```@```) operator (:issue:`10259`) for Python>=3.5
 
 .. _whatsnew_0230.api_breaking:
 

pandas/core/frame.py

@@ -863,7 +863,8 @@ def __len__(self):
 
     def dot(self, other):
         """
-        Matrix multiplication with DataFrame or Series objects
+        Matrix multiplication with DataFrame or Series objects.  Can also be
+        called using `self @ other` in Python >= 3.5.
 
         Parameters
         ----------
@@ -905,6 +906,10 @@ def dot(self, other):
         else:  # pragma: no cover
             raise TypeError('unsupported type: %s' % type(other))
 
+    def __matmul__(self, other):
+        """ Matrix multiplication using binary `@` operator in Python>=3.5 """
+        return self.dot(other)
+
     # ----------------------------------------------------------------------
     # IO methods (to / from other formats)
 

pandas/core/series.py

@@ -1949,7 +1949,7 @@ def autocorr(self, lag=1):
     def dot(self, other):
         """
         Matrix multiplication with DataFrame or inner-product with Series
-        objects
+        objects. Can also be called using `self @ other` in Python >= 3.5.
 
         Parameters
         ----------
@@ -1988,6 +1988,10 @@ def dot(self, other):
         else:  # pragma: no cover
             raise TypeError('unsupported type: %s' % type(other))
 
+    def __matmul__(self, other):
+        """ Matrix multiplication using binary `@` operator in Python>=3.5 """
+        return self.dot(other)
+
     @Substitution(klass='Series')
     @Appender(base._shared_docs['searchsorted'])
     @deprecate_kwarg(old_arg_name='v', new_arg_name='value')

pandas/tests/frame/test_analytics.py

@@ -2089,41 +2089,42 @@ def test_clip_with_na_args(self):
                               self.frame)
 
     # Matrix-like
-
-    def test_dot(self):
+    @pytest.mark.parametrize('dot_fn', [DataFrame.dot, DataFrame.__matmul__])
+    def test_dot(self, dot_fn):
+        # __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)
@@ -2133,14 +2134,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

@@ -895,28 +895,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)