diff --git a/docs/user_guide/examples/tutorial_croco_3D.ipynb b/docs/user_guide/examples/tutorial_croco_3D.ipynb new file mode 100644 index 0000000000..a9b388a17a --- /dev/null +++ b/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 +} diff --git a/docs/user_guide/examples_v3/tutorial_croco_3D.ipynb b/docs/user_guide/examples_v3/tutorial_croco_3D.ipynb deleted file mode 100644 index d46b2b706a..0000000000 --- a/docs/user_guide/examples_v3/tutorial_croco_3D.ipynb +++ /dev/null @@ -1,311 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# CROCO 3D 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 os\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import xarray as xr\n", - "\n", - "import parcels\n", - "\n", - "example_dataset_folder = parcels.download_example_dataset(\"CROCOidealized_data\")\n", - "file = os.path.join(example_dataset_folder, \"CROCO_idealized.nc\")" - ] - }, - { - "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), so we need to provide the longitudes and latitudes of the $\\rho$-points of the grid (`lon_rho` and `lat_rho`). We also need to provide the sigma levels at the depth points (`s_w`). Finally, it is important to also provide the bathymetry field (`h`), which is needed to convert the depth levels of the particles to sigma-coordinates." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "
\n", - "\n", - "__Note__ that 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": [ - "variables = {\"U\": \"u\", \"V\": \"v\", \"W\": \"w\", \"H\": \"h\", \"Zeta\": \"zeta\", \"Cs_w\": \"Cs_w\"}\n", - "\n", - "lon_rho = \"x_rho\" # Note, this would be \"lon_rho\" for a dataset on a spherical grid\n", - "lat_rho = \"y_rho\" # Note ,this would be \"lat_rho\" for a dataset on a spherical grid\n", - "\n", - "dimensions = {\n", - " \"U\": {\"lon\": lon_rho, \"lat\": lat_rho, \"depth\": \"s_w\", \"time\": \"time\"},\n", - " \"V\": {\"lon\": lon_rho, \"lat\": lat_rho, \"depth\": \"s_w\", \"time\": \"time\"},\n", - " \"W\": {\"lon\": lon_rho, \"lat\": lat_rho, \"depth\": \"s_w\", \"time\": \"time\"},\n", - " \"H\": {\"lon\": lon_rho, \"lat\": lat_rho},\n", - " \"Zeta\": {\"lon\": lon_rho, \"lat\": lat_rho, \"time\": \"time\"},\n", - " \"Cs_w\": {\"depth\": \"s_w\"},\n", - "}\n", - "fieldset = parcels.FieldSet.from_croco(\n", - " file,\n", - " variables,\n", - " dimensions,\n", - " hc=xr.open_dataset(\n", - " file\n", - " ).hc.values, # Note this stretching parameter is needed for the vertical grid\n", - " allow_time_extrapolation=True, # Note, this is only needed for this specific example dataset, that has only one snapshot\n", - " mesh=\"flat\", # Note, this is only needed for this specific example dataset, that has been created on a 'flat' mesh (i.e. in km instead of in degrees)\n", - ")" - ] - }, - { - "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": {}, - "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) * fieldset.U.grid.lat[25]\n", - "\n", - "\n", - "def DeleteParticle(particle, fieldset, time):\n", - " if particle.state >= 50:\n", - " particle.delete()\n", - "\n", - "\n", - "pset = parcels.ParticleSet(\n", - " fieldset=fieldset, pclass=parcels.Particle, lon=X, lat=Y, depth=Z\n", - ")\n", - "\n", - "outputfile = pset.ParticleFile(name=\"croco_particles3D.zarr\", outputdt=5000)\n", - "\n", - "pset.execute(\n", - " [parcels.kernels.AdvectionRK4_3D, DeleteParticle],\n", - " runtime=5e4,\n", - " dt=100,\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(ds.lon.T / 1e3, ds.z.T, \".-\")\n", - "\n", - "dsCROCO = xr.open_dataset(file)\n", - "for z in dsCROCO.s_w.values:\n", - " ax.plot(fieldset.H.lon / 1e3, fieldset.H.data[0, 25, :] * z, \"k\", linewidth=0.5)\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 3D advection\")\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` 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": {}, - "outputs": [], - "source": [ - "import copy\n", - "\n", - "fieldset_noW = copy.copy(fieldset)\n", - "fieldset_noW.W.data[:] = 0.0\n", - "\n", - "pset_noW = parcels.ParticleSet(\n", - " fieldset=fieldset_noW, pclass=parcels.Particle, lon=X, lat=Y, depth=Z\n", - ")\n", - "\n", - "outputfile = pset.ParticleFile(name=\"croco_particles_noW.zarr\", outputdt=5000)\n", - "\n", - "pset_noW.execute(\n", - " [parcels.kernels.AdvectionRK4_3D, DeleteParticle],\n", - " runtime=5e4,\n", - " dt=100,\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(ds.lon.T / 1e3, ds.z.T, \".-\")\n", - "\n", - "dsCROCO = xr.open_dataset(file)\n", - "for z in dsCROCO.s_w.values:\n", - " ax.plot(fieldset.H.lon / 1e3, fieldset.H.data[0, 25, :] * z, \"k\", linewidth=0.5)\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 `FieldSet.from_croco()`, Parcels knows that depth needs to be converted to sigma-coordinates, before doing any interpolation. This is done under the hood, using code for interpolation (in this case a `T` Field) like\n", - "```python\n", - "# First calculate local sigma level of the particle, 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\n", - "h = fieldset.H[time, 0, particle.lat, particle.lon]\n", - "zeta = fieldset.H[time, 0, particle.lat, particle.lon]\n", - "sigma_levels = fieldset.U.grid.depth\n", - "z0 = fieldset.hc * sigma_levels + (h - fieldset.hc) * fieldset.Cs_w\n", - "zvec = z0 + zeta * (1 + (z0 / h))\n", - "zinds = zvec <= z\n", - "if z >= zvec[-1]:\n", - " zi = len(zvec) - 2\n", - "else:\n", - " zi = zinds.argmin() - 1 if z >= zvec[0] else 0\n", - "\n", - "sigma = sigma_levels[zi] + (z - zvec[zi]) * (sigma_levels[zi + 1] - sigma_levels[zi]) / (zvec[zi + 1] - zvec[zi])\n", - "\n", - "# Now interpolate the field to the sigma level\n", - "temp = fieldset.T[time, sigma, particle.lat, particle.lon]\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For the `AdvectionRK4_3D` kernel, Parcels will replace the kernel with `AdvectionRK4_3D_CROCO`, 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 code shows above.\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)\n", - "\n", - "```python\n", - "(u, v, w) = fieldset.UVW[time, particle.depth, particle.lat, particle.lon, particle] \n", - "\n", - "# scaling the w with the sigma level of the particle\n", - "w_sigma = w * sigma / fieldset.H[time, particle.depth, particle.lat, particle.lon]\n", - "\n", - "lon_new = particle.lon + u*particle.dt\n", - "lat_new = particle.lat + v*particle.dt\n", - "\n", - "# calculating new sigma level\n", - "sigma_new = sigma + w_sigma*particle.dt \n", - "\n", - "# Converting back from sigma to depth, at _new_ location\n", - "depth_new = sigma_new * fieldset.H[time, particle.depth, lat_new, lon_new]\n", - "```" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "parcels", - "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.12.3" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index 35ecd9a7a8..8bf974a0bc 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -26,13 +26,13 @@ The tutorials written for Parcels v3 are currently being updated for Parcels v4. :titlesonly: examples/explanation_grids.md examples/tutorial_nemo.ipynb +examples/tutorial_croco_3D.ipynb examples/tutorial_velocityconversion.ipynb examples/tutorial_nestedgrids.ipynb examples/tutorial_manipulating_field_data.ipynb ``` - ```{toctree} diff --git a/docs/user_guide/v4-migration.md b/docs/user_guide/v4-migration.md index 263e006e8c..dbcdd36c10 100644 --- a/docs/user_guide/v4-migration.md +++ b/docs/user_guide/v4-migration.md @@ -15,6 +15,7 @@ Version 4 of Parcels is unreleased at the moment. The information in this migrat - `math` functions should be replaced with array compatible equivalents (e.g., `math.sin` -> `np.sin`). Instead of `ParcelsRandom` you should use numpy's random functions. - `particle.depth` has been changed to `particles.z` to be consistent with the [CF conventions for trajectory data](https://cfconventions.org/cf-conventions/cf-conventions.html#trajectory-data), and to make Parcels also generalizable to atmospheric contexts. - The `InteractionKernel` class has been removed. Since normal Kernels now have access to _all_ particles, particle-particle interaction can be performed within normal Kernels. +- Users need to explicitly use `convert_z_to_sigma_croco` in sampling kernels such as the `AdvectionRK4_3D_CROCO` kernel for CROCO fields, as the automatic conversion from depth to sigma grids under the hood has been removed. - We added a new AdvectionRK2 Kernel. The AdvectionRK4 kernel is still available, but RK2 is now the recommended default advection scheme as it is faster while the accuracy is comparable for most applications. See also the Choosing an integration method tutorial. ## FieldSet diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index ba7536de09..ff1e859c9b 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -162,23 +162,18 @@ def add_constant(self, name, value): name : str Name of the constant value : - Value of the constant (stored as 32-bit float) + Value of the constant - - Examples - -------- - Tutorials using fieldset.add_constant: - `Analytical advection <../examples/tutorial_analyticaladvection.ipynb>`__ - `Diffusion <../examples/tutorial_diffusion.ipynb>`__ - `Periodic boundaries <../examples/tutorial_periodic_boundaries.ipynb>`__ """ _assert_str_and_python_varname(name) if name in self.constants: raise ValueError(f"FieldSet already has a constant with name '{name}'") + if isinstance(value, xr.DataArray): + value = value.item() if not isinstance(value, (float, np.floating, int, np.integer)): raise ValueError(f"FieldSet constants have to be of type float or int, got a {type(value)}") - self.constants[name] = np.float32(value) + self.constants[name] = value @property def gridset(self) -> list[BaseGrid]: diff --git a/src/parcels/_core/index_search.py b/src/parcels/_core/index_search.py index b9b0ae497d..098e718857 100644 --- a/src/parcels/_core/index_search.py +++ b/src/parcels/_core/index_search.py @@ -104,22 +104,24 @@ def curvilinear_point_in_cell(grid, y: np.ndarray, x: np.ndarray, yi: np.ndarray ) px = np.array([grid.lon[yi, xi], grid.lon[yi, xi + 1], grid.lon[yi + 1, xi + 1], grid.lon[yi + 1, xi]]) - # Map grid and particle longitudes to [-180,180) - px = ((px + 180.0) % 360.0) - 180.0 - x = ((x + 180.0) % 360.0) - 180.0 - - # Create a mask for antimeridian cells - lon_span = px.max(axis=0) - px.min(axis=0) - antimeridian_cell = lon_span > 180.0 - - if np.any(antimeridian_cell): - # For any antimeridian cell ... - # If particle longitude is closer to 180.0, then negative cell longitudes need to be bumped by +360 - mask = (px < 0.0) & antimeridian_cell[np.newaxis, :] & (x[np.newaxis, :] > 0.0) - px[mask] += 360.0 - # If particle longitude is closer to -180.0, then positive cell longitudes need to be bumped by -360 - mask = (px > 0.0) & antimeridian_cell[np.newaxis, :] & (x[np.newaxis, :] < 0.0) - px[mask] -= 360.0 + + if grid._mesh == "spherical": + # Map grid and particle longitudes to [-180,180) + px = ((px + 180.0) % 360.0) - 180.0 + x = ((x + 180.0) % 360.0) - 180.0 + + # Create a mask for antimeridian cells + lon_span = px.max(axis=0) - px.min(axis=0) + antimeridian_cell = lon_span > 180.0 + + if np.any(antimeridian_cell): + # For any antimeridian cell ... + # If particle longitude is closer to 180.0, then negative cell longitudes need to be bumped by +360 + mask = (px < 0.0) & antimeridian_cell[np.newaxis, :] & (x[np.newaxis, :] > 0.0) + px[mask] += 360.0 + # If particle longitude is closer to -180.0, then positive cell longitudes need to be bumped by -360 + mask = (px > 0.0) & antimeridian_cell[np.newaxis, :] & (x[np.newaxis, :] < 0.0) + px[mask] -= 360.0 py = np.array([grid.lat[yi, xi], grid.lat[yi, xi + 1], grid.lat[yi + 1, xi + 1], grid.lat[yi + 1, xi]]) diff --git a/src/parcels/_core/kernel.py b/src/parcels/_core/kernel.py index 3fe8b717d2..0f6493e44a 100644 --- a/src/parcels/_core/kernel.py +++ b/src/parcels/_core/kernel.py @@ -80,10 +80,6 @@ def __init__( for f in pyfuncs: self.check_fieldsets_in_kernels(f) - # # TODO will be implemented when we support CROCO again - # if (pyfunc is AdvectionRK4_3D) and fieldset.U.gridindexingtype == "croco": - # pyfunc = AdvectionRK4_3D_CROCO - self._pyfuncs: list[Callable] = pyfuncs @property #! Ported from v3. To be removed in v4? (/find another way to name kernels in output file) diff --git a/src/parcels/convert.py b/src/parcels/convert.py index f66bc972bf..fea683670d 100644 --- a/src/parcels/convert.py +++ b/src/parcels/convert.py @@ -49,6 +49,12 @@ "T": "time", } +_CROCO_VARNAMES_MAPPING = { + "x_rho": "lon", + "y_rho": "lat", + "s_w": "depth", +} + def _maybe_bring_UV_depths_to_depth(ds): if "U" in ds.variables and "depthu" in ds.U.coords and "depth" in ds.coords: @@ -128,6 +134,29 @@ def _ds_rename_using_standard_names(ds: xr.Dataset | ux.UxDataset, name_dict: di return ds +def _maybe_convert_time_from_float_to_timedelta(ds: xr.Dataset) -> xr.Dataset: + if "time" in ds.coords: + if np.issubdtype(ds["time"].data.dtype, np.floating): + time_units = ds["time"].attrs.get("units", "").lower() + if "hour" in time_units: + factor = 3600.0 * 1e9 + elif "day" in time_units: + factor = 86400.0 * 1e9 + elif "minute" in time_units: + factor = 60.0 * 1e9 + else: + # default to seconds if unspecified + factor = 1.0 * 1e9 + + ns_int = np.rint(ds["time"].values * factor).astype("int64") + try: + ds["time"] = ns_int.astype("timedelta64[ns]") + logger.info("Converted time coordinate from float to timedelta based on units.") + except Exception as e: + logger.warning(f"Failed to convert time coordinate to timedelta: {e}") + return ds + + # TODO is this function still needed, now that we require users to provide field names explicitly? def _discover_U_and_V(ds: xr.Dataset, cf_standard_names_fallbacks) -> xr.Dataset: # Assumes that the dataset has U and V data @@ -266,6 +295,65 @@ def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Da return ds +def croco_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Dataset) -> xr.Dataset: + """Create an sgrid-compliant xarray.Dataset from a dataset of CROCO netcdf files. + + Parameters + ---------- + fields : dict[str, xr.Dataset | xr.DataArray] + Dictionary of xarray.DataArray objects as obtained from a set of Croco netcdf files. + coords : xarray.Dataset, optional + xarray.Dataset containing coordinate variables. + + Returns + ------- + xarray.Dataset + Dataset object following SGRID conventions to be (optionally) modified and passed to a FieldSet constructor. + + Notes + ----- + See the CROCO 3D tutorial for more information on how to use CROCO model outputs in Parcels + + """ + fields = fields.copy() + + for name, field_da in fields.items(): + if isinstance(field_da, xr.Dataset): + field_da = field_da[name] + # TODO: logging message, warn if multiple fields are in this dataset + else: + field_da = field_da.rename(name) + fields[name] = field_da + + ds = xr.merge(list(fields.values()) + [coords]) + ds.attrs.clear() # Clear global attributes from the merging + + ds = _maybe_rename_variables(ds, _CROCO_VARNAMES_MAPPING) + ds = _maybe_convert_time_from_float_to_timedelta(ds) + + if "grid" in ds.cf.cf_roles: + raise ValueError( + "Dataset already has a 'grid' variable (according to cf_roles). Didn't expect there to be grid metadata on copernicusmarine datasets - please open an issue with more information about your dataset." + ) + + ds["grid"] = xr.DataArray( + 0, + attrs=sgrid.Grid2DMetadata( + cf_role="grid_topology", + topology_dimension=2, + node_dimensions=("lon", "lat"), + node_coordinates=("lon", "lat"), + face_dimensions=( + sgrid.DimDimPadding("xi_u", "xi_rho", sgrid.Padding.HIGH), + sgrid.DimDimPadding("eta_v", "eta_rho", sgrid.Padding.HIGH), + ), + vertical_dimensions=(sgrid.DimDimPadding("s_rho", "depth", sgrid.Padding.HIGH),), + ).to_attrs(), + ) + + return ds + + def copernicusmarine_to_sgrid( *, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Dataset | None = None ) -> xr.Dataset: diff --git a/src/parcels/kernels/__init__.py b/src/parcels/kernels/__init__.py index de252b0030..f1c613086c 100644 --- a/src/parcels/kernels/__init__.py +++ b/src/parcels/kernels/__init__.py @@ -5,7 +5,6 @@ AdvectionRK2_3D, AdvectionRK4, AdvectionRK4_3D, - AdvectionRK4_3D_CROCO, AdvectionRK45, ) from ._advectiondiffusion import ( @@ -13,6 +12,11 @@ AdvectionDiffusionM1, DiffusionUniformKh, ) +from ._sigmagrids import ( + AdvectionRK4_3D_CROCO, + SampleOmegaCroco, + convert_z_to_sigma_croco, +) __all__ = [ # noqa: RUF022 # advection @@ -20,7 +24,6 @@ "AdvectionEE", "AdvectionRK2", "AdvectionRK2_3D", - "AdvectionRK4_3D_CROCO", "AdvectionRK4_3D", "AdvectionRK4", "AdvectionRK45", @@ -28,4 +31,8 @@ "AdvectionDiffusionEM", "AdvectionDiffusionM1", "DiffusionUniformKh", + # sigmagrids + "AdvectionRK4_3D_CROCO", + "SampleOmegaCroco", + "convert_z_to_sigma_croco", ] diff --git a/src/parcels/kernels/_advection.py b/src/parcels/kernels/_advection.py index e88d9d99c6..792c8ee17f 100644 --- a/src/parcels/kernels/_advection.py +++ b/src/parcels/kernels/_advection.py @@ -13,7 +13,6 @@ "AdvectionRK2_3D", "AdvectionRK4", "AdvectionRK4_3D", - "AdvectionRK4_3D_CROCO", "AdvectionRK45", ] @@ -90,48 +89,6 @@ def AdvectionRK4_3D(particles, fieldset): # pragma: no cover particles.dz += (w1 + 2 * w2 + 2 * w3 + w4) / 6 * dt -def AdvectionRK4_3D_CROCO(particles, fieldset): # pragma: no cover - """Advection of particles using fourth-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. - """ - dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt) - sig_dep = particles.z / fieldset.H[particles.time, 0, particles.lat, particles.lon] - - (u1, v1, w1) = fieldset.UVW[particles.time, particles.z, particles.lat, particles.lon, particles] - w1 *= sig_dep / fieldset.H[particles.time, 0, particles.lat, particles.lon] - lon1 = particles.lon + u1 * 0.5 * dt - lat1 = particles.lat + v1 * 0.5 * dt - sig_dep1 = sig_dep + w1 * 0.5 * dt - dep1 = sig_dep1 * fieldset.H[particles.time, 0, lat1, lon1] - - (u2, v2, w2) = fieldset.UVW[particles.time + 0.5 * dt, dep1, lat1, lon1, particles] - w2 *= sig_dep1 / fieldset.H[particles.time, 0, lat1, lon1] - lon2 = particles.lon + u2 * 0.5 * dt - lat2 = particles.lat + v2 * 0.5 * dt - sig_dep2 = sig_dep + w2 * 0.5 * dt - dep2 = sig_dep2 * fieldset.H[particles.time, 0, lat2, lon2] - - (u3, v3, w3) = fieldset.UVW[particles.time + 0.5 * dt, dep2, lat2, lon2, particles] - w3 *= sig_dep2 / fieldset.H[particles.time, 0, lat2, lon2] - lon3 = particles.lon + u3 * dt - lat3 = particles.lat + v3 * dt - sig_dep3 = sig_dep + w3 * dt - dep3 = sig_dep3 * fieldset.H[particles.time, 0, lat3, lon3] - - (u4, v4, w4) = fieldset.UVW[particles.time + dt, dep3, lat3, lon3, particles] - w4 *= sig_dep3 / fieldset.H[particles.time, 0, lat3, lon3] - lon4 = particles.lon + u4 * dt - lat4 = particles.lat + v4 * dt - sig_dep4 = sig_dep + w4 * dt - dep4 = sig_dep4 * fieldset.H[particles.time, 0, lat4, lon4] - - particles.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt - particles.dlat += (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 - - def AdvectionEE(particles, fieldset): # pragma: no cover """Advection of particles using Explicit Euler (aka Euler Forward) integration.""" (u1, v1) = fieldset.UV[particles] diff --git a/src/parcels/kernels/_sigmagrids.py b/src/parcels/kernels/_sigmagrids.py new file mode 100644 index 0000000000..5da1b81b20 --- /dev/null +++ b/src/parcels/kernels/_sigmagrids.py @@ -0,0 +1,87 @@ +import numpy as np + +from parcels.kernels._advection import _constrain_dt_to_within_time_interval + + +def convert_z_to_sigma_croco(fieldset, time, 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(time, np.zeros_like(z), y, x, particles=particle) + zeta = fieldset.zeta.eval(time, np.zeros_like(z), y, x, particles=particle) + sigma_levels = fieldset.U.grid.depth + cs_w = fieldset.Cs_w.data[0, :, 0, 0].values + + 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.time, particles.z, particles.lat, particles.lon, particles) + particles.omega = fieldset.omega[particles.time, sigma, particles.lat, particles.lon, 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. + 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. + """ + dt = _constrain_dt_to_within_time_interval(fieldset.time_interval, particles.time, particles.dt) + sigma = particles.z / fieldset.h[particles.time, 0, particles.lat, particles.lon] + + sig = convert_z_to_sigma_croco(fieldset, particles.time, particles.z, particles.lat, particles.lon, particles) + (u1, v1) = fieldset.UV[particles.time, sig, particles.lat, particles.lon, particles] + w1 = fieldset.W[particles.time, sig, particles.lat, particles.lon, particles] + w1 *= sigma / fieldset.h[particles.time, 0, particles.lat, particles.lon] + lon1 = particles.lon + u1 * 0.5 * dt + lat1 = particles.lat + v1 * 0.5 * dt + sig_dep1 = sigma + w1 * 0.5 * dt + dep1 = sig_dep1 * fieldset.h[particles.time, 0, lat1, lon1] + + sig1 = convert_z_to_sigma_croco(fieldset, particles.time + 0.5 * dt, dep1, lat1, lon1, particles) + (u2, v2) = fieldset.UV[particles.time + 0.5 * dt, sig1, lat1, lon1, particles] + w2 = fieldset.W[particles.time + 0.5 * dt, sig1, lat1, lon1, particles] + w2 *= sig_dep1 / fieldset.h[particles.time, 0, lat1, lon1] + lon2 = particles.lon + u2 * 0.5 * dt + lat2 = particles.lat + v2 * 0.5 * dt + sig_dep2 = sigma + w2 * 0.5 * dt + dep2 = sig_dep2 * fieldset.h[particles.time, 0, lat2, lon2] + + sig2 = convert_z_to_sigma_croco(fieldset, particles.time + 0.5 * dt, dep2, lat2, lon2, particles) + (u3, v3) = fieldset.UV[particles.time + 0.5 * dt, sig2, lat2, lon2, particles] + w3 = fieldset.W[particles.time + 0.5 * dt, sig2, lat2, lon2, particles] + w3 *= sig_dep2 / fieldset.h[particles.time, 0, lat2, lon2] + lon3 = particles.lon + u3 * dt + lat3 = particles.lat + v3 * dt + sig_dep3 = sigma + w3 * dt + dep3 = sig_dep3 * fieldset.h[particles.time, 0, lat3, lon3] + + sig3 = convert_z_to_sigma_croco(fieldset, particles.time + dt, dep3, lat3, lon3, particles) + (u4, v4) = fieldset.UV[particles.time + dt, sig3, lat3, lon3, particles] + w4 = fieldset.W[particles.time + dt, sig3, lat3, lon3, particles] + w4 *= sig_dep3 / fieldset.h[particles.time, 0, lat3, lon3] + lon4 = particles.lon + u4 * dt + lat4 = particles.lat + v4 * dt + sig_dep4 = sigma + w4 * dt + + dep4 = sig_dep4 * fieldset.h[particles.time, 0, lat4, lon4] + particles.dlon += (u1 + 2 * u2 + 2 * u3 + u4) / 6 * dt + particles.dlat += (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 diff --git a/tests-v3/test_advection.py b/tests-v3/test_advection.py index de8b004d79..3d8f06bac3 100644 --- a/tests-v3/test_advection.py +++ b/tests-v3/test_advection.py @@ -8,12 +8,10 @@ AdvectionDiffusionM1, AdvectionEE, AdvectionRK4, - AdvectionRK4_3D, AdvectionRK45, FieldSet, Particle, ParticleSet, - Variable, ) from tests.utils import TEST_DATA @@ -45,69 +43,6 @@ def depth(): return np.linspace(0, 30, zdim, dtype=np.float32) -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="When refactoring fieldfilebuffer croco support was dropped. This will be fixed in v4.") -def test_conversion_3DCROCO(): - """Test of the (SciPy) version 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] - ``` - """ - fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py") - - lat, lon = 78000.0, 38000.0 - 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, - ) - - sigma = np.zeros_like(z_xroms) - from parcels.field import _croco_from_z_to_sigma_scipy - - for zi, z in enumerate(z_xroms): - sigma[zi] = _croco_from_z_to_sigma_scipy(fieldset, 0, z, lat, lon, None) - - assert np.allclose(sigma, s_xroms, atol=1e-3) - - -@pytest.mark.v4alpha -@pytest.mark.xfail(reason="CROCO 3D interpolation is not yet implemented correctly in v4. ") -def test_advection_3DCROCO(): - fieldset = FieldSet.from_modulefile(TEST_DATA / "fieldset_CROCO3D.py") - - runtime = 1e4 - X, Z = np.meshgrid([40e3, 80e3, 120e3], [-10, -130]) - Y = np.ones(X.size) * 100e3 - - pclass = Particle.add_variable(Variable("w")) - pset = ParticleSet(fieldset=fieldset, pclass=pclass, lon=X, lat=Y, depth=Z) - - def SampleW(particle, fieldset, time): # pragma: no cover - particle.w = fieldset.W[time, particle.depth, particle.lat, particle.lon] - - pset.execute([AdvectionRK4_3D, SampleW], runtime=runtime, dt=100) - assert np.allclose(pset.depth, Z.flatten(), atol=5) # TODO lower this atol - assert np.allclose(pset.lon_nextloop, [x + runtime for x in X.flatten()], atol=1e-3) - - @pytest.mark.v4alpha @pytest.mark.xfail(reason="When refactoring fieldfilebuffer croco support was dropped. This will be fixed in v4.") def test_advection_2DCROCO(): diff --git a/tests/test_convert.py b/tests/test_convert.py index b80271a40a..0a555d82d7 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -6,6 +6,7 @@ from parcels import FieldSet from parcels._core.utils import sgrid from parcels._datasets.structured.circulation_models import datasets as datasets_circulation_models +from parcels.interpolators._xinterpolators import _get_offsets_dictionary def test_nemo_to_sgrid(): @@ -39,6 +40,21 @@ def test_nemo_to_sgrid(): }.issubset(set(ds["V"].dims)) +def test_convert_nemo_offsets(): + data_folder = parcels.download_example_dataset("NemoCurvilinear_data") + U = xr.open_mfdataset(data_folder.glob("*U.nc4")) + V = xr.open_mfdataset(data_folder.glob("*V.nc4")) + coords = xr.open_dataset(data_folder / "mesh_mask.nc4") + + ds = convert.nemo_to_sgrid(fields=dict(U=U, V=V), coords=coords) + fieldset = FieldSet.from_sgrid_conventions(ds) + + offsets = _get_offsets_dictionary(fieldset.UV.grid) + assert offsets["X"] == 1 + assert offsets["Y"] == 1 + assert offsets["Z"] == 0 + + _COPERNICUS_DATASETS = [ datasets_circulation_models["ds_copernicusmarine"], datasets_circulation_models["ds_copernicusmarine_waves"], @@ -83,3 +99,16 @@ def test_convert_copernicusmarine_no_logs(ds, caplog): assert "V" in fieldset.fields assert "UV" in fieldset.fields assert caplog.text == "" + + +def test_convert_croco_offsets(): + ds = datasets_circulation_models["ds_CROCO_idealized"] + coords = ds[["x_rho", "y_rho", "s_w", "time"]] + + ds = convert.croco_to_sgrid(fields={"U": ds["u"], "V": ds["v"]}, coords=coords) + fieldset = FieldSet.from_sgrid_conventions(ds) + + offsets = _get_offsets_dictionary(fieldset.UV.grid) + assert offsets["X"] == 0 + assert offsets["Y"] == 0 + assert offsets["Z"] == 0 diff --git a/tests/test_data/fieldset_CROCO2D.py b/tests/test_data/fieldset_CROCO2D.py deleted file mode 100644 index 74fe2c33fc..0000000000 --- a/tests/test_data/fieldset_CROCO2D.py +++ /dev/null @@ -1,23 +0,0 @@ -import os - -import parcels - - -def create_fieldset(): - example_dataset_folder = parcels.download_example_dataset("CROCOidealized_data") - file = os.path.join(example_dataset_folder, "CROCO_idealized.nc") - - variables = {"U": "u", "V": "v"} - dimensions = { - "U": {"lon": "x_rho", "lat": "y_rho", "time": "time"}, - "V": {"lon": "x_rho", "lat": "y_rho", "time": "time"}, - } - fieldset = parcels.FieldSet.from_croco( - file, - variables, - dimensions, - allow_time_extrapolation=True, - mesh="flat", - ) - - return fieldset diff --git a/tests/test_data/fieldset_CROCO3D.py b/tests/test_data/fieldset_CROCO3D.py deleted file mode 100644 index 14fd989d0b..0000000000 --- a/tests/test_data/fieldset_CROCO3D.py +++ /dev/null @@ -1,30 +0,0 @@ -import os - -import xarray as xr - -import parcels - - -def create_fieldset(): - example_dataset_folder = parcels.download_example_dataset("CROCOidealized_data") - file = os.path.join(example_dataset_folder, "CROCO_idealized.nc") - - variables = {"U": "u", "V": "v", "W": "w", "H": "h", "Zeta": "zeta", "Cs_w": "Cs_w"} - dimensions = { - "U": {"lon": "x_rho", "lat": "y_rho", "depth": "s_w", "time": "time"}, - "V": {"lon": "x_rho", "lat": "y_rho", "depth": "s_w", "time": "time"}, - "W": {"lon": "x_rho", "lat": "y_rho", "depth": "s_w", "time": "time"}, - "H": {"lon": "x_rho", "lat": "y_rho"}, - "Zeta": {"lon": "x_rho", "lat": "y_rho", "time": "time"}, - "Cs_w": {"depth": "s_w"}, - } - fieldset = parcels.FieldSet.from_croco( - file, - variables, - dimensions, - allow_time_extrapolation=True, - mesh="flat", - hc=xr.open_dataset(file).hc.values, - ) - - return fieldset diff --git a/tests/test_sigmagrids.py b/tests/test_sigmagrids.py new file mode 100644 index 0000000000..f30c11efa8 --- /dev/null +++ b/tests/test_sigmagrids.py @@ -0,0 +1,96 @@ +import numpy as np +import xarray as xr + +import parcels +from parcels import Particle, ParticleSet, Variable +from parcels.kernels import AdvectionRK4_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] + ``` + """ + data_folder = parcels.download_example_dataset("CROCOidealized_data") + ds_fields = xr.open_dataset(data_folder / "CROCO_idealized.nc") + 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"], + } + + coords = ds_fields[["x_rho", "y_rho", "s_w", "time"]] + ds_fset = parcels.convert.croco_to_sgrid(fields=fields, coords=coords) + + fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) + fieldset.add_constant("hc", ds_fields.hc) + + 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, + ) + + 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_3DCROCO(): + data_folder = parcels.download_example_dataset("CROCOidealized_data") + ds_fields = xr.open_dataset(data_folder / "CROCO_idealized.nc") + ds_fields.load() + + 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"], + } + + coords = ds_fields[["x_rho", "y_rho", "s_w", "time"]] + ds_fset = parcels.convert.croco_to_sgrid(fields=fields, coords=coords) + + fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset) + fieldset.add_constant("hc", ds_fields.hc) + + 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, lon=X, lat=Y, z=Z) + + pset.execute( + [AdvectionRK4_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.lon, [x + runtime for x in X.flatten()], atol=1e-3)