Merge pull request #424 from OceanParcels/JIT-support-power-operator

Adding support for power ('**') operator in JIT mode

Author: erikvansebille

parcels/codegenerator.py

@@ -538,7 +538,10 @@ def visit_BinOp(self, node):
         self.visit(node.left)
         self.visit(node.op)
         self.visit(node.right)
-        node.ccode = "(%s %s %s)" % (node.left.ccode, node.op.ccode, node.right.ccode)
+        if node.op.ccode == 'pow':  # catching '**' pow statements
+            node.ccode = "pow(%s, %s)" % (node.left.ccode, node.right.ccode)
+        else:
+            node.ccode = "(%s %s %s)" % (node.left.ccode, node.op.ccode, node.right.ccode)
         node.s_print = True
 
     def visit_Add(self, node):
@@ -562,6 +565,9 @@ def visit_Div(self, node):
     def visit_Mod(self, node):
         node.ccode = "%"
 
+    def visit_Pow(self, node):
+        node.ccode = "pow"
+
     def visit_Num(self, node):
         node.ccode = str(node.n)
 

tests/test_kernel_language.py

@@ -53,6 +53,7 @@ class TestParticle(ptype[mode]):
     ('Sub', '6. - 2.', 4),
     ('Mul', '3. * 5.', 15),
     ('Div', '24. / 4.', 6),
+    ('Pow', '2 ** 3', 8),
 ])
 def test_expression_float(fieldset, mode, name, expr, result, npart=10):
     """ Test basic arithmetic expressions """