def test_CodeBlock_cse():
c1 = CodeBlock(
Assignment(y, 1),
Assignment(x, sin(y)),
Assignment(z, sin(y)),
Assignment(t, x*z),
)
assert c1.cse() == CodeBlock(
Assignment(y, 1),
Assignment(x0, sin(y)),
Assignment(x, x0),
Assignment(z, x0),
Assignment(t, x*z),
)
# Multiple assignments to same symbol not supported
raises(NotImplementedError, lambda: CodeBlock(
Assignment(x, 1),
Assignment(y, 1), Assignment(y, 2)
).cse())
# Check auto-generated symbols do not collide with existing ones
c2 = CodeBlock(
Assignment(x0, sin(y) + 1),
Assignment(x1, 2 * sin(y)),
Assignment(z, x * y),
)
assert c2.cse() == CodeBlock(
Assignment(x2, sin(y)),
Assignment(x0, x2 + 1),
Assignment(x1, 2 * x2),
Assignment(z, x * y),
)
def test_CodeBlock_topological_sort():
assignments = [
Assignment(x, y + z),
Assignment(z, 1),
Assignment(t, x),
Assignment(y, 2),
]
ordered_assignments = [
# Note that the unrelated z=1 and y=2 are kept in that order
Assignment(z, 1),
Assignment(y, 2),
Assignment(x, y + z),
Assignment(t, x),
]
c = CodeBlock.topological_sort(assignments)
assert c == CodeBlock(*ordered_assignments)
# Cycle
invalid_assignments = [
Assignment(x, y + z),
Assignment(z, 1),
Assignment(y, x),
Assignment(y, 2),
]
raises(ValueError, lambda: CodeBlock.topological_sort(invalid_assignments))
# Undefined variable
invalid_assignments = [
Assignment(x, y)
]
raises(ValueError, lambda: CodeBlock.topological_sort(invalid_assignments))
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 test_CodeBlock_cse():
c = CodeBlock(
Assignment(y, 1),
Assignment(x, sin(y)),
Assignment(z, sin(y)),
Assignment(t, x*z),
)
assert c.cse() == CodeBlock(
Assignment(y, 1),
Assignment(x0, sin(y)),
Assignment(x, x0),
Assignment(z, x0),
Assignment(t, x*z),
)
raises(NotImplementedError, lambda: CodeBlock(Assignment(x, 1),
Assignment(y, 1), Assignment(y, 2)).cse())
def test_CodeBlock_topological_sort():
assignments = [
Assignment(x, y + z),
Assignment(z, 1),
Assignment(t, x),
Assignment(y, 2),
]
ordered_assignments = [
# Note that the unrelated z=1 and y=2 are kept in that order
Assignment(z, 1),
Assignment(y, 2),
Assignment(x, y + z),
Assignment(t, x),
]
c1 = CodeBlock.topological_sort(assignments)
assert c1 == CodeBlock(*ordered_assignments)
# Cycle
invalid_assignments = [
Assignment(x, y + z),
Assignment(z, 1),
Assignment(y, x),
Assignment(y, 2),
]
raises(ValueError, lambda: CodeBlock.topological_sort(invalid_assignments))
# Free symbols
free_assignments = [
Assignment(x, y + z),
Assignment(z, a * b),
Assignment(t, x),
Assignment(y, b + 3),
]
free_assignments_ordered = [
Assignment(z, a * b),
Assignment(y, b + 3),
Assignment(x, y + z),
Assignment(t, x),
]
c2 = CodeBlock.topological_sort(free_assignments)
assert c2 == CodeBlock(*free_assignments_ordered)
def main():
NUCS = DATA['nucs']
nucs = ['K40']
system = CodeBlock(*list(map(gennuc, nucs)))
sigma_map = sigma_symbol_to_indexed()
nuc_map = nuc_symbol_to_indexed()
# nuc_map = {}
system = system.xreplace({**sigma_map, **nuc_map})
sigma_array = generate_sigma_array()
code = sympy.ccode(system, order='none')
generated_code = TEMPLATE
input_data = [0.0]*len(sigma_array)
input_data[NUCS.index("K39")] = 1.0
input_time = 81.0
for val, repl in {
"I": len(sigma_array),
"T": input_time,
"PHI": 4e-10,
"SIGMA_ARRAY": str(sigma_array).replace('[', '{').replace(']', '}'),
"CODE": textwrap.indent(code, ' '),
# For testing
"N0": str(input_data).replace('[', '{').replace(']', '}'),
}.items():
generated_code = generated_code.replace("%%" + val + "%%", str(repl))
with open("sigma_array.txt", 'w') as f:
f.write('[' + ',\n'.join(map(str, sigma_array)) + ']\n')
# with open('system.txt', 'w') as f:
# for eq in system.args:
# f.write(str(eq) + '\n')
with open('system-C.txt', 'w') as f:
f.write(code)
with open('transmute.c', 'w') as f:
f.write(generated_code)
def test_ast_replace():
x = Variable('x', real)
y = Variable('y', real)
n = Variable('n', integer)
pwer = FunctionDefinition(real, 'pwer', [x, n], [pow(x.symbol, n.symbol)])
pname = pwer.name
pcall = FunctionCall('pwer', [y, 3])
tree1 = CodeBlock(pwer, pcall)
assert str(tree1.args[0].name) == 'pwer'
assert str(tree1.args[1].name) == 'pwer'
for a, b in zip(tree1, [pwer, pcall]):
assert a == b
tree2 = tree1.replace(pname, String('power'))
assert str(tree1.args[0].name) == 'pwer'
assert str(tree1.args[1].name) == 'pwer'
assert str(tree2.args[0].name) == 'power'
assert str(tree2.args[1].name) == 'power'
def _construct_body(cls, itr):
if isinstance(itr, CodeBlock):
return itr
else:
return CodeBlock(*itr)
def test_CodeBlock():
c = CodeBlock(Assignment(x, 1), Assignment(y, x + 1))
assert c.func(*c.args) == c
assert c.left_hand_sides == Tuple(x, y)
assert c.right_hand_sides == Tuple(1, x + 1)
def test_function():
c_src1 = (
'void fun1()' + '\n' +
'{' + '\n' +
'int a;' + '\n' +
'}'
)
c_src2 = (
'int fun2()' + '\n' +
'{'+ '\n' +
'int a;' + '\n' +
'return a;' + '\n' +
'}'
)
c_src3 = (
'float fun3()' + '\n' +
'{' + '\n' +
'float b;' + '\n' +
'return b;' + '\n' +
'}'
)
c_src4 = (
'float fun4()' + '\n' +
'{}'
)
res1 = SymPyExpression(c_src1, 'c').return_expr()
res2 = SymPyExpression(c_src2, 'c').return_expr()
res3 = SymPyExpression(c_src3, 'c').return_expr()
res4 = SymPyExpression(c_src4, 'c').return_expr()
assert res1[0] == FunctionDefinition(
NoneToken(),
name=String('fun1'),
parameters=(),
body=CodeBlock(
Declaration(
Variable(
Symbol('a'),
type=IntBaseType(String('integer')),
value=Integer(0)
)
)
)
)
assert res2[0] == FunctionDefinition(
IntBaseType(String('integer')),
name=String('fun2'),
parameters=(),
body=CodeBlock(
Declaration(
Variable(
Symbol('a'),
type=IntBaseType(String('integer')),
value=Integer(0)
)
),
Return('a')
)
)
assert res3[0] == FunctionDefinition(
FloatBaseType(String('real')),
name=String('fun3'),
parameters=(),
body=CodeBlock(
Declaration(
Variable(
Symbol('b'),
type=FloatBaseType(String('real')),
value=Float('0.0', precision=53)
)
),
Return('b')
)
)
assert res4[0] == FunctionPrototype(
FloatBaseType(String('real')),
name=String('fun4'),
parameters=()
)
def newtons_method(expr,
wrt,
atol=1e-12,
delta=None,
debug=False,
itermax=None,
counter=None):
""" Generates an AST for Newton-Raphson method (a root-finding algorithm).
Explanation
===========
Returns an abstract syntax tree (AST) based on ``sympy.codegen.ast`` for Netwon's
method of root-finding.
Parameters
==========
expr : expression
wrt : Symbol
With respect to, i.e. what is the variable.
atol : number or expr
Absolute tolerance (stopping criterion)
delta : Symbol
Will be a ``Dummy`` if ``None``.
debug : bool
Whether to print convergence information during iterations
itermax : number or expr
Maximum number of iterations.
counter : Symbol
Will be a ``Dummy`` if ``None``.
Examples
========
>>> from sympy import symbols, cos
>>> from sympy.codegen.ast import Assignment
>>> from sympy.codegen.algorithms import newtons_method
>>> x, dx, atol = symbols('x dx atol')
>>> expr = cos(x) - x**3
>>> algo = newtons_method(expr, x, atol, dx)
>>> algo.has(Assignment(dx, -expr/expr.diff(x)))
True
References
==========
.. [1] https://en.wikipedia.org/wiki/Newton%27s_method
"""
if delta is None:
delta = Dummy()
Wrapper = Scope
name_d = 'delta'
else:
Wrapper = lambda x: x
name_d = delta.name
delta_expr = -expr / expr.diff(wrt)
whl_bdy = [
Assignment(delta, delta_expr),
AddAugmentedAssignment(wrt, delta)
]
if debug:
prnt = Print([wrt, delta],
r"{}=%12.5g {}=%12.5g\n".format(wrt.name, name_d))
whl_bdy = [whl_bdy[0], prnt] + whl_bdy[1:]
req = Gt(Abs(delta), atol)
declars = [Declaration(Variable(delta, type=real, value=oo))]
if itermax is not None:
counter = counter or Dummy(integer=True)
v_counter = Variable.deduced(counter, 0)
declars.append(Declaration(v_counter))
whl_bdy.append(AddAugmentedAssignment(counter, 1))
req = And(req, Lt(counter, itermax))
whl = While(req, CodeBlock(*whl_bdy))
blck = declars + [whl]
return Wrapper(CodeBlock(*blck))
def newtons_method_function(expr,
wrt,
params=None,
func_name="newton",
attrs=Tuple(),
*,
delta=None,
**kwargs):
""" Generates an AST for a function implementing the Newton-Raphson method.
Parameters
==========
expr : expression
wrt : Symbol
With respect to, i.e. what is the variable
params : iterable of symbols
Symbols appearing in expr that are taken as constants during the iterations
(these will be accepted as parameters to the generated function).
func_name : str
Name of the generated function.
attrs : Tuple
Attribute instances passed as ``attrs`` to ``FunctionDefinition``.
\\*\\*kwargs :
Keyword arguments passed to :func:`sympy.codegen.algorithms.newtons_method`.
Examples
========
>>> from sympy import symbols, cos
>>> from sympy.codegen.algorithms import newtons_method_function
>>> from sympy.codegen.pyutils import render_as_module
>>> x = symbols('x')
>>> expr = cos(x) - x**3
>>> func = newtons_method_function(expr, x)
>>> py_mod = render_as_module(func) # source code as string
>>> namespace = {}
>>> exec(py_mod, namespace, namespace)
>>> res = eval('newton(0.5)', namespace)
>>> abs(res - 0.865474033102) < 1e-12
True
See Also
========
sympy.codegen.algorithms.newtons_method
"""
if params is None:
params = (wrt, )
pointer_subs = {
p.symbol: Symbol('(*%s)' % p.symbol.name)
for p in params if isinstance(p, Pointer)
}
if delta is None:
delta = Symbol('d_' + wrt.name)
if expr.has(delta):
delta = None # will use Dummy
algo = newtons_method(expr, wrt, delta=delta,
**kwargs).xreplace(pointer_subs)
if isinstance(algo, Scope):
algo = algo.body
not_in_params = expr.free_symbols.difference(
{_symbol_of(p)
for p in params})
if not_in_params:
raise ValueError("Missing symbols in params: %s" %
', '.join(map(str, not_in_params)))
declars = tuple(Variable(p, real) for p in params)
body = CodeBlock(algo, Return(wrt))
return FunctionDefinition(real, func_name, declars, body, attrs=attrs)
def assign_cse(target, expr):
cses, (new_expr, ) = cse(expr)
cse_declars = [Declaration(*args, const=True) for args in cses]
blck = cse_declars + [Assignment(target, new_expr)]
return Scope(CodeBlock(*blck))
def test_parameters():
c_src1 = "void fun1( int a)" + "\n" + "{" + "\n" + "int i;" + "\n" + "}"
c_src2 = ("int fun2(float x, float y)" + "\n" + "{" + "\n" + "int a;" +
"\n" + "return a;" + "\n" + "}")
c_src3 = ("float fun3(int p, float q, int r)" + "\n" + "{" + "\n" +
"float b;" + "\n" + "return b;" + "\n" + "}")
res1 = SymPyExpression(c_src1, "c").return_expr()
res2 = SymPyExpression(c_src2, "c").return_expr()
res3 = SymPyExpression(c_src3, "c").return_expr()
assert res1[0] == FunctionDefinition(
NoneToken(),
name=String("fun1"),
parameters=(Variable(Symbol("a"),
type=IntBaseType(String("integer")),
value=Integer(0)), ),
body=CodeBlock(
Declaration(
Variable(
Symbol("i"),
type=IntBaseType(String("integer")),
value=Integer(0),
))),
)
assert res2[0] == FunctionDefinition(
IntBaseType(String("integer")),
name=String("fun2"),
parameters=(
Variable(
Symbol("x"),
type=FloatBaseType(String("real")),
value=Float("0.0", precision=53),
),
Variable(
Symbol("y"),
type=FloatBaseType(String("real")),
value=Float("0.0", precision=53),
),
),
body=CodeBlock(
Declaration(
Variable(
Symbol("a"),
type=IntBaseType(String("integer")),
value=Integer(0),
)),
Return("a"),
),
)
assert res3[0] == FunctionDefinition(
FloatBaseType(String("real")),
name=String("fun3"),
parameters=(
Variable(Symbol("p"),
type=IntBaseType(String("integer")),
value=Integer(0)),
Variable(
Symbol("q"),
type=FloatBaseType(String("real")),
value=Float("0.0", precision=53),
),
Variable(Symbol("r"),
type=IntBaseType(String("integer")),
value=Integer(0)),
),
body=CodeBlock(
Declaration(
Variable(
Symbol("b"),
type=FloatBaseType(String("real")),
value=Float("0.0", precision=53),
)),
Return("b"),
),
)
def test_parameters():
c_src1 = (
'void fun1( int a)' + '\n' +
'{' + '\n' +
'int i;' + '\n' +
'}'
)
c_src2 = (
'int fun2(float x, float y)' + '\n' +
'{'+ '\n' +
'int a;' + '\n' +
'return a;' + '\n' +
'}'
)
c_src3 = (
'float fun3(int p, float q, int r)' + '\n' +
'{' + '\n' +
'float b;' + '\n' +
'return b;' + '\n' +
'}'
)
res1 = SymPyExpression(c_src1, 'c').return_expr()
res2 = SymPyExpression(c_src2, 'c').return_expr()
res3 = SymPyExpression(c_src3, 'c').return_expr()
assert res1[0] == FunctionDefinition(
NoneToken(),
name=String('fun1'),
parameters=(
Variable(
Symbol('a'),
type=IntBaseType(String('integer')),
value=Integer(0)
),
),
body=CodeBlock(
Declaration(
Variable(
Symbol('i'),
type=IntBaseType(String('integer')),
value=Integer(0)
)
)
)
)
assert res2[0] == FunctionDefinition(
IntBaseType(String('integer')),
name=String('fun2'),
parameters=(
Variable(
Symbol('x'),
type=FloatBaseType(String('real')),
value=Float('0.0', precision=53)
),
Variable(
Symbol('y'),
type=FloatBaseType(String('real')),
value=Float('0.0', precision=53)
)
),
body=CodeBlock(
Declaration(
Variable(
Symbol('a'),
type=IntBaseType(String('integer')),
value=Integer(0)
)
),
Return('a')
)
)
assert res3[0] == FunctionDefinition(
FloatBaseType(String('real')), name=String('fun3'),
parameters=(
Variable(
Symbol('p'),
type=IntBaseType(String('integer')),
value=Integer(0)
),
Variable(
Symbol('q'),
type=FloatBaseType(String('real')),
value=Float('0.0', precision=53)
),
Variable(
Symbol('r'),
type=IntBaseType(String('integer')),
value=Integer(0)
)
),
body=CodeBlock(
Declaration(
Variable(
Symbol('b'),
type=FloatBaseType(String('real')),
value=Float('0.0', precision=53)
)
),
Return('b')
)
)
tmp_init = VariableWithInit("n",tmp,type=Type("std::vector<T>")).as_Declaration()
bc_tmps = []
for e in bc_eqs2:
bc_tmps.append(Variable(e.lhs, type=templateT).as_Declaration(value=e.rhs))
body = [tmp_init]
body.extend(bc_tmps)
body.extend([convert_eq_to_assignment(teq1c),
convert_eq_to_assignment(teq2c),
loop1,
convert_eq_to_assignment(teq5c),
loop2])
algo = CodeBlock(*body)
# Set up to create a template function
tx = Pointer(x,type=templateT)
ty = Pointer(y,type=templateT)
ty2 = Pointer(y2,type=templateT)
yp0_var = Variable('yp0',type=templateT)
ypn_var = Variable('ypn',type=templateT)
tf = TemplateFunctionDefinition(void, "CubicSplineSolve",[tx,ty,n,yp0_var,ypn_var,ty2],[templateT],algo)
ACP = ACodePrinter()
gen_comment = Comment("Generated by gen_cubic_spline_solver.py")
cb = CodeBlock(gen_comment, tf)
s = ACP.doprint(cb)
def test_CodeBlock():
c = CodeBlock(Assignment(x, 1), Assignment(y, x + 1))
assert c.func(*c.args) == c
assert c.left_hand_sides == Tuple(x, y)
assert c.right_hand_sides == Tuple(1, x + 1)
def _construct_declarations(cls, args):
args = [Str(arg) if isinstance(arg, str) else arg for arg in args]
return CodeBlock(*args)