Add rmatmul to DataFrame, Series

Author: bnaul

pandas/core/frame.py

@@ -910,6 +910,10 @@ def __matmul__(self, other):
         """ Matrix multiplication using binary `@` operator in Python>=3.5 """
         return self.dot(other)
 
+    def __rmatmul__(self, other):
+        """ Matrix multiplication using binary `@` operator in Python>=3.5 """
+        return self.T.dot(np.transpose(other)).T
+
     # ----------------------------------------------------------------------
     # IO methods (to / from other formats)
 

pandas/core/series.py

@@ -2035,6 +2035,10 @@ def __matmul__(self, other):
         """ Matrix multiplication using binary `@` operator in Python>=3.5 """
         return self.dot(other)
 
+    def __rmatmul__(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

@@ -5,6 +5,7 @@
 import warnings
 from datetime import timedelta
 from distutils.version import LooseVersion
+import operator
 import sys
 import pytest
 
@@ -2091,42 +2092,40 @@ def test_clip_with_na_args(self):
                               self.frame)
 
     # Matrix-like
-    @pytest.mark.parametrize('dot_fn', [DataFrame.dot, DataFrame.__matmul__])
-    def test_dot(self, dot_fn):
-        # __matmul__ test is for GH #10259
+    def test_dot(self):
         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 = dot_fn(a, b)
+        result = a.dot(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 = dot_fn(a, b)
+        result = a.dot(b)
         tm.assert_frame_equal(result, expected)
 
         # Check series argument
-        result = dot_fn(a, b['one'])
+        result = a.dot(b['one'])
         tm.assert_series_equal(result, expected['one'], check_names=False)
         assert result.name is None
 
-        result = dot_fn(a, b1['one'])
+        result = a.dot(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 = dot_fn(a, row)
-        exp = dot_fn(a, a.iloc[0])
+        result = a.dot(row)
+        exp = a.dot(a.iloc[0])
         tm.assert_series_equal(result, exp)
 
         with tm.assert_raises_regex(ValueError,
                                     'Dot product shape mismatch'):
-            dot_fn(a, row[:-1])
+            a.dot(row[:-1])
 
         a = np.random.rand(1, 5)
         b = np.random.rand(5, 1)
@@ -2136,14 +2135,55 @@ def test_dot(self, dot_fn):
         B = DataFrame(b)  # noqa
 
         # it works
-        result = dot_fn(A, b)
+        result = A.dot(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'):
-            dot_fn(df, df2)
+            df.dot(df2)
+
+    @pytest.mark.skipif(sys.version_info < (3, 5),
+                        reason='matmul supported for Python>=3.5')
+    def test_matmul(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'])
+
+        # DataFrame @ DataFrame
+        result = operator.matmul(a, b)
+        expected = DataFrame(np.dot(a.values, b.values),
+                             index=['a', 'b', 'c'],
+                             columns=['one', 'two'])
+        tm.assert_frame_equal(result, expected)
+
+        # DataFrame @ Series
+        result = operator.matmul(a, b.one)
+        expected = Series(np.dot(a.values, b.one.values),
+                          index=['a', 'b', 'c'])
+        tm.assert_series_equal(result, expected)
+
+        # np.array @ DataFrame
+        result = operator.matmul(a.values, b)
+        expected = np.dot(a.values, b.values)
+        tm.assert_almost_equal(result, expected)
+
+        # nested list @ DataFrame (__rmatmul__)
+        result = operator.matmul(a.values.tolist(), b)
+        expected = DataFrame(np.dot(a.values, b.values),
+                             index=['a', 'b', 'c'],
+                             columns=['one', 'two'])
+        tm.assert_almost_equal(result.values, expected.values)
+
+        # 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'):
+            operator.matmul(df, df2)
 
 
 @pytest.fixture

pandas/tests/series/test_analytics.py

@@ -3,6 +3,8 @@
 
 from itertools import product
 from distutils.version import LooseVersion
+import operator
+import sys
 
 import pytest
 
@@ -895,30 +897,60 @@ def test_count(self):
         ts.iloc[[0, 3, 5]] = nan
         assert_series_equal(ts.count(level=1), right - 1)
 
-    @pytest.mark.parametrize('dot_fn', [Series.dot, Series.__matmul__])
-    def test_dot(self, dot_fn):
-        # __matmul__ test is for GH #10259
+    def test_dot(self):
         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 = dot_fn(a, b)
+        result = a.dot(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 = dot_fn(a, b)
+        result = a.dot(b)
         assert_series_equal(result, expected)
 
         # Check ndarray argument
-        result = dot_fn(a, b.values)
+        result = a.dot(b.values)
         assert np.all(result == expected.values)
-        assert_almost_equal(dot_fn(a, b['2'].values), expected['2'])
+        assert_almost_equal(a.dot(b['2'].values), expected['2'])
 
         # Check series argument
-        assert_almost_equal(dot_fn(a, b['1']), expected['1'])
-        assert_almost_equal(dot_fn(a, b2['1']), expected['1'])
+        assert_almost_equal(a.dot(b['1']), expected['1'])
+        assert_almost_equal(a.dot(b2['1']), expected['1'])
+
+        pytest.raises(Exception, a.dot, a.values[:3])
+        pytest.raises(ValueError, a.dot, b.T)
+
+    @pytest.mark.skipif(sys.version_info < (3, 5),
+                        reason='matmul supported for Python>=3.5')
+    def test_matmul(self):
+        # 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
+
+        # Series @ DataFrame
+        result = operator.matmul(a, b)
+        expected = Series(np.dot(a.values, b.values), index=['1', '2', '3'])
+        assert_series_equal(result, expected)
+
+        # DataFrame @ Series
+        result = operator.matmul(b.T, a)
+        expected = Series(np.dot(b.T.values, a.T.values),
+                          index=['1', '2', '3'])
+        assert_series_equal(result, expected)
+
+        # Series @ Series
+        result = operator.matmul(a, a)
+        expected = np.dot(a.values, a.values)
+        assert_almost_equal(result, expected)
+
+        # np.array @ Series  (__rmatmul__)
+        result = operator.matmul(a.values, a)
+        expected = np.dot(a.values, a.values)
+        assert_almost_equal(result, expected)
 
         pytest.raises(Exception, a.dot, a.values[:3])
         pytest.raises(ValueError, a.dot, b.T)