Making interaction tutorial v4-compatible

Not yet touching the Interaction kernels themselves...

Author: erikvansebille

docs/user_guide/examples_v3/tutorial_interaction.ipynb

@@ -15,9 +15,8 @@
    "id": "1",
    "metadata": {},
    "source": [
-    "In this notebook, we show an example of the new 'particle-particle-interaction' functionality in Parcels. Note that this functionality is still in development, and the implementation is fairly rudimentary and slow for now. Importantly:\n",
+    "In this notebook, we show an example of the new 'particle-particle-interaction' functionality in Parcels. Note that this functionality is still in development, and the implementation is fairly rudimentary for now. Importantly:\n",
     "\n",
-    "- Particle-particle interaction only works in Scipy mode\n",
     "- The type of interactions that are supported is still limited\n",
     "\n",
     "Interactions are implemented through `InteractionKernels`, which are similar to normal `Kernels`. The `InteractionKernels` are applied between particles that are located closer to each other than a specified `interaction_distance`. In general, the code structure needs three adaptations to apply particle-particle interaction:\n",
@@ -54,14 +53,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%matplotlib notebook\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import xarray as xr\n",
-    "from IPython.display import HTML\n",
     "from matplotlib.animation import FuncAnimation\n",
     "\n",
-    "import parcels"
+    "plt.rcParams[\"animation.html\"] = \"jshtml\"\n",
+    "\n",
+    "import parcels\n",
+    "from parcels._datasets.structured.generated import simple_UV_dataset"
    ]
   },
   {
@@ -71,13 +71,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def Pull(particle, fieldset, time, neighbors, mutator):\n",
+    "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 particle.attractor:\n",
+    "    if not particles.attractor:\n",
     "        return StateCode.Success\n",
     "    for n in neighbors:\n",
     "        if n.attractor:\n",
@@ -119,66 +119,84 @@
    "id": "5",
    "metadata": {},
    "outputs": [],
+   "source": [
+    "def DiffusionFieldSet():\n",
+    "    \"\"\"Define a fieldset without flow but with diffusion\"\"\"\n",
+    "    ds = simple_UV_dataset(dims=(1, 1, 1, 1), mesh=\"flat\")\n",
+    "    ds = ds.assign_coords(time=ds.time - ds.time[0])  # set time to timedelta\n",
+    "    ds[\"time\"].attrs.update({\"axis\": \"T\"})\n",
+    "    grid = parcels.XGrid.from_dataset(ds, mesh=\"flat\")\n",
+    "    U = parcels.Field(\"U\", ds[\"U\"], grid, interp_method=parcels.interpolators.XLinear)\n",
+    "    V = parcels.Field(\"V\", ds[\"V\"], grid, interp_method=parcels.interpolators.XLinear)\n",
+    "    UV = parcels.VectorField(\"UV\", U, V)\n",
+    "    fieldset = parcels.FieldSet([U, V, UV])\n",
+    "    fieldset.add_constant_field(\"Kh_zonal\", 0.0005, mesh=\"flat\")\n",
+    "    fieldset.add_constant_field(\"Kh_meridional\", 0.0005, mesh=\"flat\")\n",
+    "    return fieldset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "npart = 11\n",
-    "\n",
     "X, Y = np.meshgrid(np.linspace(-1, 1, npart), np.linspace(-1, 1, npart))\n",
     "\n",
-    "# Define a fieldset without flow\n",
-    "fieldset = parcels.FieldSet.from_data(\n",
-    "    {\"U\": 0, \"V\": 0}, {\"lon\": 0, \"lat\": 0}, mesh=\"flat\"\n",
-    ")\n",
-    "fieldset.add_constant_field(\"Kh_zonal\", 0.0005, mesh=\"flat\")\n",
-    "fieldset.add_constant_field(\"Kh_meridional\", 0.0005, mesh=\"flat\")\n",
-    "\n",
-    "\n",
     "# Create custom particle class with extra variable that indicates\n",
     "# whether the interaction kernel should be executed on this particle.\n",
     "InteractingParticle = parcels.Particle.add_variable(\n",
-    "    \"attractor\", dtype=np.bool_, to_write=\"once\"\n",
+    "    parcels.Variable(\"attractor\", dtype=np.bool_, to_write=\"once\"),\n",
     ")\n",
     "\n",
-    "\n",
     "attractor = [\n",
     "    True if i in [int(npart * npart / 2 - 3), int(npart * npart / 2 + 3)] else False\n",
     "    for i in range(npart * npart)\n",
     "]\n",
+    "\n",
     "pset = parcels.ParticleSet(\n",
-    "    fieldset=fieldset,\n",
+    "    fieldset=DiffusionFieldSet(),\n",
     "    pclass=InteractingParticle,\n",
     "    lon=X,\n",
     "    lat=Y,\n",
-    "    interaction_distance=0.5,  # note the interaction_distance argument here\n",
     "    attractor=attractor,\n",
     ")\n",
     "\n",
-    "output_file = pset.ParticleFile(name=\"InteractingParticles.zarr\", outputdt=1)\n",
+    "output_file = parcels.ParticleFile(\n",
+    "    store=\"InteractingParticles.zarr\",\n",
+    "    outputdt=np.timedelta64(1, \"s\"),\n",
+    ")\n",
+    "\n",
+    "kernels = [\n",
+    "    parcels.kernels.DiffusionUniformKh,\n",
+    "    # Pull,  # Add the Pull kernel defined above\n",
+    "]\n",
     "\n",
     "pset.execute(\n",
-    "    pyfunc=parcels.kernels.DiffusionUniformKh,\n",
-    "    pyfunc_inter=Pull,  # note the pyfunc_inter here\n",
-    "    runtime=60,\n",
-    "    dt=1,\n",
+    "    pyfunc=kernels,\n",
+    "    runtime=np.timedelta64(60, \"s\"),\n",
+    "    dt=np.timedelta64(1, \"s\"),\n",
     "    output_file=output_file,\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6",
+   "id": "7",
    "metadata": {},
    "outputs": [],
    "source": [
-    "%%capture\n",
     "data_xarray = xr.open_zarr(\"InteractingParticles.zarr\")\n",
     "data_attr = data_xarray.where(data_xarray[\"attractor\"].compute() == 1, drop=True)\n",
     "data_other = data_xarray.where(data_xarray[\"attractor\"].compute() == 0, drop=True)\n",
     "\n",
     "timerange = np.arange(\n",
     "    np.nanmin(data_xarray[\"time\"].values),\n",
     "    np.nanmax(data_xarray[\"time\"].values),\n",
-    "    np.timedelta64(1, \"s\"),  # timerange in nanoseconds\n",
+    "    np.timedelta64(1, \"s\"),\n",
     ")\n",
     "\n",
     "fig = plt.figure(figsize=(4, 4), constrained_layout=True)\n",
@@ -189,12 +207,8 @@
     "ax.set_xlim(-1.1, 1.1)\n",
     "ax.set_ylim(-1.1, 1.1)\n",
     "\n",
-    "time_id = np.where(\n",
-    "    data_other[\"time\"] == timerange[0]\n",
-    ")  # Indices of the data where time = 0\n",
-    "time_id_attr = np.where(\n",
-    "    data_attr[\"time\"] == timerange[0]\n",
-    ")  # Indices of the data where time = 0\n",
+    "time_id = np.where(data_other[\"time\"] == timerange[0])\n",
+    "time_id_attr = np.where(data_attr[\"time\"] == timerange[0])\n",
     "\n",
     "scatter = ax.scatter(\n",
     "    data_other[\"lon\"].values[time_id],\n",
@@ -247,25 +261,12 @@
     "        data_attr[\"lat\"].values[time_id_attr],\n",
     "    ):\n",
     "        c.center = (lon_a, lat_a)\n",
-    "    return (\n",
-    "        scatter,\n",
-    "        scatter_attr,\n",
-    "        circs,\n",
-    "    )\n",
     "\n",
     "\n",
-    "anim = FuncAnimation(fig, animate, frames=len(timerange), interval=200, blit=True)\n",
-    "data_xarray.close()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "HTML(anim.to_jshtml())"
+    "# Create animation\n",
+    "anim = FuncAnimation(fig, animate, frames=len(timerange), interval=100)\n",
+    "plt.close(fig)\n",
+    "anim"
    ]
   },
   {
@@ -293,39 +294,36 @@
     "X = np.random.uniform(-1, 1, size=npart)\n",
     "Y = np.random.uniform(-1, 1, size=npart)\n",
     "\n",
-    "# Define a fieldset without flow\n",
-    "fieldset = parcels.FieldSet.from_data(\n",
-    "    {\"U\": 0, \"V\": 0}, {\"lon\": 0, \"lat\": 0}, mesh=\"flat\"\n",
-    ")\n",
-    "fieldset.add_constant_field(\"Kh_zonal\", 0.0005, mesh=\"flat\")\n",
-    "fieldset.add_constant_field(\"Kh_meridional\", 0.0005, mesh=\"flat\")\n",
-    "\n",
-    "\n",
     "# Create custom InteractionParticle class\n",
     "# with extra variables nearest_neighbor and mass\n",
-    "MergeParticle = parcels.InteractionParticle.add_variables(\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",
     ")\n",
     "\n",
     "pset = parcels.ParticleSet(\n",
-    "    fieldset=fieldset,\n",
+    "    fieldset=DiffusionFieldSet(),\n",
     "    pclass=MergeParticle,\n",
     "    lon=X,\n",
     "    lat=Y,\n",
-    "    interaction_distance=0.05,  # note this argument here\n",
     ")\n",
     "\n",
-    "output_file = pset.ParticleFile(name=\"MergingParticles.zarr\", outputdt=1)\n",
+    "output_file = parcels.ParticleFile(\n",
+    "    store=\"MergingParticles.zarr\",\n",
+    "    outputdt=np.timedelta64(1, \"s\"),\n",
+    ")\n",
+    "\n",
+    "kernels = [\n",
+    "    parcels.kernels.DiffusionUniformKh,\n",
+    "    # Merge,  # Add the Merge kernel defined above\n",
+    "]\n",
     "\n",
     "pset.execute(\n",
-    "    pyfunc=parcels.kernels.DiffusionUniformKh,\n",
-    "    pyfunc_inter=pset.InteractionKernel(parcels.NearestNeighborWithinRange)\n",
-    "    + parcels.MergeWithNearestNeighbor,  # note the pyfunc_inter here\n",
-    "    runtime=60,\n",
-    "    dt=1,\n",
+    "    pyfunc=kernels,\n",
+    "    runtime=np.timedelta64(60, \"s\"),\n",
+    "    dt=np.timedelta64(1, \"s\"),\n",
     "    output_file=output_file,\n",
     ")"
    ]
@@ -337,7 +335,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%%capture\n",
     "data_xarray = xr.open_zarr(\"MergingParticles.zarr\")\n",
     "\n",
     "timerange = np.arange(\n",
@@ -354,9 +351,7 @@
     "ax.set_xlim(-1.1, 1.1)\n",
     "ax.set_ylim(-1.1, 1.1)\n",
     "\n",
-    "time_id = np.where(\n",
-    "    data_xarray[\"time\"] == timerange[0]\n",
-    ")  # Indices of the data where time = 0\n",
+    "time_id = np.where(data_xarray[\"time\"] == timerange[0])\n",
     "\n",
     "scatter = ax.scatter(\n",
     "    data_xarray[\"lon\"].values[time_id],\n",
@@ -383,29 +378,9 @@
     "    return (scatter,)\n",
     "\n",
     "\n",
-    "anim = FuncAnimation(fig, animate, frames=len(timerange), interval=200, blit=True)\n",
-    "data_xarray.close()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "11",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "HTML(anim.to_jshtml())"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "id": "12",
-   "metadata": {},
-   "source": [
-    "## Interacting with the FieldSet\n",
-    "\n",
-    "An important feature of Parcels is to evaluate a `Field` at the `Particle` location using the square bracket notation: `particle.Temperature = fieldset.T[time, depth, lat, lon]`. These types of particle-field interactions are recommended to be implemented in standard `Kernels`, since the `InteractionKernels` do not report the `StateCodes` that are used to flag particles that encounter an error in the particle-field interaction, e.g. `OutOfBoundsError`. Any variable that is needed in the particle-particle interaction can be stored in a `Variable` by sampling the field in a `Kernel` before executing the `InteractionKernel`.\n"
+    "anim = FuncAnimation(fig, animate, frames=len(timerange), interval=100)\n",
+    "plt.close(fig)\n",
+    "anim"
    ]
   }
  ],