Commit 54bd44ef authored by Martin Bauer's avatar Martin Bauer
Browse files

Merge branch 'abs' into 'master'

Correctly print sp.Abs in generated code

See merge request pycodegen/pystencils!127
parents a5ab1070 b1e9b1e4
......@@ -361,6 +361,12 @@ class CustomSympyPrinter(CCodePrinter):
result = super(CustomSympyPrinter, self)._print_Piecewise(expr)
return result.replace("\n", "")
def _print_Abs(self, expr):
if expr.is_integer:
return 'abs({0})'.format(self._print(expr.args[0]))
else:
return 'fabs({0})'.format(self._print(expr.args[0]))
def _print_Type(self, node):
return str(node)
......
......@@ -176,6 +176,10 @@ def insert_vector_casts(ast_node):
visit_expr(expr.args[4]))
elif isinstance(expr, cast_func):
return expr
elif expr.func is sp.Abs:
new_arg = visit_expr(expr.args[0])
pw = sp.Piecewise((-1 * new_arg, new_arg < 0), (new_arg, True))
return visit_expr(pw)
elif expr.func in handled_functions or isinstance(expr, sp.Rel) or isinstance(expr, BooleanFunction):
new_args = [visit_expr(a) for a in expr.args]
arg_types = [get_type_of_expression(a) for a in new_args]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment