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) a = Assignment(x, y) assert a.func(*a.args) == a # 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)) assert Relational(x, y, ':=') == Assignment(x, y)
def test_Assignment(): # 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) a = Assignment(x, y) assert a.func(*a.args) == a assert a.op == ':=' # 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))
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]) code0 = self._print(temp) lines.append(code0) return "\n".join(lines) elif self._settings.get("contract", False) 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_mul_binop(): src1 = ( src + """\ d = a + b - c c = a * b + d s = p * q / r r = p * s + q / p """ ) expr1.convert_to_expr(src1, "f") ls1 = expr1.return_expr() for iter in range(8, 12): assert isinstance(ls1[iter], Assignment) assert ls1[8] == Assignment( Variable(Symbol("d")), Symbol("a") + Symbol("b") - Symbol("c") ) assert ls1[9] == Assignment( Variable(Symbol("c")), Symbol("a") * Symbol("b") + Symbol("d") ) assert ls1[10] == Assignment( Variable(Symbol("s")), Symbol("p") * Symbol("q") / Symbol("r") ) assert ls1[11] == Assignment( Variable(Symbol("r")), Symbol("p") * Symbol("s") + Symbol("q") / Symbol("p") )
def visit_Assignment(self, node): """Visitor Function for Assignment Visits each Assignment is the LFortran ASR and creates corresponding assignment for SymPy. Notes ===== The function currently only supports variable assignment and binary operation assignments of varying multitudes. Any type of numberS or array is not supported. Raises ====== NotImplementedError() when called for Numeric assignments or Arrays """ # TODO: Arithmatic Assignment if isinstance(node.target, asr.Variable): target = node.target value = node.value if isinstance(value, asr.Variable): new_node = Assignment(Variable(target.name), Variable(value.name)) elif type(value) == asr.BinOp: exp_ast = call_visitor(value) for expr in exp_ast: new_node = Assignment(Variable(target.name), expr) else: raise NotImplementedError("Numeric assignments not supported") else: raise NotImplementedError("Arrays not supported") self._py_ast.append(new_node)
def test_CodeBlock_cse__issue_14118(): # see https://github.com/sympy/sympy/issues/14118 c = CodeBlock( Assignment(A22, Matrix([[x, sin(y)], [3, 4]])), Assignment(B22, Matrix([[sin(y), 2 * sin(y)], [sin(y)**2, 7]]))) assert c.cse() == CodeBlock( Assignment(x0, sin(y)), Assignment(A22, Matrix([[x, x0], [3, 4]])), Assignment(B22, Matrix([[x0, 2 * x0], [x0**2, 7]])))
def cdb_write_code(ex, name, filename, num): import os import sympy as sym # -- 24 sep 2021 ----------------------------------------- old # from sympy.printing.ccode import C99CodePrinter as printer # from sympy.printing.codeprinter import Assignment # -- 24 sep 2021 ----------------------------------------- new # changes due to upgrade of SymPy to v 1.7 from sympy.printing.c import C99CodePrinter as printer from sympy.codegen.ast import Assignment # -------------------------------------------------------- 24 sep 2021 # this block for scalars if num == 0: mat = sym.Matrix([ex._sympy_()]) # row vector with one term sub_exprs, simplified_rhs = sym.cse(mat) # optimise code with open(os.getcwd() + '/' + filename, 'w') as out: for lhs, rhs in sub_exprs: out.write(printer().doprint(Assignment(lhs, rhs)) + '\n') for index, rhs in enumerate(simplified_rhs[0]): lhs = sym.Symbol(name) out.write(printer().doprint(Assignment(lhs, rhs)) + '\n') # this block for tensors else: idx = [] # indices in the form [{x, x}, {x, y} ...] lst = [] # corresponding terms [termxx, termxy, ...] for i in range(len(ex[num])): # num = number of free indices idx.append(str(ex[num][i][0]._sympy_())) # indices for this term lst.append(str(ex[num][i][1]._sympy_())) # the matching term mat = sym.Matrix([lst]) # row vector of terms sub_exprs, simplified_rhs = sym.cse(mat) # optimise code with open(os.getcwd() + '/' + filename, 'w') as out: for lhs, rhs in sub_exprs: out.write(printer().doprint(Assignment(lhs, rhs)) + '\n') for index, rhs in enumerate(simplified_rhs[0]): lhs = sym.Symbol(name + ' (' + idx[index][1:-1] + ')') # lhs = sym.Symbol(name+' '+(idx[index]).replace(', ','][')) out.write(printer().doprint(Assignment(lhs, rhs)) + '\n')
def test_binop_add(): src1 = (src + """\ c = a + b d = a + c s = p + q + r """) expr1.convert_to_expr(src1, 'f') ls1 = expr1.return_expr() for iter in range(8, 11): assert isinstance(ls1[iter], Assignment) assert ls1[8] == Assignment(Variable(Symbol('c')), Symbol('a') + Symbol('b')) assert ls1[9] == Assignment(Variable(Symbol('d')), Symbol('a') + Symbol('c')) assert ls1[10] == Assignment(Variable(Symbol('s')), Symbol('p') + Symbol('q') + Symbol('r'))
def perform_operation(self, lhs, rhs, op): """Performs operation supported by the SymPy core Returns ======= combined_variable: list contains variable content and type of variable """ lhs_value = self.get_expr_for_operand(lhs) rhs_value = self.get_expr_for_operand(rhs) if op == '+': return [Add(lhs_value, rhs_value), 'expr'] if op == '-': return [Add(lhs_value, -rhs_value), 'expr'] if op == '*': return [Mul(lhs_value, rhs_value), 'expr'] if op == '/': return [Mul(lhs_value, Pow(rhs_value, Integer(-1))), 'expr'] if op == '%': return [Mod(lhs_value, rhs_value), 'expr'] if op in ['<', '<=', '>', '>=', '==', '!=']: return [Rel(lhs_value, rhs_value, op), 'expr'] if op == '&&': return [And(as_Boolean(lhs_value), as_Boolean(rhs_value)), 'expr'] if op == '||': return [Or(as_Boolean(lhs_value), as_Boolean(rhs_value)), 'expr'] if op == '=': return [Assignment(Variable(lhs_value), rhs_value), 'expr'] if op in ['+=', '-=', '*=', '/=', '%=']: return [aug_assign(Variable(lhs_value), op[0], rhs_value), 'expr']
def _print_Assignment(self, expr): from sympy.codegen.ast import Assignment 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 test_FunctionPrototype_and_FunctionDefinition(): vx = Variable(x, type=real) vn = Variable(n, type=integer) fp1 = FunctionPrototype(real, 'power', [vx, vn]) assert fp1.return_type == real assert fp1.name == String('power') assert fp1.parameters == Tuple(vx, vn) assert fp1 == FunctionPrototype(real, 'power', [vx, vn]) assert fp1 != FunctionPrototype(real, 'power', [vn, vx]) assert fp1.func(*fp1.args) == fp1 body = [Assignment(x, x**n), Return(x)] fd1 = FunctionDefinition(real, 'power', [vx, vn], body) assert fd1.return_type == real assert str(fd1.name) == 'power' assert fd1.parameters == Tuple(vx, vn) assert fd1.body == CodeBlock(*body) assert fd1 == FunctionDefinition(real, 'power', [vx, vn], body) assert fd1 != FunctionDefinition(real, 'power', [vx, vn], body[::-1]) assert fd1.func(*fd1.args) == fd1 fp2 = FunctionPrototype.from_FunctionDefinition(fd1) assert fp2 == fp1 fd2 = FunctionDefinition.from_FunctionPrototype(fp1, body) assert fd2 == fd1
def doprint(self, expr, assign_to=None): """ Print the expression as code. Parameters ---------- expr : Expression The expression to be printed. assign_to : Symbol, MatrixSymbol, or string (optional) If provided, the printed code will set the expression to a variable with name ``assign_to``. """ from sympy.matrices.expressions.matexpr import MatrixSymbol if isinstance(assign_to, string_types): if expr.is_Matrix: assign_to = MatrixSymbol(assign_to, *expr.shape) else: assign_to = Symbol(assign_to) elif not isinstance(assign_to, (Basic, type(None))): raise TypeError("{0} cannot assign to object of type {1}".format( type(self).__name__, type(assign_to))) if assign_to: expr = Assignment(assign_to, expr) else: # _sympify is not enough b/c it errors on iterables expr = sympify(expr) # keep a set of expressions that are not strictly translatable to Code # and number constants that must be declared and initialized self._not_supported = set() self._number_symbols = set() lines = self._print(expr).splitlines() # format the output if self._settings["human"]: frontlines = [] if len(self._not_supported) > 0: frontlines.append( self._get_comment("Not supported in {0}:".format( self.language))) for expr in sorted(self._not_supported, key=str): frontlines.append(self._get_comment(type(expr).__name__)) for name, value in sorted(self._number_symbols, key=str): frontlines.append(self._declare_number_const(name, value)) lines = frontlines + lines lines = self._format_code(lines) result = "\n".join(lines) else: lines = self._format_code(lines) num_syms = set([(k, self._print(v)) for k, v in self._number_symbols]) result = (num_syms, self._not_supported, "\n".join(lines)) self._not_supported = set() self._number_symbols = set() return result
def c_function_body(expr: Sequence[Expr]) -> CodeBlock: """Prepares a sympy CodeBlock consisting of scalar assignment statements to store the expression in some target variable. This will be used to generate a C function body. :meta private:""" dim = len(expr) out = array_sym("out", dim) return CodeBlock(*(Assignment(out_i, expr_i) for out_i, expr_i in zip(out, expr)))
def _mk_func1(): declars = n, inp, out = Variable('n', integer), Pointer('inp', real), Pointer( 'out', real) i = Variable('i', integer) whl = While(i < n, [Assignment(out[i], inp[i]), PreIncrement(i)]) body = CodeBlock(i.as_Declaration(value=0), whl) return FunctionDefinition(void, 'our_test_function', declars, body)
def test_Subroutine(): # Code to generate the subroutine in the example from # http://www.fortran90.org/src/best-practices.html#arrays r = Symbol("r", real=True) i = Symbol("i", integer=True) v_r = Variable.deduced(r, attrs=(dimension(assumed_extent), intent_out)) v_i = Variable.deduced(i) v_n = Variable("n", integer) do_loop = Do([Assignment(Element(r, [i]), literal_dp(1) / i**2)], i, 1, v_n) sub = Subroutine( "f", [v_r], [ Declaration(v_n), Declaration(v_i), Assignment(v_n, size(r)), do_loop ], ) x = Symbol("x", real=True) v_x3 = Variable.deduced(x, attrs=[dimension(3)]) mod = Module("mymod", definitions=[sub]) prog = Program( "foo", [ use(mod, only=[sub]), Declaration(v_x3), SubroutineCall(sub, [v_x3]), Print([sum_(v_x3), v_x3]), ], ) if not has_fortran(): skip("No fortran compiler found.") (stdout, stderr), info = compile_run_strings( [("a.f90", fcode(mod, standard=90)), ("b.f90", fcode(prog, standard=90))], clean=True, ) ref = [1.0 / i**2 for i in range(1, 4)] assert str(sum(ref))[:-3] in stdout for _ in ref: assert str(_)[:-3] in stdout assert stderr == ""
def test_assignment(): src1 = (src + """\ a = b c = d p = q r = s """) expr1.convert_to_expr(src1, 'f') ls1 = expr1.return_expr() for iter in range(0, 12): if iter < 8: assert isinstance(ls1[iter], Declaration) else: assert isinstance(ls1[iter], Assignment) assert ls1[8] == Assignment(Variable(Symbol('a')), Variable(Symbol('b'))) assert ls1[9] == Assignment(Variable(Symbol('c')), Variable(Symbol('d'))) assert ls1[10] == Assignment(Variable(Symbol('p')), Variable(Symbol('q'))) assert ls1[11] == Assignment(Variable(Symbol('r')), Variable(Symbol('s')))
def test_binop_div(): src1 = ( src + """\ c = a / b d = a / c s = p / q r = q / p """ ) expr1.convert_to_expr(src1, "f") ls1 = expr1.return_expr() for iter in range(8, 12): assert isinstance(ls1[iter], Assignment) assert ls1[8] == Assignment(Variable(Symbol("c")), Symbol("a") / Symbol("b")) assert ls1[9] == Assignment(Variable(Symbol("d")), Symbol("a") / Symbol("c")) assert ls1[10] == Assignment(Variable(Symbol("s")), Symbol("p") / Symbol("q")) assert ls1[11] == Assignment(Variable(Symbol("r")), Symbol("q") / Symbol("p"))
def _handle_assign_to(expr, assign_to): if isinstance(expr, Eq): expr = Assignment(expr.lhs, expr.rhs) if assign_to is None: return sympify(expr) if isinstance(assign_to, (list, tuple)): if len(expr) != len(assign_to): raise ValueError('Failed to assign an expression of length {} to {} variables'.format(len(expr), len(assign_to))) return CodeBlock(*[_handle_assign_to(lhs, rhs) for lhs, rhs in zip(expr, assign_to)]) if isinstance(assign_to, str): if expr.is_Matrix: assign_to = MatrixSymbol(assign_to, *expr.shape) else: assign_to = Symbol(assign_to) elif not isinstance(assign_to, Basic): raise TypeError("{} cannot assign to object of type {}".format( type(self).__name__, type(assign_to))) return Assignment(assign_to, expr)
def test_binop_add(): src1 = ( src + """\ c = a + b d = a + c s = p + q + r """ ) expr1.convert_to_expr(src1, "f") ls1 = expr1.return_expr() for iter in range(8, 11): assert isinstance(ls1[iter], Assignment) assert ls1[8] == Assignment(Variable(Symbol("c")), Symbol("a") + Symbol("b")) assert ls1[9] == Assignment(Variable(Symbol("d")), Symbol("a") + Symbol("c")) assert ls1[10] == Assignment( Variable(Symbol("s")), Symbol("p") + Symbol("q") + Symbol("r") )
def test_Scope(): assign = Assignment(x, y) incr = AddAugmentedAssignment(x, 1) scp = Scope([assign, incr]) cblk = CodeBlock(assign, incr) assert scp.body == cblk assert scp == Scope(cblk) assert scp != Scope([incr, assign]) assert scp.func(*scp.args) == scp
def test_binop_div(): src1 = (src + """\ c = a / b d = a / c s = p / q r = q / p """) expr1.convert_to_expr(src1, 'f') ls1 = expr1.return_expr() for iter in range(8, 12): assert isinstance(ls1[iter], Assignment) assert ls1[8] == Assignment(Variable(Symbol('c')), Symbol('a') / Symbol('b')) assert ls1[9] == Assignment(Variable(Symbol('d')), Symbol('a') / Symbol('c')) assert ls1[10] == Assignment(Variable(Symbol('s')), Symbol('p') / Symbol('q')) assert ls1[11] == Assignment(Variable(Symbol('r')), Symbol('q') / Symbol('p'))
def _mk_func1(): declars = n, inp, out = ( Variable("n", integer), Pointer("inp", real), Pointer("out", real), ) i = Variable("i", integer) whl = While(i < n, [Assignment(out[i], inp[i]), PreIncrement(i)]) body = CodeBlock(i.as_Declaration(value=0), whl) return FunctionDefinition(void, "our_test_function", declars, body)
def expressions(self): # The entry_template is a string of the form "...{}...{}...{}..." # and must have three sets of braces "{}" where the two indices # and the real/imaginary notations ("Re" or "Im") go. # e.g. "f{}{}_{}" -> "f00_Re", "f01_Re", "f01_Im", etc. # # Returns a list of variable assignment expressions # for real values matching entry_template that compose # the elements of M. expressions = [] for i in range(self.size): for j in range(i, self.size): assign_to = sympy.symbols(self.entry_template.format(i,j,"Re"), real=True) expressions.append(Assignment(assign_to, sympy.re(self.H[i,j]))) if j > i: assign_to = sympy.symbols(self.entry_template.format(i,j,"Im"), real=True) expressions.append(Assignment(assign_to, sympy.im(self.H[i,j]))) return expressions
def test_function(): src1 = """\ integer function f(a,b) integer :: x, y f = x + y end function """ expr1.convert_to_expr(src1, "f") for iter in expr1.return_expr(): assert isinstance(iter, FunctionDefinition) assert iter == FunctionDefinition( IntBaseType(String("integer")), name=String("f"), parameters=(Variable(Symbol("a")), Variable(Symbol("b"))), body=CodeBlock( Declaration( Variable( Symbol("a"), type=IntBaseType(String("integer")), value=Integer(0), ) ), Declaration( Variable( Symbol("b"), type=IntBaseType(String("integer")), value=Integer(0), ) ), Declaration( Variable( Symbol("f"), type=IntBaseType(String("integer")), value=Integer(0), ) ), Declaration( Variable( Symbol("x"), type=IntBaseType(String("integer")), value=Integer(0), ) ), Declaration( Variable( Symbol("y"), type=IntBaseType(String("integer")), value=Integer(0), ) ), Assignment(Variable(Symbol("f")), Add(Symbol("x"), Symbol("y"))), Return(Variable(Symbol("f"))), ), )
def test_mul_binop(): src1 = (src + """\ d = a + b - c c = a * b + d s = p * q / r r = p * s + q / p """) expr1.convert_to_expr(src1, 'f') ls1 = expr1.return_expr() for iter in range(8, 12): assert isinstance(ls1[iter], Assignment) assert ls1[8] == Assignment(Variable(Symbol('d')), Symbol('a') + Symbol('b') - Symbol('c')) assert ls1[9] == Assignment(Variable(Symbol('c')), Symbol('a') * Symbol('b') + Symbol('d')) assert ls1[10] == Assignment(Variable(Symbol('s')), Symbol('p') * Symbol('q') / Symbol('r')) assert ls1[11] == Assignment( Variable(Symbol('r')), Symbol('p') * Symbol('s') + Symbol('q') / Symbol('p'))
def moments(self, optimize=True): rho = symbols('rho') u = Matrix(symarray('u', self.descriptor.d)) exprs = [Assignment(rho, sum(self.f_curr))] for i, u_i in enumerate(u): exprs.append( Assignment( u_i, sum([(c_j * self.f_curr[j])[i] for j, c_j in enumerate(self.descriptor.c)]) / sum(self.f_curr))) if optimize: return cse(exprs, optimizations=optimizations.custom, symbols=numbered_symbols(prefix='m')) else: return ([], exprs)
def test_assignment(): src1 = ( src + """\ a = b c = d p = q r = s """ ) expr1.convert_to_expr(src1, "f") ls1 = expr1.return_expr() for iter in range(0, 12): if iter < 8: assert isinstance(ls1[iter], Declaration) else: assert isinstance(ls1[iter], Assignment) assert ls1[8] == Assignment(Variable(Symbol("a")), Variable(Symbol("b"))) assert ls1[9] == Assignment(Variable(Symbol("c")), Variable(Symbol("d"))) assert ls1[10] == Assignment(Variable(Symbol("p")), Variable(Symbol("q"))) assert ls1[11] == Assignment(Variable(Symbol("r")), Variable(Symbol("s")))
def py_write_code(lst, idx, name, filename): import os import sympy as sym # -- 24 sep 2021 ----------------------------------------- old # from sympy.printing.ccode import C99CodePrinter as printer # from sympy.printing.codeprinter import Assignment # -- 24 sep 2021 ----------------------------------------- new # changes due to upgrade of SymPy to v 1.7 from sympy.printing.c import C99CodePrinter as printer from sympy.codegen.ast import Assignment # -------------------------------------------------------- 24 sep 2021 mat = sym.Matrix([lst]) # row vector of terms sub_exprs, simplified_rhs = sym.cse(mat) # optimise code with open(os.getcwd() + '/' + filename, 'w') as out: for lhs, rhs in sub_exprs: out.write(printer().doprint(Assignment(lhs, rhs)) + '\n') # this block for scalars if len(idx) == 0: for index, rhs in enumerate(simplified_rhs[0]): lhs = sym.Symbol(name) out.write(printer().doprint(Assignment(lhs, rhs)) + '\n') # this block for matrices else: for index, rhs in enumerate(simplified_rhs[0]): lhs = sym.Symbol(name + ' (' + idx[index][1:-1] + ')') # lhs = sym.Symbol(name+' '+(idx[index]).replace(', ','][')) out.write(printer().doprint(Assignment(lhs, rhs)) + '\n')
def test_CodeBlock_free_symbols(): c1 = CodeBlock( Assignment(x, y + z), Assignment(z, 1), Assignment(t, x), Assignment(y, 2), ) assert c1.free_symbols == set() c2 = CodeBlock( Assignment(x, y + z), Assignment(z, a * b), Assignment(t, x), Assignment(y, b + 3), ) assert c2.free_symbols == {a, b}
def _print_Assignment(self, expr): from sympy.functions.elementary.piecewise import Piecewise from sympy.matrices.expressions.matexpr import MatrixSymbol from sympy.matrices import MatrixBase, MatrixSlice from sympy.codegen.ast import Assignment 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): elif isinstance(lhs, (MatrixBase, MatrixSymbol, MatrixSlice)): # 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]) code0 = self._print(temp) lines.append(code0) return "\n".join(lines) else: lhs_code = self._print(lhs) rhs_code = self._print(rhs) # hack to avoid the printing of m[i] = m[i] if lhs_code == rhs_code: return "" return self._get_statement("%s = %s" % (lhs_code, rhs_code))
def test_sym_expr(): src1 = (src + """\ d = a + b -c """) expr3 = SymPyExpression(src, 'f') expr4 = SymPyExpression(src1, 'f') ls1 = expr3.return_expr() ls2 = expr4.return_expr() for i in range(0, 7): assert isinstance(ls1[i], Declaration) assert isinstance(ls2[i], Declaration) assert isinstance(ls2[8], Assignment) assert ls1[0] == Declaration( Variable(Symbol('a'), type=IntBaseType(String('integer')), value=Integer(0))) assert ls1[1] == Declaration( Variable(Symbol('b'), type=IntBaseType(String('integer')), value=Integer(0))) assert ls1[2] == Declaration( Variable(Symbol('c'), type=IntBaseType(String('integer')), value=Integer(0))) assert ls1[3] == Declaration( Variable(Symbol('d'), type=IntBaseType(String('integer')), value=Integer(0))) assert ls1[4] == Declaration( Variable(Symbol('p'), type=FloatBaseType(String('real')), value=Float(0.0))) assert ls1[5] == Declaration( Variable(Symbol('q'), type=FloatBaseType(String('real')), value=Float(0.0))) assert ls1[6] == Declaration( Variable(Symbol('r'), type=FloatBaseType(String('real')), value=Float(0.0))) assert ls1[7] == Declaration( Variable(Symbol('s'), type=FloatBaseType(String('real')), value=Float(0.0))) assert ls2[8] == Assignment(Variable(Symbol('d')), Symbol('a') + Symbol('b') - Symbol('c'))
def test_numexpr(): # test ITE rewrite as Piecewise from sympy.logic.boolalg import ITE expr = ITE(x > 0, True, False, evaluate=False) assert NumExprPrinter().doprint(expr) == \ "numexpr.evaluate('where((x > 0), True, False)', truediv=True)" from sympy.codegen.ast import Return, FunctionDefinition, Variable, Assignment func_def = FunctionDefinition( None, 'foo', [Variable(x)], [Assignment(y, x), Return(y**2)]) print("") print(NumExprPrinter().doprint(func_def)) expected = "def foo(x):\n"\ " y = numexpr.evaluate('x', truediv=True)\n"\ " return numexpr.evaluate('y**2', truediv=True)" print(expected) assert NumExprPrinter().doprint(func_def) == expected