Use RK2 for CROCO advection

Author: erikvansebille

docs/user_guide/examples/tutorial_croco_3D.ipynb

@@ -37,14 +37,11 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import polars as pl\n",
-    "import xarray as xr\n",
     "\n",
     "import parcels\n",
     "import parcels.tutorial\n",
     "\n",
-    "ds_fields = parcels.tutorial.open_dataset(\"CROCOidealized_data/data\")\n",
-    "\n",
-    "ds_fields.load();  # Preload data to speed up access"
+    "ds_fields = parcels.tutorial.open_dataset(\"CROCOidealized_data/data\")"
    ]
   },
   {
@@ -92,7 +89,10 @@
     "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
     "\n",
     "# Add the critical depth (`hc`) as context to the fieldset\n",
-    "fieldset.add_context(\"hc\", ds_fields.hc.item())"
+    "fieldset.add_context(\"hc\", ds_fields.hc.item())\n",
+    "\n",
+    "# Convert the FieldSet to windowed arrays for better performance\n",
+    "fieldset = fieldset.to_windowed_arrays()"
    ]
   },
   {
@@ -132,7 +132,7 @@
     ")\n",
     "\n",
     "pset.execute(\n",
-    "    [parcels.kernels.AdvectionRK4_3D_CROCO, DeleteParticle],\n",
+    "    [parcels.kernels.AdvectionRK2_3D_CROCO, DeleteParticle],\n",
     "    runtime=np.timedelta64(50_000, \"s\"),\n",
     "    dt=np.timedelta64(100, \"s\"),\n",
     "    output_file=outputfile,\n",
@@ -173,7 +173,7 @@
     "ax.set_xlabel(\"X [km]\")\n",
     "ax.set_xlim(30, 170)\n",
     "ax.set_ylabel(\"Depth [m]\")\n",
-    "ax.set_title(\"Particles in idealized CROCO velocity field using AdvectionRK4_3D_CROCO\")\n",
+    "ax.set_title(\"Particles in idealized CROCO velocity field using AdvectionRK2_3D_CROCO\")\n",
     "plt.tight_layout()\n",
     "plt.show()"
    ]
@@ -189,7 +189,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "It may be insightful to compare this 3D run with the `AdvectionRK4_3D_CROCO` kernel with a run where the vertical velocity (`W`) is set to zero. In that case, the particles will not stay on their initial depth levels but instead follow sigma-layers."
+    "It may be insightful to compare this 3D run with the `AdvectionRK2_3D_CROCO` kernel with a run where the vertical velocity (`W`) is set to zero. In that case, the particles will not stay on their initial depth levels but instead follow sigma-layers."
    ]
   },
   {
@@ -221,7 +221,7 @@
     ")\n",
     "\n",
     "pset_noW.execute(\n",
-    "    [parcels.kernels.AdvectionRK4_3D_CROCO, DeleteParticle],\n",
+    "    [parcels.kernels.AdvectionRK2_3D_CROCO, DeleteParticle],\n",
     "    runtime=np.timedelta64(50_000, \"s\"),\n",
     "    dt=np.timedelta64(100, \"s\"),\n",
     "    output_file=outputfile,\n",
@@ -273,9 +273,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "For Advection, you will need to use the `AdvectionRK4_3D_CROCO` kernel, which works slightly different from the normal 3D advection kernel because it converts the vertical velocity in sigma-units. The conversion from z to sigma is done at every time step, using the `convert_z_to_sigma_croco()` function.\n",
+    "For Advection, you will need to use the `AdvectionRK2_3D_CROCO` kernel, which works slightly different from the normal 3D advection kernel because it converts the vertical velocity in sigma-units. The conversion from z to sigma is done at every time step, using the `convert_z_to_sigma_croco()` function.\n",
     "\n",
-    "In particular, the following algorithm is used (note that the RK4 version is slightly more complex than this Euler-Forward version, but the idea is identical). Also note that the vertical velocity is linearly interpolated here, which gives much better results than the default C-grid interpolation.\n",
+    "In particular, the following algorithm is used (note that the RK2 version is slightly more complex than this Euler-Forward version, but the idea is identical). Also note that the vertical velocity is linearly interpolated here, which gives much better results than the default C-grid interpolation.\n",
     "\n",
     "```python\n",
     "sigma = particles.z / fieldset.h[particles.t, 0, particles.y, particles.x]\n",
@@ -297,6 +297,11 @@
     "z_new = sigma_new * fieldset.h[particles.t, 0, y_new, x_new, particles]\n",
     "```"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {

src/parcels/kernels/__init__.py

@@ -13,7 +13,7 @@
     DiffusionUniformKh,
 )
 from ._sigmagrids import (
-    AdvectionRK4_3D_CROCO,
+    AdvectionRK2_3D_CROCO,
     SampleOmegaCroco,
     convert_z_to_sigma_croco,
 )
@@ -32,7 +32,7 @@
     "AdvectionDiffusionM1",
     "DiffusionUniformKh",
     # sigmagrids
-    "AdvectionRK4_3D_CROCO",
+    "AdvectionRK2_3D_CROCO",
     "SampleOmegaCroco",
     "convert_z_to_sigma_croco",
 ]

src/parcels/kernels/_sigmagrids.py

@@ -35,9 +35,8 @@ def SampleOmegaCroco(particles, fieldset):
     particles.omega = fieldset.omega[particles.t, sigma, particles.y, particles.x, particles]
 
 
-# TODO change to RK2 (once RK4 yields same results as v3)
-def AdvectionRK4_3D_CROCO(particles, fieldset):  # pragma: no cover
-    """Advection of particles using fourth-order Runge-Kutta integration including vertical velocity.
+def AdvectionRK2_3D_CROCO(particles, fieldset):  # pragma: no cover
+    """Advection of particles using second-order Runge-Kutta integration including vertical velocity.
     This kernel assumes the vertical velocity is the 'w' field from CROCO output and works on sigma-layers.
     It also uses linear interpolation of the W field, which gives much better results than the default C-grid interpolation.
     """
@@ -57,31 +56,7 @@ def AdvectionRK4_3D_CROCO(particles, fieldset):  # pragma: no cover
     (u2, v2) = fieldset.UV[particles.t + 0.5 * dt, sig1, y1, x1, particles]
     w2 = fieldset.W[particles.t + 0.5 * dt, sig1, y1, x1, particles]
     w2 *= sig_dep1 / fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y1, x1]
-    x2 = particles.x + u2 * 0.5 * dt
-    y2 = particles.y + v2 * 0.5 * dt
-    sig_dep2 = sigma + w2 * 0.5 * dt
-    dep2 = sig_dep2 * fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y2, x2]
 
-    sig2 = convert_z_to_sigma_croco(fieldset, particles.t + 0.5 * dt, dep2, y2, x2, particles)
-    (u3, v3) = fieldset.UV[particles.t + 0.5 * dt, sig2, y2, x2, particles]
-    w3 = fieldset.W[particles.t + 0.5 * dt, sig2, y2, x2, particles]
-    w3 *= sig_dep2 / fieldset.h[particles.t + 0.5 * dt, np.zeros_like(particles.z), y2, x2]
-    x3 = particles.x + u3 * dt
-    y3 = particles.y + v3 * dt
-    sig_dep3 = sigma + w3 * dt
-    dep3 = sig_dep3 * fieldset.h[particles.t + dt, np.zeros_like(particles.z), y3, x3]
-
-    sig3 = convert_z_to_sigma_croco(fieldset, particles.t + dt, dep3, y3, x3, particles)
-    (u4, v4) = fieldset.UV[particles.t + dt, sig3, y3, x3, particles]
-    w4 = fieldset.W[particles.t + dt, sig3, y3, x3, particles]
-    w4 *= sig_dep3 / fieldset.h[particles.t + dt, np.zeros_like(particles.z), y3, x3]
-    x4 = particles.x + u4 * dt
-    y4 = particles.y + v4 * dt
-    sig_dep4 = sigma + w4 * dt
-
-    dep4 = sig_dep4 * fieldset.h[particles.t + dt, np.zeros_like(particles.z), y4, x4]
-    particles.dx += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt
-    particles.dy += (v1 + 2 * v2 + 2 * v3 + v4) / 6 * dt
-    particles.dz += (
-        (dep1 - particles.z) * 2 + 2 * (dep2 - particles.z) * 2 + 2 * (dep3 - particles.z) + dep4 - particles.z
-    ) / 6
+    particles.dx += u2 * dt
+    particles.dy += v2 * dt
+    particles.dz += w2 * dt

tests/test_sigmagrids.py

@@ -3,7 +3,7 @@
 import parcels
 import parcels.tutorial
 from parcels import Particle, ParticleSet, Variable
-from parcels.kernels import AdvectionRK4_3D_CROCO, SampleOmegaCroco, convert_z_to_sigma_croco
+from parcels.kernels import AdvectionRK2_3D_CROCO, SampleOmegaCroco, convert_z_to_sigma_croco
 
 
 def test_conversion_3DCROCO():
@@ -71,7 +71,7 @@ def test_advection_3DCROCO():
     pset = ParticleSet(fieldset=fieldset, pclass=pclass, x=X, y=Y, z=Z)
 
     pset.execute(
-        [AdvectionRK4_3D_CROCO, SampleOmegaCroco], runtime=np.timedelta64(runtime, "s"), dt=np.timedelta64(100, "s")
+        [AdvectionRK2_3D_CROCO, SampleOmegaCroco], runtime=np.timedelta64(runtime, "s"), dt=np.timedelta64(100, "s")
     )
     np.testing.assert_allclose(pset.z, Z.flatten(), atol=5)  # TODO lower this atol
     np.testing.assert_allclose(pset.x, [x + runtime for x in X.flatten()], atol=1e-3)