def __getitem__(self, key): if not isinstance(key, tuple) and isinstance(key, slice): from diofant.matrices.expressions.slice import MatrixSlice return MatrixSlice(self, key, (0, None, 1)) if isinstance(key, tuple) and len(key) == 2: i, j = key if isinstance(i, slice) or isinstance(j, slice): from diofant.matrices.expressions.slice import MatrixSlice return MatrixSlice(self, i, j) i, j = sympify(i), sympify(j) if self.valid_index(i, j) is not False: 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) is not 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 test_entry(): B = MatrixSlice(X, (a, b), (c, d)) assert B[0, 0] == X[a, c] assert B[k, l] == X[a + k, c + l] pytest.raises(IndexError, lambda: MatrixSlice(X, 1, (2, 5))[1, 0]) assert X[1::2, :][1, 3] == X[1 + 2, 3] assert X[:, 1::2][3, 1] == X[3, 1 + 2]
def test_slicing(): assert X[1:5, 2:4] == MatrixSlice(X, (1, 5), (2, 4)) assert X[1, 2:4] == MatrixSlice(X, 1, (2, 4)) assert X[1:5, :].shape == (4, X.shape[1]) assert X[:, 1:5].shape == (X.shape[0], 4) assert X[::2, ::2].shape == (floor(n / 2), floor(m / 2)) assert X[2, :] == MatrixSlice(X, 2, (0, m)) assert X[k, :] == MatrixSlice(X, k, (0, m))
def test_indexing(): A = MatrixSymbol('A', n, m) A[1, 2] A[l, k] A[l + 1, k + 1] assert A[:] == MatrixSlice(A, (0, n, 1), (0, m, 1))
def blockcut(expr, rowsizes, colsizes): """ Cut a matrix expression into Blocks >>> from diofant import ImmutableMatrix, blockcut >>> M = ImmutableMatrix(4, 4, range(16)) >>> B = blockcut(M, (1, 3), (1, 3)) >>> type(B).__name__ 'BlockMatrix' >>> ImmutableMatrix(B.blocks[0, 1]) Matrix([[1, 2, 3]]) """ rowbounds = bounds(rowsizes) colbounds = bounds(colsizes) return BlockMatrix( [[MatrixSlice(expr, rowbound, colbound) for colbound in colbounds] for rowbound in rowbounds])
def test_inputs(): assert MatrixSlice(X, 1, (2, 5)) == MatrixSlice(X, (1, 2), (2, 5)) assert MatrixSlice(X, 1, (2, 5)).shape == (1, 3)
def test_on_diag(): assert not MatrixSlice(X, (a, b), (c, d)).on_diag assert MatrixSlice(X, (a, b), (a, b)).on_diag
def test_shape(): B = MatrixSlice(X, (a, b), (c, d)) assert B.shape == (b - a, d - c)