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 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 @@
  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 @@
  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
+ )