Repository: OceanParcels/parcels
Pull request #388 (2018-06-25 12:36:26)

Vector fields


Patch:
diff --git a/parcels/codegenerator.py b/parcels/codegenerator.py
index 89c91efd..2bec7cbe 100644
--- a/parcels/codegenerator.py
+++ b/parcels/codegenerator.py
@@ -1,4 +1,4 @@
-from parcels.field import Field
+from parcels.field import Field, VectorField
 from parcels.loggers import logger
 import ast
 import cgen as c
@@ -20,6 +20,9 @@ class FieldSetNode(IntrinsicNode):
         if isinstance(getattr(self.obj, attr), Field):
             return FieldNode(getattr(self.obj, attr),
                              ccode="%s->%s" % (self.ccode, attr))
+        elif isinstance(getattr(self.obj, attr), VectorField):
+            return VectorFieldNode(getattr(self.obj, attr),
+                                   ccode="%s->%s" % (self.ccode, attr))
         else:
             return ConstNode(getattr(self.obj, attr),
                              ccode="%s" % (attr))
@@ -35,6 +38,18 @@ class FieldEvalNode(IntrinsicNode):
         self.field = field
         self.args = args
         self.var = var  # the variable in which the interpolated field is written
+
+
+class VectorFieldNode(IntrinsicNode):
+    def __getitem__(self, attr):
+        return VectorFieldEvalNode(self.obj, attr)
+
+
+class VectorFieldEvalNode(IntrinsicNode):
+    def __init__(self, field, args, var, var2=None):
+        self.field = field
+        self.args = args
+        self.var = var  # the variable in which the interpolated field is written
         self.var2 = var2  # second variable for UV interpolation
 
 
@@ -164,11 +179,17 @@ class IntrinsicTransformer(ast.NodeTransformer):
         # temporary variable and put the evaluation call on the stack.
         if isinstance(node.value, FieldNode):
             tmp = self.get_tmp()
-            tmp2 = self.get_tmp() if node.value.obj.name == 'UV' else None
             # Insert placeholder node for field eval ...
-            self.stmt_stack += [FieldEvalNode(node.value, node.slice, tmp, tmp2)]
+            self.stmt_stack += [FieldEvalNode(node.value, node.slice, tmp)]
+            # .. and return the name of the temporary that will be populated
+            return ast.Name(id=tmp)
+        elif isinstance(node.value, VectorFieldNode):
+            tmp = self.get_tmp()
+            tmp2 = self.get_tmp()
+            # Insert placeholder node for field eval ...
+            self.stmt_stack += [VectorFieldEvalNode(node.value, node.slice, tmp, tmp2)]
             # .. and return the name of the temporary that will be populated
-            return ast.Tuple([ast.Name(id=tmp), ast.Name(id=tmp2)], ast.Load()) if tmp2 else ast.Name(id=tmp)
+            return ast.Tuple([ast.Name(id=tmp), ast.Name(id=tmp2)], ast.Load())
         else:
             return node
 
@@ -237,6 +258,7 @@ class KernelGenerator(ast.NodeVisitor):
         self.field_args = OrderedDict()
         # Hack alert: JIT requires U field to update fieldset indexes
         self.field_args['U'] = fieldset.U
+        self.vector_field_args = OrderedDict()
         self.const_args = OrderedDict()
 
     def generate(self, py_ast, funcvars):
@@ -283,19 +305,17 @@ class KernelGenerator(ast.NodeVisitor):
         args = [c.Pointer(c.Value(self.ptype.name, "particle")),
                 c.Value("double", "time"), c.Value("float", "dt")]
         for field_name, field in self.field_args.items():
-            if field_name != 'UV':
-                args += [c.Pointer(c.Value("CField", "%s" % field_name))]
-        for field_name, field in self.field_args.items():
-            if field_name == 'UV':
-                fieldset = field.fieldset
-                for f in ['U', 'V', 'cosU', 'sinU', 'cosV', 'sinV']:
-                    try:
-                        getattr(fieldset, f)
-                        if f not in self.field_args:
-                            args += [c.Pointer(c.Value("CField", "%s" % f))]
-                    except:
-                        if fieldset.U.grid.gtype in [GridCode.CurvilinearZGrid, GridCode.CurvilinearSGrid]:
-                            raise RuntimeError("cosU, sinU, cosV and sinV fields must be defined for a proper rotation of U, V fields in curvilinear grids")
+            args += [c.Pointer(c.Value("CField", "%s" % field_name))]
+        for field_name, field in self.vector_field_args.items():
+            fieldset = field.fieldset
+            for f in [field.U.name, field.V.name, 'cosU', 'sinU', 'cosV', 'sinV']:
+                try:
+                    getattr(fieldset, f)
+                    if f not in self.field_args:
+                        args += [c.Pointer(c.Value("CField", "%s" % f))]
+                except:
+                    if fieldset.U.grid.gtype in [GridCode.CurvilinearZGrid, GridCode.CurvilinearSGrid]:
+                        raise RuntimeError("cosU, sinU, cosV and sinV fields must be defined for a proper rotation of U, V fields in curvilinear grids")
         for const, _ in self.const_args.items():
             args += [c.Value("float", const)]
 
@@ -343,7 +363,7 @@ class KernelGenerator(ast.NodeVisitor):
                 if a.ccode == 'pointer_args':
                     pointer_args = True
                     continue
-                if isinstance(a, FieldNode):
+                if isinstance(a, FieldNode) or isinstance(a, VectorFieldNode):
                     a.ccode = a.obj.name
                 elif isinstance(a, ParticleNode):
                     continue
@@ -445,7 +465,7 @@ class KernelGenerator(ast.NodeVisitor):
     def visit_Subscript(self, node):
         self.visit(node.value)
         self.visit(node.slice)
-        if isinstance(node.value, FieldNode):
+        if isinstance(node.value, FieldNode) or isinstance(node.value, VectorFieldNode):
             node.ccode = node.value.__getitem__(node.slice.ccode).ccode
         elif isinstance(node.value, IntrinsicNode):
             raise NotImplementedError("Subscript not implemented for object type %s"
@@ -539,22 +559,32 @@ class KernelGenerator(ast.NodeVisitor):
         """Record intrinsic fields used in kernel"""
         self.field_args[node.obj.name] = node.obj
 
+    def visit_VectorFieldNode(self, node):
+        """Record intrinsic fields used in kernel"""
+        self.vector_field_args[node.obj.name] = node.obj
+
     def visit_ConstNode(self, node):
         self.const_args[node.ccode] = node.obj
 
     def visit_FieldEvalNode(self, node):
         self.visit(node.field)
         self.visit(node.args)
-        if node.var2:  # evaluation UV Field
-            ccode_eval = node.field.obj.ccode_evalUV(node.var, node.var2, *node.args.ccode)
-            ccode_conv1 = node.field.obj.fieldset.U.ccode_convert(*node.args.ccode)
-            ccode_conv2 = node.field.obj.fieldset.V.ccode_convert(*node.args.ccode)
-            conv_stat = c.Block([c.Statement("%s *= %s" % (node.var, ccode_conv1)),
-                                 c.Statement("%s *= %s" % (node.var2, ccode_conv2))])
-        else:
-            ccode_eval = node.field.obj.ccode_eval(node.var, *node.args.ccode)
-            ccode_conv = node.field.obj.ccode_convert(*node.args.ccode)
-            conv_stat = c.Statement("%s *= %s" % (node.var, ccode_conv))
+        ccode_eval = node.field.obj.ccode_eval(node.var, *node.args.ccode)
+        ccode_conv = node.field.obj.ccode_convert(*node.args.ccode)
+        conv_stat = c.Statement("%s *= %s" % (node.var, ccode_conv))
+        node.ccode = c.Block([c.Assign("err", ccode_eval),
+                              conv_stat, c.Statement("CHECKERROR(err)")])
+
+    def visit_VectorFieldEvalNode(self, node):
+        self.visit(node.field)
+        self.visit(node.args)
+        ccode_eval = node.field.obj.ccode_eval(node.var, node.var2,
+                                               node.field.obj.U, node.field.obj.V,
+                                               *node.args.ccode)
+        ccode_conv1 = node.field.obj.U.ccode_convert(*node.args.ccode)
+        ccode_conv2 = node.field.obj.V.ccode_convert(*node.args.ccode)
+        conv_stat = c.Block([c.Statement("%s *= %s" % (node.var, ccode_conv1)),
+                             c.Statement("%s *= %s" % (node.var2, ccode_conv2))])
         node.ccode = c.Block([c.Assign("err", ccode_eval),
                               conv_stat, c.Statement("CHECKERROR(err)")])
 
diff --git a/parcels/field.py b/parcels/field.py
index fb39d43a..a63dd4b7 100644
--- a/parcels/field.py
+++ b/parcels/field.py
@@ -12,7 +12,7 @@ from .grid import (RectilinearZGrid, RectilinearSGrid, CurvilinearZGrid,
                    CurvilinearSGrid, CGrid, GridCode)
 
 
-__all__ = ['Field', 'Geographic', 'GeographicPolar', 'GeographicSquare', 'GeographicPolarSquare']
+__all__ = ['Field', 'VectorField', 'Geographic', 'GeographicPolar', 'GeographicSquare', 'GeographicPolarSquare']
 
 
 class FieldSamplingError(RuntimeError):
@@ -179,8 +179,6 @@ class Field(object):
                  transpose=False, vmin=None, vmax=None, time_origin=None,
                  interp_method='linear', allow_time_extrapolation=None, time_periodic=False, **kwargs):
         self.name = name
-        if self.name == 'UV':
-            return
         self.data = data
         if grid:
             self.grid = grid
@@ -416,8 +414,6 @@ class Field(object):
         return (zonal, meridional)
 
     def __getitem__(self, key):
-        if self.name == 'UV':
-            return self.getUV(*key)
         return self.eval(*key)
 
     def calc_cell_edge_sizes(self):
@@ -839,16 +835,6 @@ class Field(object):
         else:
             return value
 
-    def ccode_evalUV(self, varU, varV, t, x, y, z):
-        # Casting interp_methd to int as easier to pass on in C-code
-        if self.fieldset.U.grid.gtype in [GridCode.RectilinearZGrid, GridCode.RectilinearSGrid]:
-            return "temporal_interpolationUV(%s, %s, %s, %s, U, V, particle->cxi, particle->cyi, particle->czi, particle->cti, &%s, &%s, %s)" \
-                % (x, y, z, t, varU, varV, self.fieldset.U.interp_method.upper())
-        else:
-            return "temporal_interpolationUVrotation(%s, %s, %s, %s, U, V, cosU, sinU, cosV, sinV, particle->cxi, particle->cyi, particle->czi, particle->cti, &%s, &%s, %s)" \
-                % (x, y, z, t,
-                   varU, varV, self.fieldset.U.interp_method.upper())
-
     def ccode_eval(self, var, t, x, y, z):
         # Casting interp_methd to int as easier to pass on in C-code
         return "temporal_interpolation(%s, %s, %s, %s, %s, particle->cxi, particle->cyi, particle->czi, particle->cti, &%s, %s)" \
@@ -945,7 +931,7 @@ class Field(object):
         :param data: if data is not None, the periodic halo will be achieved on data instead of self.data and data will be returned
         """
         dataNone = not isinstance(data, np.ndarray)
-        if self.name == 'UV' or (self.grid.defer_load and dataNone):
+        if self.grid.defer_load and dataNone:
             return
         data = self.data if dataNone else data
         if zonal:
@@ -979,8 +965,6 @@ class Field(object):
 
         :param filename: Basename of the file
         :param varname: Name of the field, to be appended to the filename"""
-        if self.name == 'UV':
-            return
         filepath = str(path.local('%s%s.nc' % (filename, self.name)))
         if varname is None:
             varname = self.name
@@ -1042,6 +1026,55 @@ class Field(object):
         return data
 
 
+class VectorField(object):
+    """Class VectorField stores 2 or 3 fields which defines together a vector field.
+    This enables to interpolate them as one single vector field in the kernels.
+
+    :param name: Name of the vector field
+    :param U: field defining the zonal component
+    :param V: field defining the meridional component
+    :param W: field defining the vertical component (default: None)
+    """
+    def __init__(self, name, U, V, W=None):
+        self.name = name
+        self.U = U
+        self.V = V
+        if W is not None:
+            self.W = W
+            raise NotImplementedError('3D Vector Fields are not implemented yet')
+
+    def eval(self, time, x, y, z):
+        fieldset = self.fieldset
+        U = self.U.eval(time, x, y, z, False)
+        V = self.V.eval(time, x, y, z, False)
+        if self.U.grid.gtype in [GridCode.RectilinearZGrid, GridCode.RectilinearSGrid]:
+            zonal = U
+            meridional = V
+        else:
+            cosU = fieldset.cosU.eval(time, x, y, z, False)
+            sinU = fieldset.sinU.eval(time, x, y, z, False)
+            cosV = fieldset.cosV.eval(time, x, y, z, False)
+            sinV = fieldset.sinV.eval(time, x, y, z, False)
+            zonal = U * cosU - V * sinV
+            meridional = U * sinU + V * cosV
+        zonal = fieldset.U.units.to_target(zonal, x, y, z)
+        meridional = fieldset.V.units.to_target(meridional, x, y, z)
+        return (zonal, meridional)
+
+    def __getitem__(self, key):
+        return self.eval(*key)
+
+    def ccode_eval(self, varU, varV, U, V, t, x, y, z):
+        # Casting interp_methd to int as easier to pass on in C-code
+        if self.fieldset.U.grid.gtype in [GridCode.RectilinearZGrid, GridCode.RectilinearSGrid]:
+            return "temporal_interpolationUV(%s, %s, %s, %s, %s, %s, particle->cxi, particle->cyi, particle->czi, particle->cti, &%s, &%s, %s)" \
+                % (x, y, z, t, U.name, V.name, varU, varV, self.fieldset.U.interp_method.upper())
+        else:
+            return "temporal_interpolationUVrotation(%s, %s, %s, %s, %s, %s, cosU, sinU, cosV, sinV, particle->cxi, particle->cyi, particle->czi, particle->cti, &%s, &%s, %s)" \
+                % (x, y, z, t, U.name, V.name,
+                   varU, varV, self.fieldset.U.interp_method.upper())
+
+
 class NetcdfFileBuffer(object):
     """ Class that encapsulates and manages deferred access to file data. """
 
diff --git a/parcels/fieldset.py b/parcels/fieldset.py
index d3dc0bd2..114ce15b 100644
--- a/parcels/fieldset.py
+++ b/parcels/fieldset.py
@@ -1,4 +1,4 @@
-from parcels.field import Field
+from parcels.field import Field, VectorField
 from parcels.gridset import GridSet
 from parcels.grid import RectilinearZGrid
 from parcels.scripts import compute_curvilinearGrid_rotationAngles
@@ -25,9 +25,7 @@ class FieldSet(object):
             self.add_field(U)
         if V:
             self.add_field(V)
-        UV = Field('UV', None)
-        UV.fieldset = self
-        self.UV = UV
+        self.UV = None
 
         # Add additional fields as attributes
         if fields:
@@ -95,6 +93,14 @@ class FieldSet(object):
         self.gridset.add_grid(field)
         field.fieldset = self
 
+    def add_vector_field(self, vfield):
+        """Add a :class:`parcels.field.VectorField` object to the FieldSet
+
+        :param vfield: :class:`parcels.field.VectorField` object to be added
+        """
+        setattr(self, vfield.name, vfield)
+        vfield.fieldset = self
+
     def check_complete(self):
         assert(self.U), ('U field is not defined')
         assert(self.V), ('V field is not defined')
@@ -109,6 +115,8 @@ class FieldSet(object):
                 if g.defer_load:
                     g.time_full = g.time_full + (g.time_origin - ugrid.time_origin) / np.timedelta64(1, 's')
                 g.time_origin = ugrid.time_origin
+        if getattr(self, 'UV') is None:
+            self.add_vector_field(VectorField('UV', self.U, self.V))
 
     @classmethod
     def from_netcdf(cls, filenames, variables, dimensions, indices=None,
@@ -357,7 +365,7 @@ class FieldSet(object):
             gnew.advanced = False
 
         for fnew in fieldset_new.fields:
-            if fnew.name == 'UV':
+            if isinstance(fnew, VectorField):
                 continue
             f = getattr(self, fnew.name)
             gnew = fnew.grid
@@ -377,14 +385,14 @@ class FieldSet(object):
         for g in self.gridset.grids:
             g.update_status = 'not_updated'
         for f in self.fields:
-            if f.name == 'UV' or not f.grid.defer_load:
+            if isinstance(f, VectorField) or not f.grid.defer_load:
                 continue
             if f.grid.update_status == 'not_updated':
                 nextTime_loc = f.grid.computeTimeChunk(f, time, signdt)
             nextTime = min(nextTime, nextTime_loc) if signdt >= 0 else max(nextTime, nextTime_loc)
 
         for f in self.fields:
-            if f.name == 'UV' or not f.grid.defer_load or f.is_gradient:
+            if isinstance(f, VectorField) or not f.grid.defer_load or f.is_gradient:
                 continue
             g = f.grid
             if g.update_status == 'first_updated':  # First load of data
diff --git a/parcels/kernel.py b/parcels/kernel.py
index 7c2045e6..6f2b812e 100644
--- a/parcels/kernel.py
+++ b/parcels/kernel.py
@@ -89,19 +89,19 @@ class Kernel(object):
         # Generate the kernel function and add the outer loop
         if self.ptype.uses_jit:
             kernelgen = KernelGenerator(fieldset, ptype)
-            self.field_args = kernelgen.field_args
             kernel_ccode = kernelgen.generate(deepcopy(self.py_ast),
                                               self.funcvars)
             self.field_args = kernelgen.field_args
-            if 'UV' in self.field_args:
-                fieldset = self.field_args['UV'].fieldset
-                for f in ['U', 'V', 'cosU', 'sinU', 'cosV', 'sinV']:
-                    if f not in self.field_args:
+            self.vector_field_args = kernelgen.vector_field_args
+            fieldset = self.fieldset
+            for fname in self.vector_field_args:
+                f = getattr(fieldset, fname)
+                for sF in [f.U.name, f.V.name, 'cosU', 'sinU', 'cosV', 'sinV']:
+                    if sF not in self.field_args:
                         try:
-                            self.field_args[f] = getattr(fieldset, f)
+                            self.field_args[sF] = getattr(fieldset, sF)
                         except:
                             continue
-                del self.field_args['UV']
             self.const_args = kernelgen.const_args
             loopgen = LoopGenerator(fieldset, ptype)
             if path.isfile(c_include):


Test patch:
diff --git a/tests/test_fieldset.py b/tests/test_fieldset.py
index c49800e6..dfff590e 100644
--- a/tests/test_fieldset.py
+++ b/tests/test_fieldset.py
@@ -1,5 +1,5 @@
-from parcels import FieldSet, ParticleSet, ScipyParticle, JITParticle, Variable, AdvectionRK4_3D, RectilinearZGrid
-from parcels.field import Field
+from parcels import FieldSet, ParticleSet, ScipyParticle, JITParticle, Variable, AdvectionRK4, AdvectionRK4_3D, RectilinearZGrid
+from parcels.field import Field, VectorField
 from datetime import timedelta as delta
 import datetime
 import numpy as np
@@ -201,6 +201,32 @@ def test_fieldset_constant(mode):
     assert abs(pset[0].lon - (0.5 + westval + eastval)) < 1e-4
 
 
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+@pytest.mark.parametrize('swapUV', [False, True])
+def test_vector_fields(mode, swapUV):
+    lon = np.linspace(0., 10., 12, dtype=np.float32)
+    lat = np.linspace(0., 10., 10, dtype=np.float32)
+    U = np.ones((10, 12), dtype=np.float32)
+    V = np.zeros((10, 12), dtype=np.float32)
+    data = {'U': U, 'V': V}
+    dimensions = {'U': {'lat': lat, 'lon': lon},
+                  'V': {'lat': lat, 'lon': lon}}
+    fieldset = FieldSet.from_data(data, dimensions, mesh='flat')
+    if swapUV:  # we test that we can freely edit whatever UV field
+        UV = VectorField('UV', fieldset.V, fieldset.U)
+        fieldset.add_vector_field(UV)
+
+    pset = ParticleSet.from_line(fieldset, size=1, pclass=ptype[mode],
+                                 start=(0.5, 0.5), finish=(0.5, 0.5))
+    pset.execute(AdvectionRK4, dt=1, runtime=1)
+    if swapUV:
+        assert pset[0].lon == .5
+        assert pset[0].lat == 1.5
+    else:
+        assert pset[0].lon == 1.5
+        assert pset[0].lat == .5
+
+
 @pytest.mark.parametrize('mode', ['scipy', 'jit'])
 @pytest.mark.parametrize('time_periodic', [True, False])
 @pytest.mark.parametrize('dt_sign', [-1, 1])