Merge branch 'main' into spatialhash-regional-spherical-improvements

Author: fluidnumericsJoe

docs/user_guide/additional_examples.md

@@ -57,7 +57,6 @@ import parcels.tutorial
 
 # Load the CopernicusMarine data in the Agulhas region from the example_datasets
 ds_fields = parcels.tutorial.open_dataset("CopernicusMarine_data_for_Argo_tutorial/data")
-ds_fields.load()  # load the dataset into memory
 
 # Create an idealised wind field and add it to the dataset
 tdim, ydim, xdim = (len(ds_fields.time),len(ds_fields.latitude), len(ds_fields.longitude))
@@ -83,6 +82,9 @@ fieldset = parcels.FieldSet.from_sgrid_conventions(
         "Wind": ("UWind", "VWind")
     }
 )
+
+# Convert the FieldSet to windowed arrays for better performance
+fieldset = fieldset.to_windowed_arrays()
 ```
 
 Now we define a wind Kernel that uses a forward Euler method to apply the wind forcing. Note that we update the `particles.dx` and `particles.dy` variables, rather than `particles.x` and `particles.y` directly.

docs/user_guide/examples/explanation_kernelloop.md

@@ -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."
    ]
   },
   {
@@ -205,6 +205,7 @@
     "fieldset_noW = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
     "fieldset_noW.W.data[:] = 0.0\n",
     "fieldset_noW.add_context(\"hc\", ds_fields.hc.item())\n",
+    "fieldset_noW = fieldset_noW.to_windowed_arrays()\n",
     "\n",
     "X, Z = np.meshgrid(\n",
     "    [40e3, 80e3, 120e3],\n",
@@ -221,7 +222,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 +274,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 +298,11 @@
     "z_new = sigma_new * fieldset.h[particles.t, 0, y_new, x_new, particles]\n",
     "```"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {

docs/user_guide/examples/tutorial_croco_3D.ipynb

@@ -56,12 +56,14 @@
     "ds_fields = parcels.tutorial.open_dataset(\n",
     "    \"CopernicusMarine_data_for_Argo_tutorial/data\"\n",
     ")\n",
-    "ds_fields.load()  # load the dataset into memory\n",
     "\n",
     "# Convert to SGRID-compliant dataset and create FieldSet\n",
     "fields = {\"U\": ds_fields[\"uo\"], \"V\": ds_fields[\"vo\"]}\n",
     "ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields)\n",
-    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)"
+    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
+    "\n",
+    "# Convert the FieldSet to windowed arrays for better performance\n",
+    "fieldset = fieldset.to_windowed_arrays()"
    ]
   },
   {
@@ -398,7 +400,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Parcels:docs (3.14.6)",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },

docs/user_guide/examples/tutorial_delaystart.ipynb

@@ -472,8 +472,7 @@
     "# Load the CopernicusMarine data in the Agulhas region from the example_datasets\n",
     "ds_fields = parcels.tutorial.open_dataset(\n",
     "    \"CopernicusMarine_data_for_Argo_tutorial/data\"\n",
-    ").isel(depth=slice(0, 1))\n",
-    "ds_fields.load()  # load the dataset into memory"
+    ").isel(depth=slice(0, 1))"
    ]
   },
   {
@@ -528,7 +527,10 @@
     "ds_cell_areas[\"cell_areas\"] = ([\"lat\", \"lon\"], calc_cell_areas(ds_fields))\n",
     "fset2 = parcels.FieldSet.from_sgrid_conventions(ds_cell_areas)\n",
     "\n",
-    "fieldset += fset2"
+    "fieldset += fset2\n",
+    "\n",
+    "# Convert the FieldSet to windowed arrays for better performance\n",
+    "fieldset = fieldset.to_windowed_arrays()"
    ]
   },
   {

docs/user_guide/examples/tutorial_diffusion.ipynb

@@ -67,7 +67,6 @@
     "ds_fields = parcels.tutorial.open_dataset(\n",
     "    \"CopernicusMarine_data_for_Argo_tutorial/data\"\n",
     ")\n",
-    "ds_fields.load()  # load the dataset into memory\n",
     "\n",
     "OUTPUT_FOLDER = Path(\"output\")\n",
     "OUTPUT_FOLDER.mkdir(exist_ok=True)\n",
@@ -77,6 +76,9 @@
     "ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields)\n",
     "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
     "\n",
+    "# Convert the FieldSet to windowed arrays for better performance\n",
+    "fieldset = fieldset.to_windowed_arrays()\n",
+    "\n",
     "# Check field resolution in time and space\n",
     "print(\n",
     "    f\"dt = {np.unique(np.diff(ds_fields.time.values).astype('timedelta64[h]'))} hours\"\n",

docs/user_guide/examples/tutorial_dt_integrators.ipynb

@@ -43,9 +43,6 @@
     "    \"CopernicusMarine_data_for_Argo_tutorial/data\"\n",
     ")\n",
     "\n",
-    "# TODO check how we can get good performance without loading full dataset in memory\n",
-    "ds_fields.load()  # load the dataset into memory\n",
-    "\n",
     "# Convert to SGRID-compliant dataset and create FieldSet\n",
     "fields = {\n",
     "    \"U\": ds_fields[\"uo\"],\n",
@@ -54,7 +51,10 @@
     "    \"so\": ds_fields[\"so\"],\n",
     "}\n",
     "ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields)\n",
-    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)"
+    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
+    "\n",
+    "# Convert the FieldSet to windowed arrays for better performance\n",
+    "fieldset = fieldset.to_windowed_arrays()"
    ]
   },
   {

docs/user_guide/examples/tutorial_gsw_density.ipynb

@@ -49,7 +49,6 @@
     "ds_fields = parcels.tutorial.open_dataset(\n",
     "    \"CopernicusMarine_data_for_Argo_tutorial/data\"\n",
     ")\n",
-    "ds_fields.load()  # load the dataset into memory\n",
     "\n",
     "# Create an idealised wind field and add it to the dataset\n",
     "ydim, xdim = len(ds_fields.latitude), len(ds_fields.longitude)\n",
@@ -122,6 +121,9 @@
     "ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields)\n",
     "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
     "\n",
+    "# Convert the FieldSet to windowed arrays for better performance\n",
+    "fieldset = fieldset.to_windowed_arrays()\n",
+    "\n",
     "npart = 10\n",
     "z = np.repeat(ds_fields.depth[0].values, npart)\n",
     "lons = np.repeat(32.2, npart)\n",

docs/user_guide/examples/tutorial_manipulating_field_data.ipynb

@@ -307,7 +307,10 @@
     "    fields={\"U\": ds_u[\"uo\"], \"V\": ds_v[\"vo\"], \"W\": ds_w[\"wo\"]},\n",
     "    coords=ds_coords,\n",
     ")\n",
-    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)"
+    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
+    "\n",
+    "fieldset = fieldset.to_windowed_arrays()\n",
+    "fieldset.describe()"
    ]
   },
   {
@@ -327,12 +330,12 @@
    },
    "outputs": [],
    "source": [
-    "npart = 10\n",
+    "npart = 3\n",
     "pset = parcels.ParticleSet(\n",
     "    fieldset=fieldset,\n",
     "    x=np.linspace(1.9, 3.4, npart),\n",
-    "    y=np.linspace(65, 51.6, npart),\n",
-    "    z=100 * np.ones(npart),\n",
+    "    y=np.linspace(60, 58, npart),\n",
+    "    z=20 * np.ones(npart),\n",
     ")\n",
     "\n",
     "pfile = parcels.ParticleFile(\"output_nemo3D.parquet\", outputdt=np.timedelta64(1, \"D\"))\n",
@@ -358,16 +361,30 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "field = fieldset.U.data[0, 0, :, :]\n",
+    "fig, ax = plt.subplots(figsize=(8, 4))\n",
+    "\n",
+    "field = fieldset.U.data.isel(time=0, depth_center=0).compute()\n",
     "field = field.where(field != 0, np.nan)  # Mask land values for better plotting\n",
-    "plt.pcolormesh(fieldset.U.grid.lon, fieldset.U.grid.lat, field, cmap=\"RdBu\")\n",
+    "ax.pcolormesh(fieldset.U.grid.lon, fieldset.U.grid.lat, field, cmap=\"RdBu\")\n",
     "\n",
     "df = parcels.read_particlefile(\"output_nemo3D.parquet\")\n",
     "vmin = df[\"z\"].min()\n",
     "vmax = df[\"z\"].max()\n",
     "for traj in df.partition_by(\"particle_id\", maintain_order=True):\n",
-    "    plt.scatter(traj[\"x\"], traj[\"y\"], c=-traj[\"z\"], marker=\".\", vmin=-vmax, vmax=-vmin)\n",
-    "plt.colorbar(label=\"Depth (m)\")\n",
+    "    ax.plot(traj[\"x\"], traj[\"y\"], \"-\", color=\"k\", alpha=0.5)\n",
+    "    scat = ax.scatter(\n",
+    "        traj[\"x\"],\n",
+    "        traj[\"y\"],\n",
+    "        c=-traj[\"z\"],\n",
+    "        marker=\".\",\n",
+    "        vmin=-vmax,\n",
+    "        vmax=-vmin,\n",
+    "        zorder=10,\n",
+    "    )\n",
+    "plt.colorbar(scat, label=\"Depth (m)\", ax=ax)\n",
+    "ax.set_xlim([1, 9])\n",
+    "ax.set_ylim([56, 61])\n",
+    "ax.set_aspect(\"equal\")\n",
     "plt.show()"
    ]
   }

docs/user_guide/examples/tutorial_nemo.ipynb

@@ -60,12 +60,14 @@
     "ds_fields = parcels.tutorial.open_dataset(\n",
     "    \"CopernicusMarine_data_for_Argo_tutorial/data\"\n",
     ")\n",
-    "ds_fields.load()  # load the dataset into memory\n",
     "\n",
     "# Convert to SGRID-compliant dataset and create FieldSet\n",
     "fields = {\"U\": ds_fields[\"uo\"], \"V\": ds_fields[\"vo\"], \"thetao\": ds_fields[\"thetao\"]}\n",
     "ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields)\n",
-    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)"
+    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
+    "\n",
+    "# Convert the FieldSet to windowed arrays for better performance\n",
+    "fieldset = fieldset.to_windowed_arrays()"
    ]
   },
   {