Merge pull request #7883 from jcrist/matsym_iter

Support for single indexes in MatrixExpr

sympy/matrices/expressions/matexpr.py

@@ -217,6 +217,22 @@ def __getitem__(self, key):
                 return self._entry(i, j)
             else:
                 raise IndexError("Invalid indices (%s, %s)" % (i, j))
+        elif isinstance(key, (int, Integer)):
+            # row-wise decomposition of matrix
+            rows, cols = self.shape
+            if not (isinstance(rows, Integer) and isinstance(cols, Integer)):
+                raise IndexError("Single index only supported for "
+                                 "non-symbolic matrix shapes.")
+            key = sympify(key)
+            i = key // cols
+            j = key % cols
+            if self.valid_index(i, j) != False:
+                return self._entry(i, j)
+            else:
+                raise IndexError("Invalid index %s" % key)
+        elif isinstance(key, (Symbol, Expr)):
+                raise IndexError("Single index only supported for "
+                                 "non-symbolic indices.")
         raise IndexError("Invalid index, wanted %s[i,j]" % self)
 
     def as_explicit(self):

sympy/matrices/expressions/tests/test_matrix_exprs.py

@@ -193,5 +193,17 @@ def test_indexing():
     A[l, k]
     A[l+1, k+1]
 
+
+def test_single_indexing():
+    A = MatrixSymbol('A', 2, 3)
+    assert A[1] == A[0, 1]
+    assert A[3] == A[1, 0]
+    assert list(A[:2, :2]) == [A[0, 0], A[0, 1], A[1, 0], A[1, 1]]
+    raises(IndexError, lambda: A[6])
+    raises(IndexError, lambda: A[n])
+    B = MatrixSymbol('B', n, m)
+    raises(IndexError, lambda: B[1])
+
+
 def test_MatrixElement_diff():
     assert (A[3, 0]*A[0, 0]).diff(A[0, 0]) == A[3, 0]