diff --git a/docs/user_guide/convert_croco_3D.ipynb b/docs/user_guide/convert_croco_3D.ipynb index 585d32bc1..f41ba0a3e 100644 --- a/docs/user_guide/convert_croco_3D.ipynb +++ b/docs/user_guide/convert_croco_3D.ipynb @@ -11,21 +11,35 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "This tutorial will show how to run a 3D simulation with output from the CROCO model." + "This tutorial will show how to run a 3D simulation with output from the CROCO model. \n", + "\n", + "We start with loading the relevant modules and the data." ] }, { - "cell_type": "markdown", + "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ - "## Example setup" + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import polars as pl\n", + "import xroms\n", + "\n", + "import parcels\n", + "import parcels.tutorial" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "We start with loading the relevant modules and the data." + "## Create FieldSet\n", + "\n", + "The dataset loaded below comes from the output of running the [CROCO Gravitational/Overflow example](https://croco-ocean.gitlabpages.inria.fr/croco_doc/model/model.test_cases.overflow.html), without any modifications to the CROCO configuration.\n", + "\n", + "We use the [`xroms`](https://xroms.readthedocs.io/) package to add the vertical coordinates (depth as a function of time and space) to the dataset." ] }, { @@ -34,39 +48,85 @@ "metadata": {}, "outputs": [], "source": [ - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import polars as pl\n", + "ds_fields = parcels.tutorial.open_dataset(\"CROCOoverflow_data/data\")\n", "\n", - "import parcels\n", - "import parcels.tutorial\n", + "# Add ROMS/CROCO vertical coordinates and grid metrics (note that we need to rename Cs_rho to Cs_r for `xroms` to work)\n", + "ds_fields = ds_fields.rename({\"Cs_rho\": \"Cs_r\"})\n", "\n", - "ds_fields = parcels.tutorial.open_dataset(\"CROCOidealized_data/data\")" + "# ds_fields, xgrid = xroms.roms_dataset(ds_fields)\n", + "# ds_fields.xroms.set_grid(xgrid)" ] }, { "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", + "```{note}\n", + "Note that there is currently (Nov 2026) a compatibility issue between `xroms` version 0.6.1 and `xgcm` version 0.10.1. The following code snippet provides a compatibility adapter to work around this issue. Hopefully, this will be fixed in a future release of either package.\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "from unittest.mock import patch\n", + "\n", + "import xgcm\n", + "\n", + "# Compatibility adapter for xroms 0.6.1 with xgcm 0.10.1.\n", + "OriginalGrid = xgcm.Grid\n", + "\n", + "\n", + "class XromsCompatibleGrid(OriginalGrid):\n", + " def __init__(self, *args, **kwargs):\n", + " kwargs.setdefault(\"autoparse_metadata\", False)\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." + " if \"periodic\" in kwargs:\n", + " if kwargs.pop(\"periodic\") != []:\n", + " raise ValueError(\"This adapter only supports non-periodic grids\")\n", + " kwargs.setdefault(\"padding\", \"fill\")\n", + "\n", + " super().__init__(*args, **kwargs)\n", + "\n", + " def diff(self, *args, **kwargs):\n", + " if \"boundary\" in kwargs:\n", + " kwargs[\"padding\"] = kwargs.pop(\"boundary\")\n", + " return super().diff(*args, **kwargs)\n", + "\n", + " def interp(self, *args, **kwargs):\n", + " if \"boundary\" in kwargs:\n", + " kwargs[\"padding\"] = kwargs.pop(\"boundary\")\n", + " return super().interp(*args, **kwargs)\n", + "\n", + " def derivative(self, *args, **kwargs):\n", + " if \"boundary\" in kwargs:\n", + " kwargs[\"padding\"] = kwargs.pop(\"boundary\")\n", + " return super().derivative(*args, **kwargs)" ] }, { - "cell_type": "markdown", + "cell_type": "code", + "execution_count": null, "metadata": {}, + "outputs": [], "source": [ - "Now we create a FieldSet object using the {py:func}`parcels.convert.croco_to_sgrid()` function to first create an S-Grid compliant datatset, and then use that in {py:func}`parcels.FieldSet.from_sgrid_conventions()` to create the FieldSet." + "with patch.object(xgcm, \"Grid\", XromsCompatibleGrid):\n", + " ds_fields, xgrid = xroms.roms_dataset(ds_fields)\n", + " ds_fields.xroms.set_grid(xgrid)" ] }, { "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/Parcels-code/Parcels/discussions/1728 for more information.\n", - "```" + "Now we create a FieldSet object using the {py:func}`parcels.convert.croco_to_sgrid()` function to first create an S-Grid compliant datatset, and then use that in {py:func}`parcels.FieldSet.from_sgrid_conventions()` to create the FieldSet." ] }, { @@ -79,27 +139,30 @@ " \"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", + " \"depth\": ds_fields[\"z_w_psi\"],\n", + " \"seafloor\": ds_fields[\"h\"],\n", "}\n", "\n", "ds_fset = parcels.convert.croco_to_sgrid(fields=fields, coords=ds_fields)\n", + "ds_fset.time.attrs[\"axis\"] = \"T\"\n", "\n", "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())\n", - "\n", - "# Convert the FieldSet to windowed arrays for better performance\n", - "fieldset = fieldset.to_windowed_arrays()" + "# Convert the FieldSet to ChunkCachedArray objects for better performance\n", + "fieldset = fieldset.to_chunk_cached_arrays()" ] }, { "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." + "## Create and execute ParticleSet\n", + "\n", + "Now we can use this Fieldset to advect particles as we would normally do. We first define the initial positions of the particles in the domain.\n", + "\n", + "```{note}\n", + "Note that the particle depths should be provided in (negative) meters, not in sigma-coordinates.\n", + "```" ] }, { @@ -112,29 +175,69 @@ }, "outputs": [], "source": [ - "X, Z = np.meshgrid(\n", - " [40e3, 80e3, 120e3],\n", - " [100, -10, -130, -250, -400, -850, -1400, -1550],\n", + "Y, Z = np.meshgrid(\n", + " np.arange(5e3, 60e3, 5e3),\n", + " np.arange(-41, 0, 5),\n", ")\n", - "Y = np.ones(X.size) * 98000\n", + "Y, Z = Y.flatten(), Z.flatten()\n", + "X = 1000 * np.ones(Y.size)\n", + "T = np.timedelta64(0, \"s\") * np.arange(X.size)\n", "\n", + "pset_init = parcels.ParticleSet(\n", + " fieldset=fieldset, pclass=parcels.Particle, x=X, y=Y, z=Z, t=T\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Because some particles will be located below the seafloor, we filter them out before creating the ParticleSet." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "inds = np.where(fieldset.seafloor[pset_init] != 0)\n", "\n", + "pset = parcels.ParticleSet(\n", + " fieldset=fieldset,\n", + " pclass=parcels.Particle,\n", + " x=X[inds],\n", + " y=Y[inds],\n", + " z=Z[inds],\n", + " t=T[inds],\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-output" + ] + }, + "outputs": [], + "source": [ "def DeleteParticle(particles, fieldset):\n", " any_error = particles.state >= 50\n", " particles[any_error].state = parcels.StatusCode.Delete\n", "\n", "\n", - "pset = parcels.ParticleSet(fieldset=fieldset, pclass=parcels.Particle, x=X, y=Y, z=Z)\n", - "\n", "outputfile = parcels.ParticleFile(\n", " path=\"croco_particles3D.parquet\",\n", - " outputdt=np.timedelta64(5000, \"s\"),\n", + " outputdt=np.timedelta64(3, \"h\"),\n", + " mode=\"w\",\n", ")\n", "\n", "pset.execute(\n", - " [parcels.kernels.AdvectionRK2_3D_CROCO, DeleteParticle],\n", - " runtime=np.timedelta64(50_000, \"s\"),\n", - " dt=np.timedelta64(100, \"s\"),\n", + " [parcels.kernels.AdvectionRK2_3D, DeleteParticle],\n", + " runtime=fieldset.time_interval.right - fieldset.time_interval.left,\n", + " dt=np.timedelta64(15, \"m\"),\n", " output_file=outputfile,\n", ")" ] @@ -143,7 +246,9 @@ "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)." + "## Plot trajectories\n", + "\n", + "Now we plot the particle trajectories below. Note that the particles roughly follow the sloping sigma-layers." ] }, { @@ -157,23 +262,21 @@ "outputs": [], "source": [ "fig, ax = plt.subplots(1, 1, figsize=(6, 4))\n", - "df = pl.read_parquet(\"croco_particles3D.parquet\")\n", + "df = parcels.read_particlefile(\"croco_particles3D.parquet\")\n", "\n", - "ax.plot(X / 1e3, Z, \"k.\", label=\"Initial positions\")\n", "for traj in df.partition_by(\"particle_id\", maintain_order=True):\n", - " ax.plot(traj[\"x\"] / 1e3, traj[\"z\"], \".-\")\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.plot(traj[\"y\"] / 1e3, traj[\"z\"], \".-\")\n", + "ax.plot(\n", + " fieldset.depth.grid.lat[:-1, 1] / 1e3,\n", + " fieldset.depth.data[-1, :, :, 1].T,\n", + " \"k\",\n", + " linewidth=0.5,\n", + ")\n", + "ax.set_xlabel(\"Y [km]\")\n", "ax.set_ylabel(\"Depth [m]\")\n", - "ax.set_title(\"Particles in idealized CROCO velocity field using AdvectionRK2_3D_CROCO\")\n", + "ax.set_title(\n", + " \"Particles in CROCO Gravitational/Overflow velocity field using AdvectionRK2_3D\"\n", + ")\n", "plt.tight_layout()\n", "plt.show()" ] @@ -182,14 +285,9 @@ "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 `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." + "## Animate trajectories\n", + "\n", + "We can also make an animation, which clearly shows the development of the gravity current" ] }, { @@ -197,119 +295,100 @@ "execution_count": null, "metadata": { "tags": [ - "hide-output" + "hide-input" ] }, "outputs": [], "source": [ - "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", - " [100, -10, -130, -250, -400, -850, -1400, -1550],\n", + "import matplotlib\n", + "from matplotlib.animation import FuncAnimation\n", + "from matplotlib.collections import LineCollection\n", + "\n", + "# for interactive display of animation\n", + "plt.rcParams[\"animation.html\"] = \"jshtml\"\n", + "\n", + "time_step = np.timedelta64(6, \"h\") # time step for animation frames\n", + "\n", + "# Build the time axis for the animation.\n", + "timerange = np.arange(\n", + " np.nanmin(df[\"t\"]),\n", + " np.nanmax(df[\"t\"]) + time_step,\n", + " time_step,\n", ")\n", - "Y = np.ones(X.size) * 100\n", "\n", - "pset_noW = parcels.ParticleSet(\n", - " fieldset=fieldset_noW, pclass=parcels.Particle, x=X, y=Y, z=Z\n", + "# figure setup\n", + "fig, ax = plt.subplots(figsize=(6, 4))\n", + "ax.plot(\n", + " fieldset.depth.grid.lat[:-1, 1] / 1e3,\n", + " fieldset.depth.data[-1, :, :, 1].T,\n", + " \"k\",\n", + " linewidth=0.5,\n", ")\n", + "ax.set_xlabel(\"Y [km]\")\n", + "ax.set_ylabel(\"Depth [m]\")\n", + "cmap = matplotlib.colormaps[\"tab20b\"]\n", "\n", - "outputfile = parcels.ParticleFile(\n", - " path=\"croco_particles_noW.parquet\", outputdt=np.timedelta64(5000, \"s\")\n", + "trails = LineCollection([], linewidths=0.6, alpha=0.3)\n", + "ax.add_collection(trails)\n", + "\n", + "particle_ids = df[\"particle_id\"].unique().to_list()\n", + "\n", + "# Convert the trajectory data to a time-by-particle array and assign colors.\n", + "colors = np.asarray(\n", + " [cmap(i / max(len(particle_ids) - 1, 1)) for i in range(len(particle_ids))]\n", ")\n", "\n", - "pset_noW.execute(\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", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(1, 1, figsize=(6, 4))\n", - "df = pl.read_parquet(\"croco_particles_noW.parquet\")\n", + "# Create arrays to hold the x and y positions of each particle at each time step.\n", + "y = np.full((len(particle_ids), len(timerange)), np.nan)\n", + "z = np.full((len(particle_ids), len(timerange)), np.nan)\n", "\n", - "ax.plot(X / 1e3, Z, \"k.\", label=\"Initial positions\")\n", - "for traj in df.partition_by(\"particle_id\", maintain_order=True):\n", - " ax.plot(traj[\"x\"] / 1e3, traj[\"z\"], \".-\")\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.z 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", + "traj = df.with_columns(\n", + " pl.col(\"particle_id\")\n", + " .replace(particle_ids, range(len(particle_ids)))\n", + " .alias(\"p_idx\"),\n", + " pl.int_range(pl.len()).alias(\"row_idx\"),\n", + ")\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.t, particles.z, particles.y, particles.x, particles)\n", - " particles.temp = fieldset.T[particles.t, sigma, particles.y, particles.x, particles]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "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", + "# Map each observation to its nearest animation time index.\n", + "traj = traj.with_columns(\n", + " pl.Series(\"t_idx\", np.searchsorted(timerange, traj[\"t\"].to_numpy())).alias(\"t_idx\")\n", + ")\n", "\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", + "y[traj[\"p_idx\"].to_numpy(), traj[\"t_idx\"].to_numpy()] = traj[\"y\"].to_numpy() / 1e3\n", + "z[traj[\"p_idx\"].to_numpy(), traj[\"t_idx\"].to_numpy()] = traj[\"z\"].to_numpy()\n", "\n", - "```python\n", - "sigma = particles.z / fieldset.h[particles.t, 0, particles.y, particles.x]\n", + "# Plot first timestep\n", + "scatter = ax.scatter(y[:, 0], z[:, 0], s=10, c=colors, zorder=2)\n", "\n", - "sigma_levels = convert_z_to_sigma_croco(fieldset, particles.t, particles.z, particles.y, particles.x, particles)\n", - "(u, v) = fieldset.UV[particles.t, sigma_levels, particles.y, particles.x, particles]\n", - "w = fieldset.W[particles.t, sigma_levels, particles.y, particles.x, particles]\n", + "# Set initial title\n", + "t_str = timerange[0].astype(\"timedelta64[h]\")\n", + "title = ax.set_title(f\"Particles at t = {t_str}\")\n", "\n", - "# scaling the w with the sigma level of the particle\n", - "w_sigma = w * sigma / fieldset.h[particles.t, 0, particles.y, particles.x, particles]\n", "\n", - "x_new = particles.x + u*particles.dt\n", - "y_new = particles.y + v*particles.dt\n", + "def animate(i):\n", + " t_str = timerange[i].astype(\"timedelta64[h]\")\n", + " title.set_text(f\"Particles at t = {t_str}\")\n", "\n", - "# calculating new sigma level\n", - "sigma_new = sigma + w_sigma*particles.dt \n", + " scatter.set_offsets(np.column_stack((y[:, i], z[:, i])))\n", "\n", - "# converting back from sigma to z, at _new_ location\n", - "z_new = sigma_new * fieldset.h[particles.t, 0, y_new, x_new, particles]\n", - "```" + " trail_length = min(10, i) # trails have max length of 10 time steps\n", + " start = max(0, i - trail_length)\n", + " y_slice = y[:, start : i + 1]\n", + " z_slice = z[:, start : i + 1]\n", + " trails.set_segments(np.dstack((y_slice, z_slice)))\n", + " trails.set_color(colors)\n", + "\n", + "\n", + "# Create animation\n", + "anim = FuncAnimation(fig, animate, frames=len(timerange), interval=100)\n", + "plt.close(fig)\n", + "anim" ] } ], "metadata": { "kernelspec": { - "display_name": "Parcels:docs (3.14.6)", + "display_name": "Parcels:default (3.14.8)", "language": "python", "name": "python3" }, @@ -323,7 +402,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.14.6" + "version": "3.14.8" }, "mystnb": { "execution_in_temp": true diff --git a/docs/user_guide/fieldsets_grids.ipynb b/docs/user_guide/fieldsets_grids.ipynb index a6fa91b17..5201c28db 100644 --- a/docs/user_guide/fieldsets_grids.ipynb +++ b/docs/user_guide/fieldsets_grids.ipynb @@ -333,7 +333,7 @@ "source": [ "### Vertical coordinates\n", "\n", - "Lastly, a short note on vertical coordinates. Parcels can handle two styles of vertical coordinates; z-levels which define fixed depth levels in physical space, and sigma-levels which define varying depth levels in physical space as a function of the water column depth. As sigma-levels are effectively \"terrain-following\", the grid cell faces may no longer be orthogonal to the domain surface. The bottom row of Figure 1 visualises these differences, and Parcels handles this all under the hood. See the [CROCO Tutorial](./convert_croco_3D.ipynb) for more details on the sigma-levels implementation.\n" + "Lastly, a short note on vertical coordinates. Parcels can handle two styles of vertical coordinates; z-levels which define fixed depth levels in physical space, and sigma-levels which define varying depth levels in physical space as a function of the water column depth. As sigma-levels are effectively \"terrain-following\", the grid cell faces may no longer be orthogonal to the domain surface. The bottom row of Figure 1 visualises these differences, and Parcels handles this all under the hood. See the [Sigma Coordinate Vertical Grids tutorial](./fieldsets_sigma_coordinates.ipynb) for more details on the sigma-levels implementation.\n" ] } ], diff --git a/docs/user_guide/v4-migration.md b/docs/user_guide/v4-migration.md index c7196ef32..929604cbd 100644 --- a/docs/user_guide/v4-migration.md +++ b/docs/user_guide/v4-migration.md @@ -176,10 +176,10 @@ Particle-particle interaction can be performe
-Automatic conversion from depth to sigma grids under the hood is removed +The Advection_CROCO methods have been removed
-Explicitly use convert_z_to_sigma_croco in sampling kernels (such as the AdvectionRK4_3D_CROCO or SampleOMegaCroco kernels) when working with CROCO data +use the 4-D depth support when working with CROCO data
diff --git a/pixi.toml b/pixi.toml index 0c0251866..0b6b2a3fc 100644 --- a/pixi.toml +++ b/pixi.toml @@ -74,6 +74,7 @@ pytest-html = "*" pytest-cov = "*" pytest-reportlog = "*" pytest-timeout = "*" +xroms = "*" [feature.test.pypi-dependencies] re-assert = "*" @@ -92,6 +93,7 @@ gsw = "*" py-triangle = "*" cmocean = "*" cartopy = "*" +xroms = "*" xgcm = ">=0.10.0" [feature.devtools.dependencies] diff --git a/src/parcels/_datasets/remote.py b/src/parcels/_datasets/remote.py index 0cb9e3299..3dc8d402f 100644 --- a/src/parcels/_datasets/remote.py +++ b/src/parcels/_datasets/remote.py @@ -119,7 +119,7 @@ def _get_data_home() -> Path: "data/SWASH_data/field_0065557.nc", ] # + [f"data/WOA_data/woa18_decav_t{m:02d}_04.nc" for m in range(1, 13)] - + ["data/CROCOidealized_data/CROCO_idealized.nc"] + + ["data/CROCOOverflow_data/over_his.nc"] # These datasets are from v4 of Parcels where we're opting for Zipped zarr datasets + [ "data-zarr/Benchmarks_FESOM2-baroclinic-gyre/data.zip", @@ -264,7 +264,7 @@ class _Purpose(enum.Enum): # "POPSouthernOcean_data/t.x1_SAMOC_flux.16900*.nc", # TODO v4: In v3 but should not be in v4 https://github.com/Parcels-code/Parcels/issues/2571#issuecomment-4214476973 # ("SWASH_data/data", (_V3Dataset(_ODIE,"data/SWASH_data/field_00655*.nc"), _Purpose.TUTORIAL)), # ("WOA_data/data", (_V3Dataset(_ODIE,"data/WOA_data/woa18_decav_t*_04.nc", _preprocess_set_cf_calendar_360_day), _Purpose.TUTORIAL)), - ("CROCOidealized_data/data", (_V3Dataset(_ODIE,"data/CROCOidealized_data/CROCO_idealized.nc"), _Purpose.TUTORIAL)), + ("CROCOoverflow_data/data", (_V3Dataset(_ODIE,"data/CROCOOverflow_data/over_his.nc"), _Purpose.TUTORIAL)), ("AIFSmeteo/data", (_V3Dataset(_ODIE,"data/AIFS-meteo-data/*.nc"), _Purpose.TUTORIAL)) ] + [ ("Benchmarks_FESOM2-baroclinic-gyre/data", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/data.zip', zarr_format=2), _Purpose.TESTING)), diff --git a/src/parcels/_typing.py b/src/parcels/_typing.py index 62b470bf8..d5d9f2653 100644 --- a/src/parcels/_typing.py +++ b/src/parcels/_typing.py @@ -17,7 +17,7 @@ from parcels._sgrid.core import Padding PathLike = str | os.PathLike -VectorType = Literal["3D", "3DSigma", "2D"] | None # corresponds with `vector_type` +VectorType = Literal["3D", "2D"] | None # corresponds with `vector_type` TimeLike = datetime | cftime_datetime | np.datetime64 CfAxisSpatial = Literal["X", "Y", "Z"] diff --git a/src/parcels/convert.py b/src/parcels/convert.py index ec82ada43..a92cc7d69 100644 --- a/src/parcels/convert.py +++ b/src/parcels/convert.py @@ -145,7 +145,6 @@ class _ModelConventions: varnames_mapping={ "x_rho": "lon", "y_rho": "lat", - "s_w": "depth", }, ) @@ -605,7 +604,7 @@ def croco_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.D sgrid.FaceNodePadding("xi_u", "xi_rho", sgrid.Padding.NONE), sgrid.FaceNodePadding("eta_v", "eta_rho", sgrid.Padding.NONE), ), - vertical_dimensions=(sgrid.FaceNodePadding("s_rho", "depth", sgrid.Padding.NONE),), + vertical_dimensions=(sgrid.FaceNodePadding("s_rho", "s_w", sgrid.Padding.NONE),), ), ) diff --git a/src/parcels/kernels/__init__.py b/src/parcels/kernels/__init__.py index 8ea1d0686..265eac714 100644 --- a/src/parcels/kernels/__init__.py +++ b/src/parcels/kernels/__init__.py @@ -13,11 +13,6 @@ DiffusionUniformKh, ) from ._integration import RK2, RK4 -from ._sigmagrids import ( - AdvectionRK2_3D_CROCO, - SampleOmegaCroco, - convert_z_to_sigma_croco, -) __all__ = [ # noqa: RUF022 # advection @@ -32,10 +27,6 @@ "AdvectionDiffusionEM", "AdvectionDiffusionM1", "DiffusionUniformKh", - # sigmagrids - "AdvectionRK2_3D_CROCO", - "SampleOmegaCroco", - "convert_z_to_sigma_croco", # integration "RK2", "RK4", diff --git a/src/parcels/kernels/_sigmagrids.py b/src/parcels/kernels/_sigmagrids.py deleted file mode 100644 index 72ea991d7..000000000 --- a/src/parcels/kernels/_sigmagrids.py +++ /dev/null @@ -1,72 +0,0 @@ -import warnings - -import numpy as np - - -def convert_z_to_sigma_croco(fieldset, t, z, y, x, particle): - """Calculate local sigma level of the particles, by linearly interpolating the - scaling function that maps sigma to depth (using local ocean depth h, - sea-surface Zeta and stretching parameters Cs_w and hc). - See also https://croco-ocean.gitlabpages.inria.fr/croco_doc/model/model.grid.html#vertical-grid-parameters - """ - h = fieldset.h.eval(t, np.zeros_like(z), y, x, particles=particle) - zeta = fieldset.zeta.eval(t, np.zeros_like(z), y, x, particles=particle) - sigma_levels = fieldset.U.grid.depth - cs_w = fieldset.Cs_w.data.values.flatten() - - z0 = fieldset.hc * sigma_levels[None, :] + (h[:, None] - fieldset.hc) * cs_w[None, :] - zvec = z0 + zeta[:, None] * (1.0 + (z0 / h[:, None])) - zinds = zvec <= z[:, None] - zi = np.argmin(zinds, axis=1) - 1 - zi = np.where(zinds.all(axis=1), zvec.shape[1] - 2, zi) - idx = np.arange(zi.shape[0]) - return sigma_levels[zi] + (z - zvec[idx, zi]) * (sigma_levels[zi + 1] - sigma_levels[zi]) / ( - zvec[idx, zi + 1] - zvec[idx, zi] - ) - - -def SampleOmegaCroco(particles, fieldset): - """Sample omega field on a CROCO sigma grid by first converting z to sigma levels. - - This Kernel can be adapted to sample any other field on a CROCO sigma grid by - replacing 'omega' with the desired field name. - """ - sigma = convert_z_to_sigma_croco(fieldset, particles.t, particles.z, particles.y, particles.x, particles) - particles.omega = fieldset.omega[particles.t, sigma, particles.y, particles.x, particles] - - -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. - """ - with warnings.catch_warnings(): - warnings.filterwarnings( - "ignore", - message=r"^Sampling of velocities should normally be done using fieldset\.UV or fieldset\.UVW object; tread carefully$", - category=RuntimeWarning, - ) # Needed because of linear sampling of W with sigma conversion - - sigma = particles.z / fieldset.h[particles.t, np.zeros_like(particles.z), particles.y, particles.x] - - sig = convert_z_to_sigma_croco(fieldset, particles.t, particles.z, particles.y, particles.x, particles) - (u1, v1) = fieldset.UV[particles.t, sig, particles.y, particles.x, particles] - w1 = fieldset.W[particles.t, sig, particles.y, particles.x, particles] - w1 *= sigma / fieldset.h[particles.t, np.zeros_like(particles.z), particles.y, particles.x] - x1 = particles.x + u1 * 0.5 * particles.dt - y1 = particles.y + v1 * 0.5 * particles.dt - sig_dep1 = sigma + w1 * 0.5 * particles.dt - dep1 = sig_dep1 * fieldset.h[particles.t, np.zeros_like(particles.z), y1, x1] - - sig1 = convert_z_to_sigma_croco(fieldset, particles.t + 0.5 * particles.dt, dep1, y1, x1, particles) - (u2, v2) = fieldset.UV[particles.t + 0.5 * particles.dt, sig1, y1, x1, particles] - w2 = fieldset.W[particles.t + 0.5 * particles.dt, sig1, y1, x1, particles] - w2 *= sig_dep1 / fieldset.h[particles.t + 0.5 * particles.dt, np.zeros_like(particles.z), y1, x1] - x2 = particles.x + u2 * 0.5 * particles.dt - y2 = particles.y + v2 * 0.5 * particles.dt - sig_dep2 = sigma + w2 * 0.5 * particles.dt - dep2 = sig_dep2 * fieldset.h[particles.t + 0.5 * particles.dt, np.zeros_like(particles.z), y2, x2] - - particles.dx += u2 * particles.dt - particles.dy += v2 * particles.dt - particles.dz += (dep1 - particles.z) + (dep2 - particles.z) diff --git a/tests/test_convert.py b/tests/test_convert.py index 6d233b1a0..6317007ce 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -2,6 +2,7 @@ import pytest import uxarray as ux import xarray as xr +import xroms import parcels import parcels.convert as convert @@ -116,9 +117,12 @@ def test_convert_mitgcm_offsets(): assert offsets["Z"] == 0 +@pytest.mark.xfail(reason="Compatibility issue between xroms v0.6.1 and xgcm v0.10.1") def test_convert_croco_offsets(): ds = datasets_circulation_models["ds_CROCO_idealized"] coords = ds[["x_rho", "y_rho", "s_w", "time"]] + ds, xgrid = xroms.roms_dataset(ds) + ds.xroms.set_grid(xgrid) ds = convert.croco_to_sgrid(fields={"U": ds["u"], "V": ds["v"]}, coords=coords) fieldset = FieldSet.from_sgrid_conventions(ds) diff --git a/tests/test_sigmagrids.py b/tests/test_sigmagrids.py deleted file mode 100644 index 8337091eb..000000000 --- a/tests/test_sigmagrids.py +++ /dev/null @@ -1,99 +0,0 @@ -import numpy as np - -import parcels -import parcels.tutorial -from parcels import Particle, ParticleSet, Variable -from parcels.kernels import AdvectionRK2, AdvectionRK2_3D_CROCO, SampleOmegaCroco, convert_z_to_sigma_croco - - -def test_conversion_3DCROCO(): - """Test of the conversion from depth to sigma in CROCO - - Values below are retrieved using xroms and hardcoded in the method (to avoid dependency on xroms): - ```py - x, y = 10, 20 - s_xroms = ds.s_w.values - z_xroms = ds.z_w.isel(time=0).isel(eta_rho=y).isel(xi_rho=x).values - lat, lon = ds.y_rho.values[y, x], ds.x_rho.values[y, x] - ``` - """ - ds_fields = parcels.tutorial.open_dataset("CROCOidealized_data/data") - fields = { - "U": ds_fields["u"], - "V": ds_fields["v"], - "W": ds_fields["w"], - "h": ds_fields["h"], - "zeta": ds_fields["zeta"], - "Cs_w": ds_fields["Cs_w"], - } - - ds_fset = parcels.convert.croco_to_sgrid(fields=fields, coords=ds_fields) - - fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) - fieldset.add_context("hc", ds_fields.hc.item()) - - s_xroms = np.array([-1.0, -0.9, -0.8, -0.7, -0.6, -0.5, -0.4, -0.3, -0.2, -0.1, 0.0], dtype=np.float32) - z_xroms = np.array([-1.26000000e02, -1.10585846e02, -9.60985413e01, -8.24131317e01, -6.94126511e01, -5.69870148e01, -4.50318756e01, -3.34476166e01, -2.21383114e01, -1.10107975e01, 2.62768921e-02,], dtype=np.float32,) # fmt: skip - - time = np.zeros_like(z_xroms) - lon = np.full_like(z_xroms, 38000.0) - lat = np.full_like(z_xroms, 78000.0) - - sigma = convert_z_to_sigma_croco(fieldset, time, z_xroms, lat, lon, None) - - np.testing.assert_allclose(sigma, s_xroms, atol=1e-3) - - -def test_advection_2DCROCO(): - ds_fields = parcels.tutorial.open_dataset("CROCOidealized_data/data") - - fields = { - "U": ds_fields["u"], - "V": ds_fields["v"], - } - ds_fset = parcels.convert.croco_to_sgrid(fields=fields, coords=ds_fields) - fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) - fieldset = fieldset.to_windowed_arrays() - - runtime = 10_000 - X = np.array([40e3, 80e3, 120e3]) - Y = np.ones(X.size) * 100e3 - Z = np.zeros(X.size) - pset = ParticleSet(fieldset=fieldset, x=X, y=Y, z=Z) - - pset.execute([AdvectionRK2], runtime=runtime, dt=100) - assert np.allclose(pset.z, Z.flatten(), atol=1e-3) - assert np.allclose(pset.x, [x + runtime for x in X], atol=1e-3) - - -def test_advection_3DCROCO(): - ds_fields = parcels.tutorial.open_dataset("CROCOidealized_data/data") - - fields = { - "U": ds_fields["u"], - "V": ds_fields["v"], - "W": ds_fields["w"], - "h": ds_fields["h"], - "zeta": ds_fields["zeta"], - "Cs_w": ds_fields["Cs_w"], - "omega": ds_fields["omega"], - } - - ds_fset = parcels.convert.croco_to_sgrid(fields=fields, coords=ds_fields) - - fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) - fieldset = fieldset.to_windowed_arrays() - fieldset.add_context("hc", ds_fields.hc.item()) - - runtime = 10_000 - X, Z = np.meshgrid([40e3, 80e3, 120e3], [-10, -130]) - Y = np.ones(X.size) * 100e3 - - pclass = Particle.add_variable(Variable("omega")) - pset = ParticleSet(fieldset=fieldset, pclass=pclass, x=X, y=Y, z=Z) - - pset.execute( - [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)