Create convert.croco_to_sgrid function

docs/user_guide/examples/tutorial_croco_3D.ipynb

@@ -0,0 +1,324 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 🖥️ CROCO tutorial"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This tutorial will show how to run a 3D simulation with output from the CROCO model."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Example setup"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We start with loading the relevant modules and the data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import xarray as xr\n",
+    "\n",
+    "import parcels\n",
+    "\n",
+    "data_folder = parcels.download_example_dataset(\"CROCOidealized_data\")\n",
+    "ds_fields = xr.open_dataset(data_folder / \"CROCO_idealized.nc\")\n",
+    "\n",
+    "ds_fields.load();  # Preload data to speed up access"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The simulation will be a very simple, idealised flow: a purely zonal flow over a sloping bottom. This flow (which is somewhat unrealistic of course) nicely showcases that particles stay on their initial depth levels, even though the sigma-layers slope down. \n",
+    "\n",
+    "This flow has been created by first running the example from the [Shelf front example on the CROCO website](https://croco-ocean.gitlabpages.inria.fr/croco_doc/model/model.test_cases.shelfront.html). Then, we took the restart file are manually replaced all the `u`-velocities with `1` m/s and all the `v`-velocities with `0` m/s. This way we get a purely zonal flow. We then started a new simulation from the restart file, and CROCO then automatically calculated the `w` velocities to match the new zonal field. We saved the `time=0` snapshot from this new run and use it below."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we create a FieldSet object using the `FieldSet.from_croco()` method. Note that CROCO is a C-grid (with similar indexing at MITgcm)."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "```{note}\n",
+    "In the code below, we use the `w` velocity field for vertical velocity. However, it is unclear whether this is always the right choice. CROCO (and ROMS) also output an `omega` field, which may be more appropriate to use. The idealised simulation below only works when using `w`, though. In other simulations, it is recommended to test whether `omega` provides more realistic results. See https://github.com/OceanParcels/Parcels/discussions/1728 for more information.\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fields = {\n",
+    "    \"U\": ds_fields[\"u\"],\n",
+    "    \"V\": ds_fields[\"v\"],\n",
+    "    \"W\": ds_fields[\"w\"],\n",
+    "    \"h\": ds_fields[\"h\"],\n",
+    "    \"zeta\": ds_fields[\"zeta\"],\n",
+    "    \"Cs_w\": ds_fields[\"Cs_w\"],\n",
+    "}\n",
+    "\n",
+    "coords = ds_fields[[\"x_rho\", \"y_rho\", \"s_w\", \"time\"]]\n",
+    "ds_fset = parcels.convert.croco_to_sgrid(fields=fields, coords=coords)\n",
+    "\n",
+    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
+    "\n",
+    "# Add the critical depth (`hc`) as a constant to the fieldset\n",
+    "fieldset.add_constant(\"hc\", ds_fields.hc)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can use this Fieldset to advect particles as we would normally do. Note that the particle depths should be provided in (negative) meters, not in sigma-coordinates."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-output"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "X, Z = np.meshgrid(\n",
+    "    [40e3, 80e3, 120e3],\n",
+    "    [100, -10, -130, -250, -400, -850, -1400, -1550],\n",
+    ")\n",
+    "Y = np.ones(X.size) * 98000\n",
+    "\n",
+    "\n",
+    "def DeleteParticle(particles, fieldset):\n",
+    "    any_error = particles.state >= 50\n",
+    "    particles[any_error].state = parcels.StatusCode.Delete\n",
+    "\n",
+    "\n",
+    "pset = parcels.ParticleSet(\n",
+    "    fieldset=fieldset, pclass=parcels.Particle, lon=X, lat=Y, z=Z\n",
+    ")\n",
+    "\n",
+    "outputfile = parcels.ParticleFile(\n",
+    "    store=\"croco_particles3D.zarr\",\n",
+    "    outputdt=np.timedelta64(5000, \"s\"),\n",
+    ")\n",
+    "\n",
+    "pset.execute(\n",
+    "    [parcels.kernels.AdvectionRK4_3D_CROCO, DeleteParticle],\n",
+    "    runtime=np.timedelta64(50_000, \"s\"),\n",
+    "    dt=np.timedelta64(100, \"s\"),\n",
+    "    output_file=outputfile,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we plot the particle trajectories below. Note that the particles stay on their initial depth levels, even though the sigma-layers slope down. Also note that particles released above the surface (where depth >0) or below the bathymetry are not advected (due to the `DeleteParticle` kernel)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "nbsphinx-thumbnail"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(1, 1, figsize=(6, 4))\n",
+    "ds = xr.open_zarr(\"croco_particles3D.zarr\")\n",
+    "\n",
+    "ax.plot(X / 1e3, Z, \"k.\", label=\"Initial positions\")\n",
+    "ax.plot(ds.lon.T / 1e3, ds.z.T, \".-\")\n",
+    "\n",
+    "for z in ds_fields.s_w.values:\n",
+    "    ax.plot(\n",
+    "        fieldset.h.grid.lon[0, :] / 1e3,\n",
+    "        fieldset.h.data[0, 0, 25, :] * z,\n",
+    "        \"k\",\n",
+    "        linewidth=0.5,\n",
+    "    )\n",
+    "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",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### A CROCO simulation with no vertical velocities"
+   ]
+  },
+  {
+   "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."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "hide-output"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "fieldset_noW = parcels.FieldSet.from_sgrid_conventions(ds_fset)\n",
+    "fieldset_noW.W.data[:] = 0.0\n",
+    "fieldset_noW.add_constant(\"hc\", ds_fields.hc)\n",
+    "\n",
+    "X, Z = np.meshgrid(\n",
+    "    [40e3, 80e3, 120e3],\n",
+    "    [100, -10, -130, -250, -400, -850, -1400, -1550],\n",
+    ")\n",
+    "Y = np.ones(X.size) * 100\n",
+    "\n",
+    "pset_noW = parcels.ParticleSet(\n",
+    "    fieldset=fieldset_noW, pclass=parcels.Particle, lon=X, lat=Y, z=Z\n",
+    ")\n",
+    "\n",
+    "outputfile = parcels.ParticleFile(\n",
+    "    store=\"croco_particles_noW.zarr\", outputdt=np.timedelta64(5000, \"s\")\n",
+    ")\n",
+    "\n",
+    "pset_noW.execute(\n",
+    "    [parcels.kernels.AdvectionRK4_3D_CROCO, DeleteParticle],\n",
+    "    runtime=np.timedelta64(50_000, \"s\"),\n",
+    "    dt=np.timedelta64(100, \"s\"),\n",
+    "    output_file=outputfile,\n",
+    ")\n",
+    "\n",
+    "fig, ax = plt.subplots(1, 1, figsize=(6, 4))\n",
+    "ds = xr.open_zarr(\"croco_particles_noW.zarr\")\n",
+    "\n",
+    "ax.plot(X / 1e3, Z, \"k.\", label=\"Initial positions\")\n",
+    "ax.plot(ds.lon.T / 1e3, ds.z.T, \".-\")\n",
+    "\n",
+    "for z in ds_fields.s_w.values:\n",
+    "    ax.plot(\n",
+    "        fieldset.h.grid.lon[0, :] / 1e3,\n",
+    "        fieldset.h.data[0, 0, 25, :] * z,\n",
+    "        \"k\",\n",
+    "        linewidth=0.5,\n",
+    "    )\n",
+    "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 with W=0\")\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## The algorithms used"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "When using Croco data, Parcels needs to convert the particles.depth to sigma-coordinates, before doing any interpolation. This is done with the `convert_z_to_sigma_croco()`  function. Interpolating onto a Field is then done like:\n",
+    "\n",
+    "```python\n",
+    "def SampleTempCroco(particles, fieldset):\n",
+    "    \"\"\"Sample temperature field on a CROCO sigma grid by first converting z to sigma levels.\"\"\"\n",
+    "    sigma = convert_z_to_sigma_croco(fieldset, particles.time, particles.z, particles.lat, particles.lon, particles)\n",
+    "    particles.temp = fieldset.T[particles.time, sigma, particles.lat, particles.lon, particles]\n",
+    "```"
+   ]
+  },
+  {
+   "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 depth 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",
+    "\n",
+    "```python\n",
+    "sigma = particles.z / fieldset.h[particles.time, 0, particles.lat, particles.lon]\n",
+    "\n",
+    "sigma_levels = convert_z_to_sigma_croco(fieldset, particles.time, particles.z, particles.lat, particles.lon, particles)\n",
+    "(u, v) = fieldset.UV[time, sigma_levels, particle.lat, particle.lon, particle]\n",
+    "w = fieldset.W[time, sigma_levels, particle.lat, particle.lon, particle]\n",
+    "\n",
+    "# scaling the w with the sigma level of the particle\n",
+    "w_sigma = w * sigma / fieldset.h[particles.time, 0, particles.lat, particles.lon]\n",
+    "\n",
+    "lon_new = particles.lon + u*particles.dt\n",
+    "lat_new = particles.lat + v*particles.dt\n",
+    "\n",
+    "# calculating new sigma level\n",
+    "sigma_new = sigma + w_sigma*particles.dt \n",
+    "\n",
+    "# converting back from sigma to depth, at _new_ location\n",
+    "depth_new = sigma_new * fieldset.h[particles.time, 0, lat_new, lon_new]\n",
+    "```"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "docs",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.14.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}