Analytical Advection in jit

docs/examples/example_nemo_curvilinear.py

@@ -34,6 +34,34 @@ def run_nemo_curvilinear(mode, outfile, advtype="RK4"):
     fieldset = parcels.FieldSet.from_nemo(
         filenames, variables, dimensions, chunksize=chunksize
     )
+    if advtype == "AA":
+        mesh_mask = f"{data_folder}/mesh_mask_edges.nc4"
+        fieldset.add_field(
+            parcels.Field.from_netcdf(
+                mesh_mask, "e2u", dimensions, interp_method="nearest"
+            )
+        )
+        fieldset.add_field(
+            parcels.Field.from_netcdf(
+                mesh_mask, "e1v", dimensions, interp_method="nearest"
+            )
+        )
+        fieldset.add_field(
+            parcels.Field.from_netcdf(
+                mesh_mask, "e1t", dimensions, interp_method="nearest"
+            )
+        )
+        fieldset.add_field(
+            parcels.Field.from_netcdf(
+                mesh_mask, "e2t", dimensions, interp_method="nearest"
+            )
+        )
+        fieldset.add_field(
+            parcels.Field(
+                "e3t", np.array(1), lon=0, lat=0, depth=0, interp_method="nearest"
+            )
+        )
+
     assert fieldset.U.chunksize == chunksize
 
     # Now run particles as normal

docs/examples/example_peninsula.py

@@ -199,13 +199,40 @@ def test_peninsula_fieldset(mode, mesh, tmpdir):
     assert (err_smpl <= 1.0e-3).all()
 
 
-@pytest.mark.parametrize(
-    "mode", ["scipy"]
-)  # Analytical Advection only implemented in Scipy mode
-@pytest.mark.parametrize("mesh", ["flat", "spherical"])
-def test_peninsula_fieldset_AnalyticalAdvection(mode, mesh, tmpdir):
+@pytest.mark.parametrize("mode", ["jit", "scipy"])
+def test_peninsula_fieldset_AnalyticalAdvection(mode, tmpdir):
     """Execute peninsula test using Analytical Advection on C grid."""
     fieldset = peninsula_fieldset(101, 51, "flat", grid_type="C")
+    lon = fieldset.U.lon
+    lat = fieldset.U.lat
+    dx = lon[1] - lon[0]
+    dy = lat[1] - lat[0]
+    fieldset.add_field(
+        parcels.Field(
+            "e2u", dy * np.ones((len(lat), len(lon))), lon, lat, interp_method="nearest"
+        )
+    )
+    fieldset.add_field(
+        parcels.Field(
+            "e1v", dx * np.ones((len(lat), len(lon))), lon, lat, interp_method="nearest"
+        )
+    )
+    fieldset.add_field(
+        parcels.Field(
+            "e1t", dx * np.ones((len(lat), len(lon))), lon, lat, interp_method="nearest"
+        )
+    )
+    fieldset.add_field(
+        parcels.Field(
+            "e2t", dy * np.ones((len(lat), len(lon))), lon, lat, interp_method="nearest"
+        )
+    )
+    fieldset.add_field(
+        parcels.Field(
+            "e3t", np.array(1), lon=0, lat=0, depth=0, interp_method="nearest"
+        )
+    )
+
     outfile = tmpdir.join("PeninsulaAA")
     pset = peninsula_example(
         fieldset, outfile, npart=10, mode=mode, method=parcels.AdvectionAnalytical

parcels/application_kernels/advection.py

@@ -110,7 +110,8 @@ def AdvectionRK45(particle, fieldset, time):
 
 
 def AdvectionAnalytical(particle, fieldset, time):
-    """Advection of particles using 'analytical advection' integration.
+    """Advection of particles using 'analytical advection' integration. This code only works in SciPy mode.
+    The JIT version of this code is hardcoded in C in advectionanalytical.h, and is called by the JIT version of this kernel.
 
     Based on Ariane/TRACMASS algorithm, as detailed in e.g. Doos et al (https://doi.org/10.5194/gmd-10-1733-2017).
     Note that the time-dependent scheme is currently implemented with 'intermediate timesteps'
@@ -136,27 +137,53 @@ def AdvectionAnalytical(particle, fieldset, time):
         particle.lon, particle.lat, particle.depth, particle=particle
     )
     if withW:
+        updatexi = 0
+        updateyi = 0
         if abs(xsi - 1) < tol:
-            if fieldset.U.data[0, zi + 1, yi + 1, xi + 1] > 0:
-                xi += 1
+            if fieldset.U.data[0, zi + 1, yi + 1, xi + 1] > 0:  # TODO update scipy code here to match JIT kernel
+                updatexi = 1
                 xsi = 0
+        elif abs(xsi) < tol:
+            if fieldset.U.data[0, zi, yi, xi] < 0:
+                updatexi = -1
+                xsi = 1
         if abs(eta - 1) < tol:
             if fieldset.V.data[0, zi + 1, yi + 1, xi + 1] > 0:
-                yi += 1
+                updateyi = 1
                 eta = 0
+        elif abs(eta) < tol:
+            if fieldset.V.data[0, zi, yi, xi] < 0:
+                yi -= 1
+                eta = 1
         if abs(zeta - 1) < tol:
             if fieldset.W.data[0, zi + 1, yi + 1, xi + 1] > 0:
                 zi += 1
                 zeta = 0
+        if abs(zeta) < tol:
+            if fieldset.W.data[0, zi, yi, xi] < 0:
+                zi -= 1
+                zeta = 1
+        xi += updatexi
+        yi += updateyi
     else:
+        updatexi = 0
         if abs(xsi - 1) < tol:
             if fieldset.U.data[0, yi + 1, xi + 1] > 0:
-                xi += 1
+                updatexi = 1
                 xsi = 0
+        elif abs(xsi) < tol:
+            if fieldset.U.data[0, yi, xi] < 0:
+                updatexi = -1
+                xsi = 1
         if abs(eta - 1) < tol:
             if fieldset.V.data[0, yi + 1, xi + 1] > 0:
                 yi += 1
                 eta = 0
+        elif abs(eta) < tol:
+            if fieldset.V.data[0, yi, xi] < 0:
+                yi -= 1
+                eta = 1
+        xi += updatexi
 
     particle.xi[:] = xi
     particle.yi[:] = yi
@@ -284,3 +311,69 @@ def compute_rs(r, B, delta, s_min):
         particle.dt = max(direction * s_min * (dxdy * dz), 1e-7)
     else:
         particle.dt = min(direction * s_min * (dxdy * dz), -1e-7)
+
+
+def AdvectionAnalytical_2D_JIT(particle, fieldset, time):
+    """JIT version of 2D AdvectionAnalytical kernel.
+
+    The code itself is in advectionanalytical.h and is automatically included
+    when using AdvectionAnalytical in JIT mode.
+    """
+    flow3D = 0
+    calcAdvectionAnalytical_JIT(  # noqa
+        "parcels_customed_Cfunc_pointer_args",
+        fieldset.U,
+        fieldset.V,
+        fieldset.V,  # Note repeating V field here; W is not used in 2D
+        flow3D,
+        fieldset.e2u,
+        fieldset.e1v,
+        fieldset.e1t,
+        fieldset.e2t,
+        fieldset.e3t,
+        particle.xi,
+        particle.yi,
+        particle.zi,
+        particle.ti,
+        particle.lon,
+        particle.lat,
+        particle.depth,
+        time,
+        particle.dt,
+        particle_dlon,  # noqa
+        particle_dlat,  # noqa
+        particle_ddepth,  # noqa
+    )
+
+
+def AdvectionAnalytical_3D_JIT(particle, fieldset, time):
+    """JIT version of 3D AdvectionAnalytical kernel.
+
+    The code itself is in advectionanalytical.h and is automatically included
+    when using AdvectionAnalytical in JIT mode.
+    """
+    flow3D = 1
+    calcAdvectionAnalytical_JIT(  # noqa
+        "parcels_customed_Cfunc_pointer_args",
+        fieldset.U,
+        fieldset.V,
+        fieldset.W,
+        flow3D,
+        fieldset.e2u,
+        fieldset.e1v,
+        fieldset.e1t,
+        fieldset.e2t,
+        fieldset.e3t,
+        particle.xi,
+        particle.yi,
+        particle.zi,
+        particle.ti,
+        particle.lon,
+        particle.lat,
+        particle.depth,
+        time,
+        particle.dt,
+        particle_dlon,  # noqa
+        particle_dlat,  # noqa
+        particle_ddepth,  # noqa
+    )