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 3e924ab6b1..370bbfc287 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -27,13 +27,13 @@ The tutorials written for Parcels v3 are currently being updated for Parcels v4. examples/explanation_grids.md examples/tutorial_nemo_curvilinear.ipynb examples/tutorial_nemo_3D.ipynb +examples/tutorial_croco_3D.ipynb examples/tutorial_velocityconversion.ipynb examples/tutorial_nestedgrids.ipynb examples/tutorial_summingfields.ipynb ``` - ```{toctree} diff --git a/docs/user_guide/v4-migration.md b/docs/user_guide/v4-migration.md index 6e2f4743df..3be852ae80 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. ## FieldSet diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index bed9245c88..63861b2d9e 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -160,23 +160,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 170871d92d..dca352c2e1 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)