diff --git a/docs/user_guide/examples_v3/tutorial_croco_3D.ipynb b/docs/user_guide/examples_v3/tutorial_croco_3D.ipynb index d46b2b706a..5a3e499b8f 100644 --- a/docs/user_guide/examples_v3/tutorial_croco_3D.ipynb +++ b/docs/user_guide/examples_v3/tutorial_croco_3D.ipynb @@ -34,16 +34,16 @@ "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\")" + "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" ] }, { @@ -59,17 +59,16 @@ "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." + "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": [ - "
\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", - "
" + "```{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", + "```" ] }, { @@ -78,29 +77,12 @@ "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", + "vars_to_use = [\"u\", \"v\", \"w\", \"h\", \"zeta\", \"Cs_w\"]\n", + "ds_fields = ds_fields[vars_to_use]\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", - ")" + "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", + "fieldset = parcels.FieldSet.from_croco(ds_fields, mesh=mesh)\n", + "fieldset.add_constant(\"hc\", ds_fields.hc)" ] }, { @@ -120,24 +102,27 @@ " [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", + "Y = np.ones(X.size) * 98000\n", "\n", "\n", - "def DeleteParticle(particle, fieldset, time):\n", - " if particle.state >= 50:\n", - " particle.delete()\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, depth=Z\n", + " fieldset=fieldset, pclass=parcels.Particle, lon=X, lat=Y, z=Z\n", ")\n", "\n", - "outputfile = pset.ParticleFile(name=\"croco_particles3D.zarr\", outputdt=5000)\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, DeleteParticle],\n", - " runtime=5e4,\n", - " dt=100,\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", ")" ] @@ -162,15 +147,20 @@ "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", - "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", + "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 3D advection\")\n", + "ax.set_title(\"Particles in idealized CROCO velocity field using AdvectionRK4_3D_CROCO\")\n", "plt.tight_layout()\n", "plt.show()" ] @@ -186,7 +176,7 @@ "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." + "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." ] }, { @@ -195,32 +185,44 @@ "metadata": {}, "outputs": [], "source": [ - "import copy\n", - "\n", - "fieldset_noW = copy.copy(fieldset)\n", + "fieldset_noW = parcels.FieldSet.from_croco(ds_fields, mesh=mesh)\n", + "fieldset_noW.add_constant(\"hc\", ds_fields.hc)\n", "fieldset_noW.W.data[:] = 0.0\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, depth=Z\n", + " fieldset=fieldset_noW, pclass=parcels.Particle, lon=X, lat=Y, z=Z\n", ")\n", "\n", - "outputfile = pset.ParticleFile(name=\"croco_particles_noW.zarr\", outputdt=5000)\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, DeleteParticle],\n", - " runtime=5e4,\n", - " dt=100,\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", - "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", + "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", @@ -240,24 +242,13 @@ "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", + "When using `FieldSet.from_croco()`, 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", - "# Now interpolate the field to the sigma level\n", - "temp = fieldset.T[time, sigma, particle.lat, particle.lon]\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", "```" ] }, @@ -265,31 +256,35 @@ "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", + "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)\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", - "(u, v, w) = fieldset.UVW[time, particle.depth, particle.lat, particle.lon, particle] \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[time, particle.depth, particle.lat, particle.lon]\n", + "w_sigma = w * sigma / fieldset.h[particles.time, 0, particles.lat, particles.lon]\n", "\n", - "lon_new = particle.lon + u*particle.dt\n", - "lat_new = particle.lat + v*particle.dt\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*particle.dt \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[time, particle.depth, lat_new, lon_new]\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": "parcels", + "display_name": "docs", "language": "python", "name": "python3" }, @@ -303,7 +298,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.3" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index 7e76285be8..1fdbb08869 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -31,11 +31,11 @@ TODO: Add links to Reference API throughout examples/explanation_grids.md examples/tutorial_nemo_curvilinear.ipynb examples/tutorial_nemo_3D.ipynb +examples/tutorial_croco_3D.ipynb examples/tutorial_unitconverters.ipynb ``` - ```{toctree} diff --git a/docs/user_guide/v4-migration.md b/docs/user_guide/v4-migration.md index a5685e58e0..603a575164 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 2892779a1e..bb0d1d8b18 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -21,6 +21,7 @@ from parcels._logger import logger from parcels._typing import Mesh from parcels.interpolators import ( + CGrid_Tracer, CGrid_Velocity, UxPiecewiseConstantFace, UxPiecewiseLinearNode, @@ -316,6 +317,74 @@ def from_nemo(ds: xr.Dataset): fieldset.UVW.vector_interp_method = CGrid_Velocity return fieldset + def from_croco(ds: xr.Dataset, mesh="spherical"): + """Create a FieldSet from a xarray.Dataset from CROCO netcdf files. + + Parameters + ---------- + ds : xarray.Dataset + xarray.Dataset as obtained from a set of CROCO netcdf files. + mesh : str + String indicating the type of mesh coordinates and units used during the simulation. + Options are "spherical" or "flat". Default is "spherical". + + Returns + ------- + FieldSet + FieldSet object containing the fields from the dataset that can be used for a Parcels simulation. + + Notes + ----- + TODO + + """ + ds = ds.copy() + ds = _maybe_rename_variables(ds, _CROCO_VARNAMES_MAPPING) + ds = _discover_U_and_V(ds, _CROCO_CF_STANDARD_NAME_FALLBACKS) + ds = _assign_dims_as_coords(ds, _CROCO_DIMENSION_NAMES) + ds = _set_coords(ds, _CROCO_DIMENSION_NAMES) + ds = _set_axis_attrs(ds, _CROCO_AXIS_VARNAMES) + + expected_axes = set("XYZT") # TODO: Update after we have support for 2D spatial fields + if missing_axes := (expected_axes - set(ds.cf.axes)): + raise ValueError( + f"Dataset missing CF compliant metadata for axes " + f"{missing_axes}. Expected 'axis' attribute to be set " + f"on all dimension axes {expected_axes}. " + "HINT: Add xarray metadata attribute 'axis' to dimension - e.g., ds['lat'].attrs['axis'] = 'Y'" + ) + + 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"), + face_dimensions=( + sgrid.DimDimPadding("xi_u", "xi_rho", sgrid.Padding.BOTH), + sgrid.DimDimPadding("eta_v", "eta_rho", sgrid.Padding.BOTH), + ), + vertical_dimensions=(sgrid.DimDimPadding("s_rho", "depth", sgrid.Padding.NONE),), + ).to_attrs(), + ) + fieldset = FieldSet.from_sgrid_conventions(ds, mesh=mesh) + if "UV" in fieldset.fields: + fieldset.UV.vector_interp_method = CGrid_Velocity + if "UVW" in fieldset.fields: + fieldset.UVW.vector_interp_method = CGrid_Velocity + + # Reset interpolation method for fields that use XLinear + for field in fieldset.fields.values(): + if hasattr(field, "interp_method") and field.interp_method == XLinear: + if field.name not in ["h", "W"]: + field.interp_method = CGrid_Tracer + + return fieldset + def from_fesom2(ds: ux.UxDataset): """Create a FieldSet from a FESOM2 uxarray.UxDataset. @@ -540,6 +609,31 @@ def _format_calendar_error_message(field: Field, reference_datetime: TimeLike) - "wo": "W", } +_CROCO_CF_STANDARD_NAME_FALLBACKS = { + "UV": [ + ( + "sea_water_x_velocity_at_u_location", + "sea_water_y_velocity_at_v_location", + ), + ], + "W": ["upward_sea_water_velocity"], +} + +_CROCO_DIMENSION_NAMES = ["xi_rho", "xi_u", "eta_rho", "eta_v", "lon", "lat", "s_w", "s_rho", "time"] + +_CROCO_AXIS_VARNAMES = { + "X": "xi_rho", + "Y": "eta_rho", + "Z": "s_w", + "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: diff --git a/src/parcels/_core/index_search.py b/src/parcels/_core/index_search.py index a225621fcc..d7178c1f7d 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/kernels/__init__.py b/src/parcels/kernels/__init__.py index b4635869af..e6ad6f39a2 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,9 @@ AdvectionDiffusionM1, DiffusionUniformKh, ) +from .sigmagrids import ( + AdvectionRK4_3D_CROCO, +) __all__ = [ # noqa: RUF022 # advection @@ -20,7 +22,6 @@ "AdvectionEE", "AdvectionRK2", "AdvectionRK2_3D", - "AdvectionRK4_3D_CROCO", "AdvectionRK4_3D", "AdvectionRK4", "AdvectionRK45", @@ -28,4 +29,6 @@ "AdvectionDiffusionEM", "AdvectionDiffusionM1", "DiffusionUniformKh", + # sigmagrids + "AdvectionRK4_3D_CROCO", ] diff --git a/src/parcels/kernels/advection.py b/src/parcels/kernels/advection.py index 29666cba08..150d77900d 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..459d22b4db --- /dev/null +++ b/src/parcels/kernels/sigmagrids.py @@ -0,0 +1,93 @@ +import numpy as np + +from parcels.kernels.advection import _constrain_dt_to_within_time_interval + +__all__ = [ + "AdvectionRK4_3D_CROCO", + "SampleOmegaCroco", + "convert_z_to_sigma_croco", +] + + +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, applyConversion=False) + zeta = fieldset.zeta.eval(time, np.zeros_like(z), y, x, particles=particle, applyConversion=False) + 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_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..1be3503913 --- /dev/null +++ b/tests/test_sigmagrids.py @@ -0,0 +1,75 @@ +import numpy as np +import xarray as xr + +import parcels +from parcels import Particle, ParticleSet, Variable +from parcels.kernels import AdvectionRK4_3D_CROCO +from parcels.kernels.sigmagrids import 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") + ds_fields = ds_fields[["u", "v", "w", "h", "zeta", "Cs_w"]] + + fieldset = parcels.FieldSet.from_croco(ds_fields, mesh="flat") + 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 = ds_fields[["u", "v", "w", "h", "zeta", "Cs_w", "omega"]] + ds_fields.load() + + fieldset = parcels.FieldSet.from_croco(ds_fields, mesh="flat") + 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)