def _print_Assignment(self, expr): from sympy.functions.elementary.piecewise import Piecewise from sympy.tensor.indexed import IndexedBase # Copied from codeprinter, but remove special MatrixSymbol treatment lhs = expr.lhs rhs = expr.rhs # We special case assignments that take multiple lines if not self._settings["inline"] and isinstance(expr.rhs, Piecewise): # Here we modify Piecewise so each expression is now # an Assignment, and then continue on the print. expressions = [] conditions = [] for (e, c) in rhs.args: expressions.append(Assignment(lhs, e)) conditions.append(c) temp = Piecewise(*zip(expressions, conditions)) return self._print(temp) if self._settings["contract"] and (lhs.has(IndexedBase) or rhs.has(IndexedBase)): # Here we check if there is looping to be done, and if so # print the required loops. return self._doprint_loops(rhs, lhs) else: lhs_code = self._print(lhs) rhs_code = self._print(rhs) return self._get_statement("%s = %s" % (lhs_code, rhs_code))
def _print_Assignment(self, expr): from sympy.functions.elementary.piecewise import Piecewise from sympy.matrices.expressions.matexpr import MatrixSymbol from sympy.tensor.indexed import IndexedBase lhs = expr.lhs rhs = expr.rhs # We special case assignments that take multiple lines if isinstance(expr.rhs, Piecewise): # Here we modify Piecewise so each expression is now # an Assignment, and then continue on the print. expressions = [] conditions = [] for (e, c) in rhs.args: expressions.append(Assignment(lhs, e)) conditions.append(c) temp = Piecewise(*zip(expressions, conditions)) return self._print(temp) elif isinstance(lhs, MatrixSymbol): # Here we form an Assignment for each element in the array, # printing each one. lines = [] for (i, j) in self._traverse_matrix_indices(lhs): temp = Assignment(lhs[i, j], rhs[i, j]) if not (rhs[i, j] == 0): code0 = self._print(temp) lines.append(code0) return "\n".join(lines) elif self._settings["contract"] and (lhs.has(IndexedBase) or rhs.has(IndexedBase)): # Here we check if there is looping to be done, and if so # print the required loops. return self._doprint_loops(rhs, lhs) else: lhs_code = self._print(lhs) rhs_code = self._print(rhs) return self._get_statement("%s = %s" % (lhs_code, rhs_code))
def test_Assignment(): x, y = symbols("x, y") A = MatrixSymbol('A', 3, 1) mat = Matrix([1, 2, 3]) B = IndexedBase('B') n = symbols("n", integer=True) i = Idx("i", n) # Here we just do things to show they don't error Assignment(x, y) Assignment(x, 0) Assignment(A, mat) Assignment(A[1, 0], 0) Assignment(A[1, 0], x) Assignment(B[i], x) Assignment(B[i], 0) # Here we test things to show that they error # Matrix to scalar raises(ValueError, lambda: Assignment(B[i], A)) raises(ValueError, lambda: Assignment(B[i], mat)) raises(ValueError, lambda: Assignment(x, mat)) raises(ValueError, lambda: Assignment(x, A)) raises(ValueError, lambda: Assignment(A[1, 0], mat)) # Scalar to matrix raises(ValueError, lambda: Assignment(A, x)) raises(ValueError, lambda: Assignment(A, 0)) # Non-atomic lhs raises(TypeError, lambda: Assignment(mat, A)) raises(TypeError, lambda: Assignment(0, x)) raises(TypeError, lambda: Assignment(x * x, 1)) raises(TypeError, lambda: Assignment(A + A, mat)) raises(TypeError, lambda: Assignment(B, 0))