Merge pull request #362 from bjodah/matmul-operator

Add support for Python's matrix multiplication operator

Author: isuruf

symengine/lib/symengine_wrapper.pyx

@@ -3230,6 +3230,8 @@ cdef class DenseMatrixBase(MatrixBase):
                 raise ValueError("sizes don't match.")
             else:
                 self.thisptr = new symengine.DenseMatrix(0, 0, v_)
+        elif col is not None and (row*col != v_.size()):
+            raise ValueError("Number of elements should equal rows*columns.")
         else:
             self.thisptr = new symengine.DenseMatrix(row, v_.size() / row, v_)
 
@@ -3271,6 +3273,13 @@ cdef class DenseMatrixBase(MatrixBase):
         else:
             return NotImplemented
 
+    def __matmul__(a, b):
+        a = _sympify(a, False)
+        b = _sympify(b, False)
+        if (a.ncols() != b.nrows()):
+            raise ShapeError("Invalid shapes for matrix multiplication. Got %s %s" % (a.shape, b.shape))
+        return a.mul_matrix(b)
+
     def __truediv__(a, b):
         return div_matrices(a, b)
 
@@ -4889,7 +4898,7 @@ cdef class LambdaDouble(_Lambdify):
         c_inp = np.ascontiguousarray(inp.ravel(order=self.order), dtype=self.numpy_dtype)
         c_out = out
         for idx in range(nbroadcast):
-            self.lambda_double[0].call(&c_out[idx*self.tot_out_size], &c_inp[idx*self.args_size]) 
+            self.lambda_double[0].call(&c_out[idx*self.tot_out_size], &c_inp[idx*self.args_size])
 
     cpdef as_scipy_low_level_callable(self):
         from ctypes import c_double, c_void_p, c_int, cast, POINTER, CFUNCTYPE
@@ -5135,7 +5144,7 @@ def Lambdify(args, *exprs, cppbool real=True, backend=None, order='C',
                     raise ValueError("Long double not supported on this platform")
             else:
                 raise ValueError("Unknown numpy dtype.")
-                
+
             if as_scipy:
                 return ret.as_scipy_low_level_callable()
             return ret

symengine/tests/test_matrices.py

@@ -12,6 +12,10 @@
     HAVE_NUMPY = False
 
 
+def test_init():
+    raises(ValueError, lambda: DenseMatrix(2, 1, [0]*4))
+
+
 def test_get():
     A = DenseMatrix([[1, 2], [3, 4]])
 
@@ -229,6 +233,8 @@ def test_mul_matrix():
 
     assert A.mul_matrix(B) == DenseMatrix(2, 2, [a + b, 0, c + d, 0])
     assert A * B == DenseMatrix(2, 2, [a + b, 0, c + d, 0])
+    assert A @ B == DenseMatrix(2, 2, [a + b, 0, c + d, 0])
+    assert (A @ DenseMatrix(2, 1, [0]*2)).shape == (2, 1)
 
     C = DenseMatrix(2, 3, [1, 2, 3, 2, 3, 4])
     D = DenseMatrix(3, 2, [3, 4, 4, 5, 5, 6])