Merge remote-tracking branch 'upstream/main' into pr/maureenjcohen/2739

Author: VeckoTheGecko

docs/development/unstructured_grid_search.md

@@ -42,12 +42,12 @@ The algorithm has two phases: **initialisation** (once, when the grid is first u
 
 The hash grid is three-dimensional regardless of the source grid:
 
-| Source grid | Mesh type | Hash-grid space                       |
-| ----------- | --------- | ------------------------------------- |
-| `XGrid`     | spherical | Cartesian unit cube (lon/lat → x,y,z) |
-| `XGrid`     | flat      | 2-D lon/lat bounding box (z set to 0) |
-| `UxGrid`    | spherical | Cartesian unit cube                   |
-| `UxGrid`    | flat      | 2-D lon/lat bounding box              |
+| Source grid | Mesh type | Hash-grid space                          |
+| ----------- | --------- | ---------------------------------------- |
+| `XGrid`     | spherical | Cartesian bounding box (lon/lat → x,y,z) |
+| `XGrid`     | flat      | 2-D lon/lat bounding box (z set to 0)    |
+| `UxGrid`    | spherical | Cartesian bounding box (lon/lat → x,y,z) |
+| `UxGrid`    | flat      | 2-D lon/lat bounding box                 |
 
 For spherical grids, node coordinates are converted from degrees to radians and then to Cartesian (x, y, z) on the unit sphere using `parcels._core.index_search._latlon_rad_to_xyz`:
 
@@ -57,7 +57,7 @@ y = sin(lon) * cos(lat)
 z = sin(lat)
 ```
 
-The hash grid then spans the unit cube `[-1, 1]³`. Working in Cartesian space avoids the longitude wrap-around discontinuity that would otherwise cause the bounding boxes of cells crossing the antimeridian to erroneously span the entire domain.
+The hash grid then spans the axis-aligned Cartesian bounding box of the transformed grid nodes. For a global grid this is (nearly) the unit cube `[-1, 1]³`; for a regional grid it is much tighter, so the full quantization resolution (1024 bins per axis) is spent on the region actually covered by the grid rather than the whole sphere. Working in Cartesian space avoids the longitude wrap-around discontinuity that would otherwise cause the bounding boxes of cells crossing the antimeridian to erroneously span the entire domain.
 
 For flat meshes the hash grid simply spans `[lon_min, lon_max] × [lat_min, lat_max]`, with z fixed at 0.
 

docs/user_guide/additional_examples.md

@@ -57,17 +57,16 @@ 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))
+ydim, xdim = len(ds_fields.latitude), len(ds_fields.longitude)
 ds_fields["UWind"] = xr.DataArray(
-    data=np.ones((tdim, ydim, xdim)) * np.sin(ds_fields.latitude.values)[None, :, None],
-    coords=[ds_fields.time, ds_fields.latitude, ds_fields.longitude])
+    data=0.5 * np.ones((ydim, xdim)) * np.sin(ds_fields.latitude.values - ds_fields.latitude.values.mean())[:, None],
+    coords=[ds_fields.latitude, ds_fields.longitude])
 
 ds_fields["VWind"] = xr.DataArray(
-    data=np.zeros((tdim, ydim, xdim)),
-    coords=[ds_fields.time, ds_fields.latitude, ds_fields.longitude])
+    data=np.zeros((ydim, xdim)),
+    coords=[ds_fields.latitude, ds_fields.longitude])
 
 fields = {
     "U": ds_fields["uo"],
@@ -83,15 +82,20 @@ 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.
 
 ```{code-cell}
-def wind_kernel(particles, fieldset):
-    uwind, vwind = fieldset.Wind[particles]
-    particles.dx += uwind * particles.dt
-    particles.dy += vwind * particles.dt
+def windRK2(particles, fieldset):
+    (u1, v1) = fieldset.Wind[particles]
+    x1, y1 = (particles.x + u1 * 0.5 * particles.dt, particles.y + v1 * 0.5 * particles.dt)
+    (u2, v2) = fieldset.Wind[particles.t + 0.5 * particles.dt, particles.z, y1, x1, particles]
+    particles.dx += u2 * particles.dt
+    particles.dy += v2 * particles.dt
 ```
 
 First run a simulation where we apply Kernels as `[AdvectionRK2, wind_kernel]`
@@ -100,22 +104,22 @@ First run a simulation where we apply Kernels as `[AdvectionRK2, wind_kernel]`
 :tags: [hide-output]
 npart = 10
 z = np.repeat(ds_fields.depth[0].values, npart)
-lons = np.repeat(31, npart)
+lons = np.repeat(32.2, npart)
 lats = np.linspace(-32.5, -30.5, npart)
 
 pset = parcels.ParticleSet(fieldset, pclass=parcels.Particle, z=z, y=lats, x=lons)
 output_file = parcels.ParticleFile(
     path="advection_then_wind.parquet", outputdt=np.timedelta64(6,'h')
 )
 pset.execute(
-    [parcels.kernels.AdvectionRK2, wind_kernel],
+    [parcels.kernels.AdvectionRK2, windRK2],
     runtime=np.timedelta64(5,'D'),
     dt=np.timedelta64(1,'h'),
     output_file=output_file,
 )
 ```
 
-Then also run a simulation where we apply the Kernels in the reverse order as `[wind_kernel, AdvectionRK2]`
+Then also run a simulation where we apply the Kernels in the reverse order as `[windRK2, AdvectionRK2]`
 
 ```{code-cell}
 :tags: [hide-output]
@@ -126,7 +130,7 @@ output_file_reverse = parcels.ParticleFile(
     path="wind_then_advection.parquet", outputdt=np.timedelta64(6,"h")
 )
 pset_reverse.execute(
-    [wind_kernel, parcels.kernels.AdvectionRK2],
+    [windRK2, parcels.kernels.AdvectionRK2],
     runtime=np.timedelta64(5,"D"),
     dt=np.timedelta64(1,"h"),
     output_file=output_file_reverse,

docs/user_guide/examples/explanation_kernelloop.md

@@ -107,9 +107,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",
     "# Select fields\n",
     "fields = {\n",
     "    \"U\": ds_fields[\"uo\"],\n",
@@ -120,6 +117,7 @@
     "# Convert to SGRID-compliant dataset and create FieldSet\n",
     "ds_fset = parcels.convert.copernicusmarine_to_sgrid(fields=fields)\n",
     "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
+    "fieldset.to_windowed_arrays()\n",
     "\n",
     "# Define a new Particle type including extra Variables\n",
     "ArgoParticle = parcels.Particle.add_variable(\n",

docs/user_guide/examples/tutorial_Argofloats.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."
    ]
   },
   {
@@ -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()"
    ]
   },
   {