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