Vectorising Pull Kernel

Introducing Pull as normal Kernel (but not performance-optimised yet)

Author: erikvansebille

docs/user_guide/examples_v3/tutorial_interaction.ipynb

@@ -58,9 +58,10 @@
     "import xarray as xr\n",
     "from matplotlib.animation import FuncAnimation\n",
     "\n",
-    "plt.rcParams[\"animation.html\"] = \"jshtml\"\n",
+    "import parcels\n",
     "\n",
-    "import parcels"
+    "# for interactive display of animations\n",
+    "plt.rcParams[\"animation.html\"] = \"jshtml\""
    ]
   },
   {
@@ -73,43 +74,28 @@
     "def Pull(particles, fieldset):\n",
     "    \"\"\"InterActionKernel that \"pulls\" all neighbor particles\n",
     "    toward the attracting particle with a constant velocity\"\"\"\n",
-    "    distances = []\n",
-    "    na_neighbors = []\n",
-    "    # only execute kernel for particles that are 'attractor'\n",
-    "    if not particles.attractor:\n",
-    "        return StateCode.Success\n",
-    "    for n in neighbors:\n",
-    "        if n.attractor:\n",
-    "            continue\n",
-    "        x_p = np.array([particle.depth, particle.lat, particle.lon])\n",
-    "        x_n = np.array([n.depth, n.lat, n.lon])\n",
-    "\n",
-    "        # compute distance between attracted and attracting particle\n",
-    "        distances.append(np.linalg.norm(x_p - x_n))\n",
-    "        na_neighbors.append(n)\n",
-    "\n",
-    "    velocity = 0.04  # predefined attracting velocity\n",
-    "    for n in na_neighbors:\n",
-    "        assert n.dt == particle.dt\n",
-    "        dx = np.array(\n",
-    "            [particle.lat - n.lat, particle.lon - n.lon, particle.depth - n.depth]\n",
-    "        )\n",
-    "        dx_norm = np.linalg.norm(dx)\n",
+    "    interaction_distance = 0.5\n",
+    "    velocity = -0.04  # predefined attracting velocity\n",
     "\n",
-    "        # calculate vector of position change\n",
-    "        distance = velocity * n.dt\n",
-    "        d_vec = distance * dx / dx_norm\n",
+    "    lon_a = particles.lon[np.where(particles.attractor)[0]]\n",
+    "    dx = particles.lon - lon_a[:, None]\n",
+    "    lat_a = particles.lat[np.where(particles.attractor)[0]]\n",
+    "    dy = particles.lat - lat_a[:, None]\n",
     "\n",
-    "        # define mutation function for mutator\n",
-    "        def f(n, dlat, dlon, ddepth):\n",
-    "            n.lat_nextloop += (\n",
-    "                dlat  # note that we need to change the locations for the next loop\n",
-    "            )\n",
-    "            n.lon_nextloop += dlon\n",
-    "            n.depth_nextloop += ddepth\n",
+    "    distances = np.sqrt(dx**2 + dy**2)  # TODO use np.linalg.norm?\n",
+    "\n",
+    "    dx = np.where(distances < interaction_distance, dx, 0.0)\n",
+    "    dy = np.where(distances < interaction_distance, dy, 0.0)\n",
+    "\n",
+    "    dx_norm = np.where(distances > 0, dx / distances, 0)\n",
+    "    particles.dlon += (\n",
+    "        np.sum(dx_norm, axis=0) * velocity\n",
+    "    )  # * particles.dt / np.timedelta64(1, \"s\")\n",
     "\n",
-    "        # add mutation to the mutator\n",
-    "        mutator[n.id].append((f, d_vec))"
+    "    dy_norm = np.where(distances > 0, dy / distances, 0)\n",
+    "    particles.dlat += (\n",
+    "        np.sum(dy_norm, axis=0) * velocity\n",
+    "    )  # * particles.dt / np.timedelta64(1, \"s\")"
    ]
   },
   {
@@ -163,7 +149,7 @@
     "\n",
     "kernels = [\n",
     "    parcels.kernels.DiffusionUniformKh,\n",
-    "    # Pull,  # Add the Pull kernel defined above\n",
+    "    Pull,  # Add the Pull kernel defined above\n",
     "]\n",
     "\n",
     "pset.execute(\n",

src/parcels/_core/particleset.py

@@ -261,7 +261,7 @@ def _active_particles_mask(self, time, dt):
         return active_indices
 
     def _compute_neighbor_tree(self, time, dt):
-        active_mask = self._active_particles_mask(time, dt)
+        active_mask = self._active_particles_mask(time, dt)  # TODO still needed with KernelParticles?
 
         self._values = np.vstack(
             (