Add Custom Merge Kernel

Author: erikvansebille

docs/user_guide/examples_v3/tutorial_interaction.ipynb

@@ -241,9 +241,7 @@
    "source": [
     "## Merging particles\n",
     "\n",
-    "Another type of interaction that is supported is the merging of particles. The supported merging functions also comes with limitations (only mutual-nearest particles can be accurately merged), so this is really just a prototype. Nevertheless, the example below shows the possibilities that merging of particles can provide for more complex simulations.\n",
-    "\n",
-    "In the example below, we use two build-in Kernels: `NearestNeighborWithinRange` and `MergeWithNearestNeighbor`.\n"
+    "Another type of interaction is the merging of particles. The merging Kernel also comes with limitations (only mutual-nearest particles can be accurately merged), so this is really just a prototype. Nevertheless, the example below shows the possibilities that merging of particles can provide for more complex simulations."
    ]
   },
   {
@@ -253,25 +251,81 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "npart = 800\n",
+    "def Merge(particles, fieldset):\n",
+    "    \"\"\"Kernel that \"merges\" two particles that are within interaction_distance range\n",
+    "    Particles must have a 'mass' Variable, and during merging the mass of the smallest\n",
+    "    particle is transferred to the largest particle. The smallest particle is then deleted.\n",
+    "    \"\"\"\n",
+    "    interaction_distance = 0.05\n",
+    "\n",
+    "    N = len(particles.lon)\n",
+    "\n",
+    "    # calculate pairwise distances (n_particles × n_particles)\n",
+    "    dx = particles.lon[None, :] - particles.lon[:, None]\n",
+    "    dy = particles.lat[None, :] - particles.lat[:, None]\n",
+    "    distances = np.sqrt(dx**2 + dy**2)\n",
     "\n",
-    "X = np.random.uniform(-1, 1, size=npart)\n",
-    "Y = np.random.uniform(-1, 1, size=npart)\n",
+    "    # mask distances by interaction range\n",
+    "    within = (distances > 0) & (distances < interaction_distance)\n",
+    "    dist_masked = np.where(within, distances, np.inf)\n",
+    "\n",
+    "    # nearest neighbour for each particle (index) and corresponding distance\n",
+    "    nn = np.argmin(dist_masked, axis=1)\n",
+    "    nn_dist = dist_masked[np.arange(N), nn]\n",
+    "    has_nn = np.isfinite(nn_dist)\n",
+    "\n",
+    "    # mutual nearest: nn[nn[i]] == i\n",
+    "    i_idx = np.arange(N)\n",
+    "    mutual = has_nn & (nn[nn] == i_idx)\n",
+    "\n",
+    "    # process each mutual pair only once (i < j)\n",
+    "    pair_mask = mutual & (i_idx < nn)\n",
+    "    pair_i = i_idx[pair_mask]\n",
+    "    pair_j = nn[pair_mask]\n",
+    "\n",
+    "    if pair_i.size == 0:\n",
+    "        return\n",
+    "\n",
+    "    # larger = where mass_j >= mass_i -> j else i\n",
+    "    mass_i = particles.mass[pair_i]\n",
+    "    mass_j = particles.mass[pair_j]\n",
+    "    larger_idx = np.where(mass_j > mass_i, pair_j, pair_i)\n",
+    "    smaller_idx = np.where(mass_j > mass_i, pair_i, pair_j)\n",
+    "\n",
+    "    # perform transfer and mark deletions\n",
+    "    # note that we use temporary arrays for indexing because of KernelParticle bug (GH #2143)\n",
+    "    masses = particles.mass\n",
+    "    states = particles.state\n",
+    "\n",
+    "    # transfer mass from smaller to larger and mark smaller for deletion\n",
+    "    masses[larger_idx] += particles.mass[smaller_idx]\n",
+    "    states[smaller_idx] = parcels.StatusCode.Delete\n",
+    "\n",
+    "    # TODO use particle variables directly KernelParticle bug (GH #2143) is fixed\n",
+    "    particles.mass = masses\n",
+    "    particles.state = states"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "npart = 800\n",
     "\n",
-    "# Create custom InteractionParticle class\n",
-    "# with extra variables nearest_neighbor and mass\n",
+    "# Create custom Particle class with extra variable mass\n",
     "MergeParticle = parcels.Particle.add_variable(\n",
-    "    [\n",
-    "        parcels.Variable(\"nearest_neighbor\", dtype=np.int64, to_write=False),\n",
-    "        parcels.Variable(\"mass\", initial=1, dtype=np.float32),\n",
-    "    ]\n",
+    "    parcels.Variable(\"mass\", dtype=np.float32)\n",
     ")\n",
     "\n",
     "pset = parcels.ParticleSet(\n",
     "    fieldset=DiffusionFieldSet(),\n",
     "    pclass=MergeParticle,\n",
-    "    lon=X,\n",
-    "    lat=Y,\n",
+    "    lon=np.random.uniform(-1, 1, size=npart),\n",
+    "    lat=np.random.uniform(-1, 1, size=npart),\n",
+    "    mass=np.random.uniform(0.5, 1.5, size=npart),\n",
     ")\n",
     "\n",
     "output_file = parcels.ParticleFile(\n",
@@ -281,7 +335,7 @@
     "\n",
     "kernels = [\n",
     "    parcels.kernels.DiffusionUniformKh,\n",
-    "    # Merge,  # Add the Merge kernel defined above\n",
+    "    Merge,  # Add the Merge kernel defined above\n",
     "]\n",
     "\n",
     "pset.execute(\n",
@@ -295,7 +349,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "10",
+   "id": "11",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -305,7 +359,7 @@
     "    np.nanmin(data_xarray[\"time\"].values),\n",
     "    np.nanmax(data_xarray[\"time\"].values),\n",
     "    np.timedelta64(1, \"s\"),\n",
-    ")  # timerange in nanoseconds\n",
+    ")\n",
     "\n",
     "fig = plt.figure(figsize=(4, 4), constrained_layout=True)\n",
     "ax = fig.add_subplot()\n",
@@ -350,7 +404,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "test-latest",
    "language": "python",
    "name": "python3"
   },
@@ -364,7 +418,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6"
+   "version": "3.11.0"
   }
  },
  "nbformat": 4,