Import `sympy.ccode` inside templates instead of as argument

1 files changed, 7 insertions, 5 deletions

diff --git a/boltzgen/kernel/template/momenta_boundary.cl.mako b/boltzgen/kernel/template/momenta_boundary.cl.mako
index 7c1e3df..b0b4c9e 100644
--- a/boltzgen/kernel/template/momenta_boundary.cl.mako
+++ b/boltzgen/kernel/template/momenta_boundary.cl.mako
@@ -1,4 +1,6 @@
 <%
+import sympy
+
 moments_subexpr, moments_assignment = model.moments()
 collision_subexpr, collision_assignment = model.collision(f_eq = model.equilibrium(resolve_moments = False))
 %>
@@ -17,17 +19,17 @@ __kernel void ${name}_momenta_boundary_gid(
 % endfor
 
 % for i, expr in enumerate(moments_subexpr):
-    const ${float_type} ${expr[0]} = ${ccode(expr[1])};
+    const ${float_type} ${expr[0]} = ${sympy.ccode(expr[1])};
 % endfor
 
     ${caller.body()}
 
 % for i, expr in enumerate(collision_subexpr):
-    const ${float_type} ${expr[0]} = ${ccode(expr[1])};
+    const ${float_type} ${expr[0]} = ${sympy.ccode(expr[1])};
 % endfor
 
 % for i, expr in enumerate(collision_assignment):
-    const ${float_type} ${ccode(expr)}
+    const ${float_type} ${sympy.ccode(expr)}
 % endfor
 
 % for i, expr in enumerate(collision_assignment):
@@ -37,7 +39,7 @@ __kernel void ${name}_momenta_boundary_gid(
 </%def>
 
 <%call expr="momenta_boundary('velocity', '%s%d velocity' % (float_type, descriptor.d))">
-    ${float_type} ${ccode(moments_assignment[0])}
+    ${float_type} ${sympy.ccode(moments_assignment[0])}
 % for i, expr in enumerate(moments_assignment[1:]):
     ${float_type} ${expr.lhs} = velocity.${['x', 'y', 'z'][i]};
 % endfor
@@ -46,7 +48,7 @@ __kernel void ${name}_momenta_boundary_gid(
 <%call expr="momenta_boundary('density', '%s density' % float_type)">
     ${float_type} ${moments_assignment[0].lhs} = density;
 % for i, expr in enumerate(moments_assignment[1:]):
-    ${float_type} ${ccode(expr)}
+    ${float_type} ${sympy.ccode(expr)}
 % endfor
 </%call>