Merge pull request #1175 from OceanParcels/fieldset_velocitysampling_fix

Fixing sampling of fieldset.UV

Author: erikvansebille

parcels/application_kernels/advectiondiffusion.py

@@ -33,14 +33,13 @@ def AdvectionDiffusionM1(particle, fieldset, time):
     Kxm1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon - fieldset.dres]
     dKdx = (Kxp1 - Kxm1) / (2 * fieldset.dres)
 
-    u = fieldset.U[time, particle.depth, particle.lat, particle.lon]
+    u, v = fieldset.UV[time, particle.depth, particle.lat, particle.lon]
     bx = math.sqrt(2 * fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon])
 
     Kyp1 = fieldset.Kh_meridional[time, particle.depth, particle.lat + fieldset.dres, particle.lon]
     Kym1 = fieldset.Kh_meridional[time, particle.depth, particle.lat - fieldset.dres, particle.lon]
     dKdy = (Kyp1 - Kym1) / (2 * fieldset.dres)
 
-    v = fieldset.V[time, particle.depth, particle.lat, particle.lon]
     by = math.sqrt(2 * fieldset.Kh_meridional[time, particle.depth, particle.lat, particle.lon])
 
     # Particle positions are updated only after evaluating all terms.
@@ -66,16 +65,18 @@ def AdvectionDiffusionEM(particle, fieldset, time):
     dWx = ParcelsRandom.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
     dWy = ParcelsRandom.normalvariate(0, math.sqrt(math.fabs(particle.dt)))
 
+    u, v = fieldset.UV[time, particle.depth, particle.lat, particle.lon]
+
     Kxp1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon + fieldset.dres]
     Kxm1 = fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon - fieldset.dres]
     dKdx = (Kxp1 - Kxm1) / (2 * fieldset.dres)
-    ax = fieldset.U[time, particle.depth, particle.lat, particle.lon] + dKdx
+    ax = u + dKdx
     bx = math.sqrt(2 * fieldset.Kh_zonal[time, particle.depth, particle.lat, particle.lon])
 
     Kyp1 = fieldset.Kh_meridional[time, particle.depth, particle.lat + fieldset.dres, particle.lon]
     Kym1 = fieldset.Kh_meridional[time, particle.depth, particle.lat - fieldset.dres, particle.lon]
     dKdy = (Kyp1 - Kym1) / (2 * fieldset.dres)
-    ay = fieldset.V[time, particle.depth, particle.lat, particle.lon] + dKdy
+    ay = v + dKdy
     by = math.sqrt(2 * fieldset.Kh_meridional[time, particle.depth, particle.lat, particle.lon])
 
     # Particle positions are updated only after evaluating all terms.

parcels/compilation/codegenerator.py

@@ -78,17 +78,31 @@ def __init__(self, field, args, var, convert=True):
 
 
 class VectorFieldNode(IntrinsicNode):
+    def __getattr__(self, attr):
+        if attr == "eval":
+            return VectorFieldEvalCallNode(self)
+        else:
+            raise NotImplementedError('Access to VectorField attributes are not (yet) implemented in JIT mode')
+
     def __getitem__(self, attr):
         return VectorFieldEvalNode(self.obj, attr)
 
 
+class VectorFieldEvalCallNode(IntrinsicNode):
+    def __init__(self, field):
+        self.field = field
+        self.obj = field.obj
+        self.ccode = ""
+
+
 class VectorFieldEvalNode(IntrinsicNode):
-    def __init__(self, field, args, var, var2, var3):
+    def __init__(self, field, args, var, var2, var3, convert=True):
         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
         self.var3 = var3  # third variable for UVW interpolation
+        self.convert = convert  # whether to convert the result (like field.applyConversion)
 
 
 class SummedFieldNode(IntrinsicNode):
@@ -412,6 +426,26 @@ def visit_Call(self, node):
             self.stmt_stack += [FieldEvalNode(node.func.field, args, tmp, convert)]
             return ast.Name(id=tmp)
 
+        elif isinstance(node.func, VectorFieldEvalCallNode):
+            # get a temporary value to assign result to
+            tmp1 = self.get_tmp()
+            tmp2 = self.get_tmp()
+            tmp3 = self.get_tmp() if node.func.field.obj.vector_type == '3D' else None
+            # whether to convert
+            convert = True
+            if "applyConversion" in node.keywords:
+                k = node.keywords["applyConversion"]
+                if isinstance(k, ast.NameConstant):
+                    convert = k.value
+
+            # convert args to Index(Tuple(*args))
+            args = ast.Index(value=ast.Tuple(node.args, ast.Load()))
+
+            self.stmt_stack += [VectorFieldEvalNode(node.func.field, args, tmp1, tmp2, tmp3, convert)]
+            if tmp3:
+                return ast.Tuple([ast.Name(id=tmp1), ast.Name(id=tmp2), ast.Name(id=tmp3)], ast.Load())
+            else:
+                return ast.Tuple([ast.Name(id=tmp1), ast.Name(id=tmp2)], ast.Load())
         return node
 
 
@@ -907,14 +941,14 @@ def visit_VectorFieldEvalNode(self, node):
         args = self._check_FieldSamplingArguments(node.args.ccode)
         ccode_eval = node.field.obj.ccode_eval_array(node.var, node.var2, node.var3,
                                                      node.field.obj.U, node.field.obj.V, node.field.obj.W, *args)
-        if node.field.obj.U.interp_method != 'cgrid_velocity':
+        if node.convert and node.field.obj.U.interp_method != 'cgrid_velocity':
             ccode_conv1 = node.field.obj.U.ccode_convert(*args)
             ccode_conv2 = node.field.obj.V.ccode_convert(*args)
             statements = [c.Statement("%s *= %s" % (node.var, ccode_conv1)),
                           c.Statement("%s *= %s" % (node.var2, ccode_conv2))]
         else:
             statements = []
-        if node.field.obj.vector_type == '3D':
+        if node.convert and node.field.obj.vector_type == '3D':
             ccode_conv3 = node.field.obj.W.ccode_convert(*args)
             statements.append(c.Statement("%s *= %s" % (node.var3, ccode_conv3)))
         conv_stat = c.Block(statements)
@@ -1058,14 +1092,14 @@ def visit_VectorFieldEvalNode(self, node):
         self.visit(node.args)
         args = self._check_FieldSamplingArguments(node.args.ccode)
         ccode_eval = node.field.obj.ccode_eval_object(node.var, node.var2, node.var3, node.field.obj.U, node.field.obj.V, node.field.obj.W, *args)
-        if node.field.obj.U.interp_method != 'cgrid_velocity':
+        if node.convert and node.field.obj.U.interp_method != 'cgrid_velocity':
             ccode_conv1 = node.field.obj.U.ccode_convert(*args)
             ccode_conv2 = node.field.obj.V.ccode_convert(*args)
             statements = [c.Statement("%s *= %s" % (node.var, ccode_conv1)),
                           c.Statement("%s *= %s" % (node.var2, ccode_conv2))]
         else:
             statements = []
-        if node.field.obj.vector_type == '3D':
+        if node.convert and node.field.obj.vector_type == '3D':
             ccode_conv3 = node.field.obj.W.ccode_convert(*args)
             statements.append(c.Statement("%s *= %s" % (node.var3, ccode_conv3)))
         conv_stat = c.Block(statements)

parcels/examples/example_globcurrent.py

@@ -105,7 +105,8 @@ class MyParticle(ptype[mode]):
         pset = ParticleSet(fieldset, pclass=MyParticle, lon=25, lat=-35, time=fieldset.U.grid.time[0])
 
         def SampleU(particle, fieldset, time):
-            particle.sample_var += fieldset.U[time, particle.depth, particle.lat, particle.lon]
+            u, v = fieldset.UV[time, particle.depth, particle.lat, particle.lon]
+            particle.sample_var += u
 
         pset.execute(SampleU, runtime=delta(days=rundays), dt=delta(days=1))
         sample_var.append(pset[0].sample_var)

parcels/examples/tutorial_diffusion.ipynb

@@ -400,10 +400,16 @@
     "def smagdiff(particle, fieldset, time):\n",
     "    dx = 0.01\n",
     "    # gradients are computed by using a local central difference.\n",
-    "    dudx = (fieldset.U[time, particle.depth, particle.lat, particle.lon+dx]-fieldset.U[time, particle.depth, particle.lat, particle.lon-dx]) / (2*dx)\n",
-    "    dudy = (fieldset.U[time, particle.depth, particle.lat+dx, particle.lon]-fieldset.U[time, particle.depth, particle.lat-dx, particle.lon]) / (2*dx)\n",
-    "    dvdx = (fieldset.V[time, particle.depth, particle.lat, particle.lon+dx]-fieldset.V[time, particle.depth, particle.lat, particle.lon-dx]) / (2*dx)\n",
-    "    dvdy = (fieldset.V[time, particle.depth, particle.lat+dx, particle.lon]-fieldset.V[time, particle.depth, particle.lat-dx, particle.lon]) / (2*dx)\n",
+    "    updx, vpdx = fieldset.UV[time, particle.depth, particle.lat, particle.lon+dx]\n",
+    "    umdx, vmdx = fieldset.UV[time, particle.depth, particle.lat, particle.lon-dx]\n",
+    "    updy, vpdy = fieldset.UV[time, particle.depth, particle.lat+dx, particle.lon]\n",
+    "    umdy, vmdy = fieldset.UV[time, particle.depth, particle.lat-dx, particle.lon]\n",
+    "\n",
+    "    dudx = (updx - umdx) / (2*dx)\n",
+    "    dudy = (updy - umdy) / (2*dx)\n",
+    "    \n",
+    "    dvdx = (vpdx - vmdx) / (2*dx)\n",
+    "    dvdy = (vpdy - vmdy) / (2*dx)\n",
     "\n",
     "    A = fieldset.cell_areas[time, 0, particle.lat, particle.lon]\n",
     "    sq_deg_to_sq_m = (1852*60)**2*math.cos(particle.lat*math.pi/180)\n",
@@ -631,7 +637,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -645,7 +651,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.8"
+   "version": "3.8.13"
   }
  },
  "nbformat": 4,