Merge pull request #278 from OceanParcels/ckernels

For JITParticles, Kernels can now also be written directly in C (instead of Python translation to C)

Author: erikvansebille

parcels/codegenerator.py

@@ -300,10 +300,21 @@ def visit_Call(self, node):
         """Generate C code for simple C-style function calls. Please
         note that starred and keyword arguments are currently not
         supported."""
+        pointer_args = False
         for a in node.args:
             self.visit(a)
-        ccode_args = ", ".join([a.ccode for a in node.args])
+            if a.ccode == 'pointer_args':
+                pointer_args = True
+                continue
+            if isinstance(a, FieldNode):
+                a.ccode = a.obj.name
+            elif isinstance(a, ParticleNode):
+                continue
+            elif pointer_args:
+                a.ccode = "&%s" % a.ccode
+        ccode_args = ", ".join([a.ccode for a in node.args[pointer_args:]])
         try:
+            self.visit(node.func)
             node.ccode = "%s(%s)" % (node.func.ccode, ccode_args)
         except:
             raise RuntimeError("Error in converting Kernel to C. See http://oceanparcels.org/#writing-parcels-kernels for hints and tips")
@@ -543,7 +554,7 @@ def __init__(self, fieldset, ptype=None):
         self.fieldset = fieldset
         self.ptype = ptype
 
-    def generate(self, funcname, field_args, const_args, kernel_ast):
+    def generate(self, funcname, field_args, const_args, kernel_ast, c_include):
         ccode = []
 
         # Add include for Parcels and math header
@@ -560,6 +571,9 @@ def generate(self, funcname, field_args, const_args, kernel_ast):
 
         ccode += [str(c.Typedef(c.GenerableStruct("", vdecl, declname=self.ptype.name)))]
 
+        if c_include:
+            ccode += [c_include]
+
         # Insert kernel code
         ccode += [str(kernel_ast)]
 

parcels/include/parcels.h

@@ -57,6 +57,7 @@ typedef struct
   CGridIndex *gridIndices;
 } CGridIndexSet;  
 
+
 static inline ErrorCode search_indices_vertical_z(float z, int zdim, float *zvals, int *k, double *zeta)
 {
   if (z < zvals[0] || z > zvals[zdim-1]) {return ERROR_OUT_OF_BOUNDS;}

parcels/kernel.py

@@ -48,7 +48,7 @@ class Kernel(object):
     """
 
     def __init__(self, fieldset, ptype, pyfunc=None, funcname=None,
-                 funccode=None, py_ast=None, funcvars=None):
+                 funccode=None, py_ast=None, funcvars=None, c_include=""):
         self.fieldset = fieldset
         self.ptype = ptype
 
@@ -104,8 +104,13 @@ def __init__(self, fieldset, ptype, pyfunc=None, funcname=None,
                 del self.field_args['UV']
             self.const_args = kernelgen.const_args
             loopgen = LoopGenerator(fieldset, ptype)
+            if path.isfile(c_include):
+                with open(c_include, 'r') as f:
+                    c_include_str = f.read()
+            else:
+                c_include_str = c_include
             self.ccode = loopgen.generate(self.funcname, self.field_args, self.const_args,
-                                          kernel_ccode)
+                                          kernel_ccode, c_include_str)
 
             basename = path.join(get_cache_dir(), self._cache_key)
             self.src_file = "%s.c" % basename

parcels/particleset.py

@@ -545,10 +545,10 @@ def density(self, field=None, particle_val=None, relative=False, area_scale=True
 
         return Density
 
-    def Kernel(self, pyfunc):
+    def Kernel(self, pyfunc, c_include=""):
         """Wrapper method to convert a `pyfunc` into a :class:`parcels.kernel.Kernel` object
         based on `fieldset` and `ptype` of the ParticleSet"""
-        return Kernel(self.fieldset, self.ptype, pyfunc=pyfunc)
+        return Kernel(self.fieldset, self.ptype, pyfunc=pyfunc, c_include=c_include)
 
     def ParticleFile(self, *args, **kwargs):
         """Wrapper method to initialise a :class:`parcels.particlefile.ParticleFile`

tests/customed_header.h

@@ -0,0 +1,7 @@
+
+static inline void func(CField *f, float *lon, float *dt)
+{
+  float (*data)[f->xdim] = (float (*)[f->xdim]) f->data;
+  float u = data[2][1];
+  *lon += u * *dt;
+}

tests/test_kernel_language.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pytest
 import random as py_random
+from os import path
 
 
 ptype = {'scipy': ScipyParticle, 'jit': JITParticle}
@@ -150,3 +151,38 @@ class TestParticle(ptype[mode]):
                          'random.%s(%s)' % (rngfunc, ', '.join([str(a) for a in rngargs])))
     pset.execute(kernel, endtime=1., dt=1.)
     assert np.allclose(np.array([p.p for p in pset]), series, rtol=1e-12)
+
+
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+@pytest.mark.parametrize('c_inc', ['str', 'file'])
+def test_c_kernel(fieldset, mode, c_inc):
+    pset = ParticleSet(fieldset, pclass=ptype[mode], lon=[0.5], lat=[0])
+
+    def func(U, lon, dt):
+        u = U.data[0, 2, 1]
+        return lon + u * dt
+
+    if c_inc == 'str':
+        c_include = """
+                 static inline void func(CField *f, float *lon, float *dt)
+                 {
+                   float (*data)[f->xdim] = (float (*)[f->xdim]) f->data;
+                   float u = data[2][1];
+                   *lon += u * *dt;
+                 }
+                 """
+    else:
+        c_include = path.join(path.dirname(__file__), 'customed_header.h')
+
+    def ckernel(particle, fieldset, time, dt):
+        func('pointer_args', fieldset.U, particle.lon, dt)
+
+    def pykernel(particle, fieldset, time, dt):
+        particle.lon = func(fieldset.U, particle.lon, dt)
+
+    if mode == 'scipy':
+        kernel = pset.Kernel(pykernel)
+    else:
+        kernel = pset.Kernel(ckernel, c_include=c_include)
+    pset.execute(kernel, endtime=3., dt=3.)
+    assert np.allclose(pset[0].lon, 0.81578948)