Fix Lambdify for heterogeneous input

Author: bjodah

symengine/lib/symengine_wrapper.pyx

@@ -2117,6 +2117,9 @@ cdef class DenseMatrixBase(MatrixBase):
         def __get__(self):
             return self.nrows()*self.ncols()
 
+    def ravel(self):
+        return [self.get(i, j) for i in range(self.nrows()) for j in range(self.ncols())]
+
     def reshape(self, rows, cols):
         if len(self) != rows*cols:
             raise ValueError("Invalid reshape parameters %d %d" % (rows, cols))
@@ -2132,9 +2135,9 @@ cdef class DenseMatrixBase(MatrixBase):
         if j < 0:
             j += nc
         if i < 0 or i >= nr:
-            raise IndexError
+            raise IndexError("Row index out of bounds: %d" % i)
         if j < 0 or j >= nc:
-            raise IndexError
+            raise IndexError("Column index out of bounds: %d" % j)
         return i, j
 
     def get(self, i, j):
@@ -3140,19 +3143,20 @@ cdef class _Lambdify(object):
                 e_ = _sympify(e)
                 args_.push_back(e_.thisptr)
 
-        if isinstance(exprs, DenseMatrixBase):
-            nr = exprs.nrows()
-            nc = exprs.ncols()
-            mtx = (<DenseMatrixBase>exprs).thisptr
-            for ri in range(nr):
-                for ci in range(nc):
-                    b_ = deref(mtx).get(ri, ci)
-                    outs_.push_back(b_)
-        else:
-            for e in ravel(exprs):
-                e_ = _sympify(e)
-                outs_.push_back(e_.thisptr)
 
+        for curr_expr in exprs:
+            if isinstance(curr_expr, DenseMatrixBase):
+                nr = curr_expr.nrows()
+                nc = curr_expr.ncols()
+                mtx = (<DenseMatrixBase>curr_expr).thisptr
+                for ri in range(nr):
+                    for ci in range(nc):
+                        b_ = deref(mtx).get(ri, ci)
+                        outs_.push_back(b_)
+            else:
+                for e in ravel(curr_expr):
+                    e_ = _sympify(e)
+                    outs_.push_back(e_.thisptr)
         self._init(args_, outs_)
 
     cdef _init(self, symengine.vec_basic& args_, symengine.vec_basic& outs_):
@@ -3391,7 +3395,8 @@ def LambdifyCSE(args, *exprs, real=True, cse=None, concatenate=None):
     if concatenate is None:
         from numpy import concatenate
     from sympy import sympify as s_sympify
-    subs, new_exprs = cse([s_sympify(expr) for expr in exprs])
+    flat_exprs = list(itertools.chain(*map(ravel, exprs)))
+    subs, flat_new_exprs = cse([s_sympify(expr) for expr in flat_exprs])
     if subs:
         cse_symbs, cse_exprs = zip(*subs)
         new_exprs = []

symengine/tests/test_lambdify.py

@@ -50,6 +50,19 @@ def test_get_shape():
     assert get_shape([[1], [1], [1]]) == (3, 1)
     assert get_shape([[1, 1, 1]]) == (1, 3)
 
+    x = se.symbols('x')
+    exprs = [x+1, x+2, x+3, 1/x, 1/(x*x), 1/(x**3.0)]
+    A = se.DenseMatrix(2, 3, exprs)
+    assert get_shape(A) == (2, 3)
+
+
+def test_ravel():
+    x = se.symbols('x')
+    ravel = se.lib.symengine_wrapper.ravel
+    exprs = [x+1, x+2, x+3, 1/x, 1/(x*x), 1/(x**3.0)]
+    A = se.DenseMatrix(2, 3, exprs)
+    assert ravel(A) == exprs
+
 
 def test_Lambdify():
     n = 7