diff --git a/docs/user_guide/examples/tutorial_nemo_curvilinear.ipynb b/docs/user_guide/examples/tutorial_nemo_curvilinear.ipynb index 795743b86..ba26f3558 100644 --- a/docs/user_guide/examples/tutorial_nemo_curvilinear.ipynb +++ b/docs/user_guide/examples/tutorial_nemo_curvilinear.ipynb @@ -64,43 +64,19 @@ "outputs": [], "source": [ "data_folder = parcels.download_example_dataset(\"NemoCurvilinear_data\")\n", - "files = data_folder.glob(\"*.nc4\")\n", "ds_fields = xr.open_mfdataset(\n", - " files, combine=\"nested\", data_vars=\"minimal\", coords=\"minimal\", compat=\"override\"\n", + " data_folder.glob(\"*.nc4\"),\n", + " data_vars=\"minimal\",\n", + " coords=\"minimal\",\n", + " compat=\"override\",\n", ")\n", "\n", - "\n", - "# TODO: replace manual fieldset creation with FieldSet.from_nemo() once available\n", - "ds_fields = (\n", - " ds_fields.isel(time_counter=0, drop=True)\n", - " .isel(time=0, drop=True)\n", - " .isel(z_a=0, drop=True)\n", - " .rename({\"glamf\": \"lon\", \"gphif\": \"lat\", \"z\": \"depth\"})\n", - ")\n", - "\n", - "import xgcm\n", - "\n", - "xgcm_grid = xgcm.Grid(\n", - " ds_fields,\n", - " coords={\n", - " \"X\": {\"left\": \"x\"},\n", - " \"Y\": {\"left\": \"y\"},\n", - " },\n", - " periodic=False,\n", - " autoparse_metadata=False,\n", + "ds_coords = xr.open_dataset(data_folder / \"mesh_mask.nc4\", decode_times=False)\n", + "ds_fset = parcels.convert.nemo_to_sgrid(\n", + " fields=dict(U=ds_fields[\"U\"], V=ds_fields[\"V\"]), coords=ds_coords\n", ")\n", - "grid = parcels.XGrid(xgcm_grid, mesh=\"spherical\")\n", "\n", - "U = parcels.Field(\n", - " \"U\", ds_fields[\"U\"], grid, interp_method=parcels.interpolators.XLinear\n", - ")\n", - "V = parcels.Field(\n", - " \"V\", ds_fields[\"V\"], grid, interp_method=parcels.interpolators.XLinear\n", - ")\n", - "UV = parcels.VectorField(\n", - " \"UV\", U, V, vector_interp_method=parcels.interpolators.CGrid_Velocity\n", - ")\n", - "fieldset = parcels.FieldSet([U, V, UV])" + "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)" ] }, { @@ -118,7 +94,7 @@ "outputs": [], "source": [ "fig, ax = plt.subplots(1, 2, figsize=(10, 4))\n", - "pc1 = ds_fields.U.plot(cmap=\"viridis\", ax=ax[0], vmin=0)\n", + "pc1 = fieldset.U.data.plot(cmap=\"viridis\", ax=ax[0], vmin=0)\n", "pc2 = ax[1].pcolormesh(\n", " fieldset.U.grid.lon,\n", " fieldset.U.grid.lat,\n", @@ -295,7 +271,7 @@ ], "metadata": { "kernelspec": { - "display_name": "default", + "display_name": "test-notebooks-latest", "language": "python", "name": "python3" }, @@ -309,7 +285,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.14.2" + "version": "3.13.9" } }, "nbformat": 4, diff --git a/docs/user_guide/examples_v3/tutorial_nemo_3D.ipynb b/docs/user_guide/examples_v3/tutorial_nemo_3D.ipynb index 21af7ad97..a70f5a58e 100644 --- a/docs/user_guide/examples_v3/tutorial_nemo_3D.ipynb +++ b/docs/user_guide/examples_v3/tutorial_nemo_3D.ipynb @@ -18,26 +18,13 @@ "More information about C-grid interpolation can be found in [Delandmeter et al., 2019](https://www.geosci-model-dev-discuss.net/gmd-2018-339/).\n", "An example of such a discretisation is the NEMO model, which is one of the models supported in Parcels. A tutorial teaching how to [interpolate 2D data on a NEMO grid](https://docs.oceanparcels.org/en/latest/examples/tutorial_nemo_curvilinear.html) is available within Parcels.\n", "\n", - "Here, we focus on 3D fields. Basically, it is a straightforward extension of the 2D example, but it is very easy to do a mistake in the setup of the vertical discretisation that would affect the interpolation scheme.\n", - "\n", - "## Preliminary comments\n", - "\n", + "```{note}\n", "_How to know if your data is discretised on a C grid?_ The best way is to read the documentation that comes with the data. Alternatively, an easy check is to assess the coordinates of the U, V and W fields: for an A grid, U, V and W are distributed on the same nodes, such that the coordinates are the same. For a C grid, there is a shift of half a cell between the different variables.\n", + "```\n", "\n", - "_What about grid indexing?_ Since the C-grid variables are not located on the same nodes, there is not one obvious way to define the indexing, i.e. where is `u[k,j,i]` compared to `v[k,j,i]` and `w[k,j,i]`. In Parcels, we use the same notation as in NEMO: see [horizontal indexing](https://www.nemo-ocean.eu/doc/img360.png) and [vertical indexing](https://www.nemo-ocean.eu/doc/img362.png).\n", - "It is important that you check if your data is following the same notation. Otherwise, you should re-index your data properly (this can be done within Parcels, there is no need to regenerate new netcdf files).\n", - "\n", - "_What about the accuracy?_ By default in Parcels, particle coordinates (i.e. longitude, latitude and depth) are stored using single-precision `np.float32` numbers. The advantage of this is that it saves some memory resources for the computation. In some applications, especially where particles travel very close to the coast, the single precision accuracy can lead to uncontrolled particle beaching, due to numerical rounding errors. In such case, you may want to double the coordinate precision to `np.float64`. This can be done by adding the parameter `lonlatdepth_dtype=np.float64` to the ParticleSet constructor. Note that for C-grid fieldsets such as in NEMO, the coordinates precision is set by default to `np.float64`.\n", - "\n", - "## How to create a 3D NEMO `dimensions` dictionary?\n", - "\n", - "In the following, we will show how to create the `dimensions` dictionary for 3D NEMO simulations. What you require is a 'mesh_mask' file, which in our case is called `coordinates.nc` but in some other versions of NEMO has a different name. In any case, it will have to contain the variables `glamf`, `gphif` and `depthw`, which are the longitude, latitude and depth of the mesh nodes, respectively. Note that `depthw` is not part of the mesh_mask file, but is in the same file as the w data (`wfiles[0]`).\n", - "\n", - "For the C-grid interpolation in Parcels to work properly, it is important that `U`, `V` and `W` are on the same grid. \n", - "\n", - "All other tracers (e.g. `temperature`, `salinity`) should also be on this same grid. So even though these tracers are computed by NEMO on the T-points, Parcels expects them on the f-points (`glamf`, `gphif` and `depthw`). Parcels then under the hood makes sure the interpolation of these tracers is done correctly.\n", + "Here, we focus on 3D fields. Basically, it is a straightforward extension of the 2D example, but it is very easy to make a mistake in the setup of the vertical discretisation that would affect the interpolation scheme.\n", "\n", - "The code below is an example of how to create a 3D simulation with particles, starting in the mouth of the river Rhine at 1m depth, and advecting them through the North Sea using the `AdvectionRK4_3D`\n" + "For the C-grid interpolation in Parcels to work properly, it is important that `U`, `V` and `W` are on the same grid. All other tracers (e.g. `temperature`, `salinity`) should also be on this same grid. So even though these tracers are computed by NEMO on the T-points, Parcels expects them on the f-points (`glamf`, `gphif` and `depthw`). Parcels then under the hood makes sure the interpolation of these tracers is done correctly." ] }, { @@ -46,69 +33,18 @@ "metadata": {}, "outputs": [], "source": [ - "import warnings\n", - "from datetime import timedelta\n", - "from glob import glob\n", - "\n", "import matplotlib.pyplot as plt\n", + "import numpy as np\n", "import xarray as xr\n", "\n", - "import parcels\n", - "from parcels import FileWarning\n", - "\n", - "# Add a filter for the xarray decoding warning\n", - "warnings.simplefilter(\"ignore\", FileWarning)\n", - "\n", - "example_dataset_folder = parcels.download_example_dataset(\n", - " \"NemoNorthSeaORCA025-N006_data\"\n", - ")\n", - "ufiles = sorted(glob(f\"{example_dataset_folder}/ORCA*U.nc\"))\n", - "vfiles = sorted(glob(f\"{example_dataset_folder}/ORCA*V.nc\"))\n", - "wfiles = sorted(glob(f\"{example_dataset_folder}/ORCA*W.nc\"))\n", - "# tfiles = sorted(glob(f\"{example_dataset_folder}/ORCA*T.nc\")) # Not used in this example\n", - "mesh_mask = f\"{example_dataset_folder}/coordinates.nc\"\n", - "\n", - "filenames = {\n", - " \"U\": {\"lon\": mesh_mask, \"lat\": mesh_mask, \"depth\": wfiles[0], \"data\": ufiles},\n", - " \"V\": {\"lon\": mesh_mask, \"lat\": mesh_mask, \"depth\": wfiles[0], \"data\": vfiles},\n", - " \"W\": {\"lon\": mesh_mask, \"lat\": mesh_mask, \"depth\": wfiles[0], \"data\": wfiles},\n", - " # \"T\": {\"lon\": mesh_mask, \"lat\": mesh_mask, \"depth\": wfiles[0], \"data\": tfiles}, # Not used in this example\n", - "}\n", - "\n", - "variables = {\n", - " \"U\": \"uo\",\n", - " \"V\": \"vo\",\n", - " \"W\": \"wo\",\n", - " # \"T\": \"thetao\", # Not used in this example\n", - "}\n", - "\n", - "# Note that all variables need the same dimensions in a C-Grid\n", - "c_grid_dimensions = {\n", - " \"lon\": \"glamf\",\n", - " \"lat\": \"gphif\",\n", - " \"depth\": \"depthw\",\n", - " \"time\": \"time_counter\",\n", - "}\n", - "dimensions = {\n", - " \"U\": c_grid_dimensions,\n", - " \"V\": c_grid_dimensions,\n", - " \"W\": c_grid_dimensions,\n", - " # \"T\": c_grid_dimensions, # Not used in this example\n", - "}\n", - "\n", - "fieldset = parcels.FieldSet.from_nemo(filenames, variables, dimensions)\n", - "\n", - "pset = parcels.ParticleSet.from_line(\n", - " fieldset=fieldset,\n", - " pclass=parcels.Particle,\n", - " size=10,\n", - " start=(1.9, 52.5),\n", - " finish=(3.4, 51.6),\n", - " depth=1,\n", - ")\n", - "\n", - "kernels = pset.Kernel(parcels.kernels.AdvectionRK4_3D)\n", - "pset.execute(kernels, runtime=timedelta(days=4), dt=timedelta(hours=6))" + "import parcels" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Parcels v4 comes with a `from_nemo` method that automatically sets up the correct `Grid` for both 2D and 3D NEMO data. However, you may need to do some data-wrangling on your NEMO data files, such as below:" ] }, { @@ -117,47 +53,70 @@ "metadata": {}, "outputs": [], "source": [ - "depth_level = 8\n", - "print(\n", - " f\"Level[{int(depth_level)}] depth is: \"\n", - " f\"[{fieldset.W.grid.depth[depth_level]:g} \"\n", - " f\"{fieldset.W.grid.depth[depth_level + 1]:g}]\"\n", + "data_folder = parcels.download_example_dataset(\"NemoNorthSeaORCA025-N006_data\")\n", + "\n", + "# Open field datasets (with minimal loading to speed up)\n", + "ds_fields = xr.open_mfdataset(\n", + " data_folder.glob(\"ORCA*.nc\"),\n", + " data_vars=\"minimal\",\n", + " coords=\"minimal\",\n", + " compat=\"override\",\n", ")\n", "\n", - "plt.pcolormesh(\n", - " fieldset.U.grid.lon, fieldset.U.grid.lat, fieldset.U.data[0, depth_level, :, :]\n", - ")\n", - "plt.show()" + "# Open coordinates dataset (with no time decoding because xarray otherwise complains)\n", + "ds_coords = xr.open_dataset(data_folder / \"coordinates.nc\", decode_times=False)\n", + "\n", + "# Remove time dimension from coordinates (otherwise xarray complains)\n", + "ds_coords = ds_coords.isel(time=0, drop=True)\n", + "\n", + "# Combine field and coordinate datasets\n", + "ds = xr.merge([ds_fields, ds_coords[[\"glamf\", \"gphif\"]]])\n", + "\n", + "fieldset = parcels.FieldSet.from_nemo(ds)" ] }, { - "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ - "## Adding other fields like cell edges\n", - "\n", - "It is quite straightforward to add other gridded data, on the same curvilinear or any other type of grid, to the fieldset. Because it is good practice to make no changes to a `FieldSet` once a `ParticleSet` has been defined in it, we redefine the fieldset and add the fields with the cell edges from the coordinates file using `FieldSet.add_field()`.\n", - "\n", - "Note that including tracers like `temperature` and `salinity` needs to be done at the f-points, so on the same grid as the velocity fields. See also the section above.\n" + "The code below is an example of how to then create a 3D simulation with particles, starting on a meridional line through the North Sea from the mouth of the river Rhine at 100m depth, and advecting them through the North Sea using the `AdvectionRK2_3D`" ] }, { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "tags": [ + "hide-output" + ] + }, "outputs": [], "source": [ - "from parcels import Field" + "npart = 10\n", + "pset = parcels.ParticleSet(\n", + " fieldset=fieldset,\n", + " lon=np.linspace(1.9, 3.4, npart),\n", + " lat=np.linspace(65, 51.6, npart),\n", + " z=100 * np.ones(npart),\n", + ")\n", + "\n", + "pfile = parcels.ParticleFile(\n", + " store=\"output_nemo3D.zarr\", outputdt=np.timedelta64(1, \"D\")\n", + ")\n", + "\n", + "pset.execute(\n", + " parcels.kernels.AdvectionRK2_3D,\n", + " endtime=fieldset.time_interval.right,\n", + " dt=np.timedelta64(6, \"h\"),\n", + " output_file=pfile,\n", + ")" ] }, { - "cell_type": "code", - "execution_count": null, + "cell_type": "markdown", "metadata": {}, - "outputs": [], "source": [ - "fieldset = parcels.FieldSet.from_nemo(filenames, variables, dimensions)" + "We can then plot the trajectories on top of the surface U field" ] }, { @@ -166,34 +125,14 @@ "metadata": {}, "outputs": [], "source": [ - "e1u = Field.from_netcdf(\n", - " filenames=mesh_mask,\n", - " variable=\"e1u\",\n", - " dimensions={\"lon\": \"glamu\", \"lat\": \"gphiu\"},\n", - " interp_method=\"nearest\",\n", - ")\n", - "e2u = Field.from_netcdf(\n", - " filenames=mesh_mask,\n", - " variable=\"e2u\",\n", - " dimensions={\"lon\": \"glamu\", \"lat\": \"gphiu\"},\n", - " interp_method=\"nearest\",\n", - ")\n", - "e1v = Field.from_netcdf(\n", - " filenames=mesh_mask,\n", - " variable=\"e1v\",\n", - " dimensions={\"lon\": \"glamv\", \"lat\": \"gphiv\"},\n", - " interp_method=\"nearest\",\n", - ")\n", - "e2v = Field.from_netcdf(\n", - " filenames=mesh_mask,\n", - " variable=\"e2v\",\n", - " dimensions={\"lon\": \"glamv\", \"lat\": \"gphiv\"},\n", - " interp_method=\"nearest\",\n", - ")\n", - "fieldset.add_field(e1u)\n", - "fieldset.add_field(e2u)\n", - "fieldset.add_field(e1v)\n", - "fieldset.add_field(e2v)" + "field = fieldset.U.data[0, 0, :, :]\n", + "field = field.where(field != 0, np.nan) # Mask land values for better plotting\n", + "plt.pcolormesh(fieldset.U.grid.lon, fieldset.U.grid.lat, field, cmap=\"RdBu\")\n", + "\n", + "ds_out = xr.open_zarr(\"output_nemo3D.zarr\")\n", + "plt.scatter(ds_out.lon.T, ds_out.lat.T, c=-ds_out.z.T, marker=\".\")\n", + "plt.colorbar(label=\"Depth (m)\")\n", + "plt.show()" ] } ], diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index b445dcd1d..c97864c69 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -26,12 +26,12 @@ The tutorials written for Parcels v3 are currently being updated for Parcels v4. :titlesonly: examples/explanation_grids.md examples/tutorial_nemo_curvilinear.ipynb +examples_v3/tutorial_nemo_3D.ipynb # TODO move to examples folder examples/tutorial_unitconverters.ipynb examples/tutorial_nestedgrids.ipynb ``` - diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index f11241294..26565929f 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -20,7 +20,9 @@ from parcels._logger import logger from parcels._reprs import fieldset_repr from parcels._typing import Mesh +from parcels.convert import _discover_U_and_V, _ds_rename_using_standard_names, _maybe_rename_coords from parcels.interpolators import ( + CGrid_Velocity, Ux_Velocity, UxPiecewiseConstantFace, UxPiecewiseLinearNode, @@ -209,7 +211,7 @@ def from_copernicusmarine(cls, ds: xr.Dataset): """ ds = ds.copy() - ds = _discover_copernicusmarine_U_and_V(ds) + ds = _discover_U_and_V(ds, _COPERNICUS_MARINE_CF_STANDARD_NAME_FALLBACKS) 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( @@ -219,10 +221,10 @@ def from_copernicusmarine(cls, ds: xr.Dataset): "HINT: Add xarray metadata attribute 'axis' to dimension - e.g., ds['lat'].attrs['axis'] = 'Y'" ) - ds = _rename_coords_copernicusmarine(ds) + ds = _maybe_rename_coords(ds, _COPERNICUS_MARINE_AXIS_VARNAMES) if "W" in ds.data_vars: # Negate W to convert from up positive to down positive (as that's the direction of positive z) - ds["W"] -= ds["W"] + ds["W"].data *= -1 if "grid" in ds.cf.cf_roles: raise ValueError( @@ -286,7 +288,7 @@ def from_fesom2(cls, ds: ux.UxDataset): @classmethod def from_sgrid_conventions( - cls, ds: xr.Dataset, mesh: Mesh + cls, ds: xr.Dataset, mesh: Mesh | None = None ): # TODO: Update mesh to be discovered from the dataset metadata """Create a FieldSet from a dataset using SGRID convention metadata. @@ -316,6 +318,8 @@ def from_sgrid_conventions( See https://sgrid.github.io/ for more information on the SGRID conventions. """ ds = ds.copy() + if mesh is None: + mesh = _get_mesh_type_from_sgrid_dataset(ds) # Ensure time dimension has axis attribute if present if "time" in ds.dims and "time" in ds.coords: @@ -342,6 +346,11 @@ def from_sgrid_conventions( if "T" not in xgcm_kwargs["coords"]: xgcm_kwargs["coords"]["T"] = {"center": "time"} + if "lon" not in ds.coords or "lat" not in ds.coords: + node_dimensions = sgrid.load_mappings(ds.grid.node_dimensions) + ds["lon"] = ds[node_dimensions[0]] + ds["lat"] = ds[node_dimensions[1]] + # Create xgcm Grid object xgcm_grid = xgcm.Grid(ds, autoparse_metadata=False, **xgcm_kwargs, **_DEFAULT_XGCM_KWARGS) @@ -357,20 +366,22 @@ def from_sgrid_conventions( fields = {} if "U" in ds.data_vars and "V" in ds.data_vars: + vector_interp_method = XLinear_Velocity if _is_agrid(ds) else CGrid_Velocity fields["U"] = Field("U", ds["U"], grid, XLinear) fields["V"] = Field("V", ds["V"], grid, XLinear) - fields["UV"] = VectorField("UV", fields["U"], fields["V"], vector_interp_method=XLinear_Velocity) + fields["UV"] = VectorField("UV", fields["U"], fields["V"], vector_interp_method=vector_interp_method) if "W" in ds.data_vars: fields["W"] = Field("W", ds["W"], grid, XLinear) fields["UVW"] = VectorField( - "UVW", fields["U"], fields["V"], fields["W"], vector_interp_method=XLinear_Velocity + "UVW", fields["U"], fields["V"], fields["W"], vector_interp_method=vector_interp_method ) for varname in set(ds.data_vars) - set(fields.keys()) - skip_vars: fields[varname] = Field(varname, ds[varname], grid, XLinear) - return cls(list(fields.values())) + fieldset = cls(list(fields.values())) + return fieldset class CalendarError(Exception): # TODO: Move to a parcels errors module @@ -414,19 +425,8 @@ def _format_calendar_error_message(field: Field, reference_datetime: TimeLike) - } -def _rename_coords_copernicusmarine(ds): - try: - for axis, [coord] in ds.cf.axes.items(): - ds = ds.rename({coord: _COPERNICUS_MARINE_AXIS_VARNAMES[axis]}) - except ValueError as e: - raise ValueError(f"Multiple coordinates found for Copernicus dataset on axis '{axis}'. Check your data.") from e - return ds - - -def _discover_copernicusmarine_U_and_V(ds: xr.Dataset) -> xr.Dataset: - # Assumes that the dataset has U and V data - - cf_UV_standard_name_fallbacks = [ +_COPERNICUS_MARINE_CF_STANDARD_NAME_FALLBACKS = { + "UV": [ ( "eastward_sea_water_velocity", "northward_sea_water_velocity", @@ -448,36 +448,9 @@ def _discover_copernicusmarine_U_and_V(ds: xr.Dataset) -> xr.Dataset: "eastward_sea_water_velocity_vertical_mean_over_pelagic_layer", "northward_sea_water_velocity_vertical_mean_over_pelagic_layer", ), # GLOBAL_MULTIYEAR_BGC_001_033 - ] - cf_W_standard_name_fallbacks = ["upward_sea_water_velocity", "vertical_sea_water_velocity"] - - if "W" not in ds: - for cf_standard_name_W in cf_W_standard_name_fallbacks: - if cf_standard_name_W in ds.cf.standard_names: - ds = _ds_rename_using_standard_names(ds, {cf_standard_name_W: "W"}) - break - - if "U" in ds and "V" in ds: - return ds # U and V already present - elif "U" in ds or "V" in ds: - raise ValueError( - "Dataset has only one of the two variables 'U' and 'V'. Please rename the appropriate variable in your dataset to have both 'U' and 'V' for Parcels simulation." - ) - - for cf_standard_name_U, cf_standard_name_V in cf_UV_standard_name_fallbacks: - if cf_standard_name_U in ds.cf.standard_names: - if cf_standard_name_V not in ds.cf.standard_names: - raise ValueError( - f"Dataset has variable with CF standard name {cf_standard_name_U!r}, " - f"but not the matching variable with CF standard name {cf_standard_name_V!r}. " - "Please rename the appropriate variables in your dataset to have both 'U' and 'V' for Parcels simulation." - ) - else: - continue - - ds = _ds_rename_using_standard_names(ds, {cf_standard_name_U: "U", cf_standard_name_V: "V"}) - break - return ds + ], + "W": ["upward_sea_water_velocity", "vertical_sea_water_velocity"], +} def _discover_fesom2_U_and_V(ds: ux.UxDataset) -> ux.UxDataset: @@ -522,16 +495,6 @@ def _discover_fesom2_U_and_V(ds: ux.UxDataset) -> ux.UxDataset: return ds -def _ds_rename_using_standard_names(ds: xr.Dataset | ux.UxDataset, name_dict: dict[str, str]) -> xr.Dataset: - for standard_name, rename_to in name_dict.items(): - name = ds.cf[standard_name].name - ds = ds.rename({name: rename_to}) - logger.info( - f"cf_xarray found variable {name!r} with CF standard name {standard_name!r} in dataset, renamed it to {rename_to!r} for Parcels simulation." - ) - return ds - - def _select_uxinterpolator(da: ux.UxDataArray): """Selects the appropriate uxarray interpolator for a given uxarray UxDataArray""" supported_uxinterp_mapping = { @@ -563,3 +526,43 @@ def _select_uxinterpolator(da: ux.UxDataArray): return supported_uxinterp_mapping[key] return None + + +# TODO: Refactor later into something like `parcels._metadata.discover(dataset)` helper that can be used to discover important metadata like this. I think this whole metadata handling should be refactored into its own module. +def _get_mesh_type_from_sgrid_dataset(ds_sgrid: xr.Dataset) -> Mesh: + """Small helper to inspect SGRID metadata and dataset metadata to determine mesh type.""" + grid_da = sgrid.get_grid_topology(ds_sgrid) + if grid_da is None: + raise ValueError("Dataset does not contain SGRID grid topology metadata (cf_role='grid_topology').") + + sgrid_metadata = sgrid.parse_grid_attrs(grid_da.attrs) + + fpoint_x, fpoint_y = sgrid_metadata.node_coordinates + + if _is_coordinate_in_degrees(ds_sgrid[fpoint_x]) ^ _is_coordinate_in_degrees(ds_sgrid[fpoint_x]): + msg = ( + f"Mismatch in units between X and Y coordinates.\n" + f" Coordinate {ds_sgrid[fpoint_x]!r} attrs: {ds_sgrid[fpoint_x].attrs}\n" + f" Coordinate {ds_sgrid[fpoint_y]!r} attrs: {ds_sgrid[fpoint_y].attrs}\n" + ) + raise ValueError(msg) + + return "spherical" if _is_coordinate_in_degrees(ds_sgrid[fpoint_x]) else "flat" + + +def _is_agrid(ds: xr.Dataset) -> bool: + # check if U and V are defined on the same dimensions + # if yes, interpret as A grid + return set(ds["U"].dims) == set(ds["V"].dims) + + +def _is_coordinate_in_degrees(da: xr.DataArray) -> bool: + match da.attrs.get("units"): + case None: + raise ValueError( + f"Coordinate {da.name!r} of your dataset has no 'units' attribute - we don't know what the spatial units are." + ) + case "degrees": + return True + case _: + return False diff --git a/src/parcels/_core/particleset.py b/src/parcels/_core/particleset.py index 09bb56d86..c4cd5ffd8 100644 --- a/src/parcels/_core/particleset.py +++ b/src/parcels/_core/particleset.py @@ -544,6 +544,9 @@ def _get_simulation_start_and_end_times( if runtime is None and time_interval is None: raise ValueError("The runtime must be provided when the time_interval is not defined for a fieldset.") + if runtime is None and endtime is None: + raise ValueError("Either runtime or endtime must be provided.") + if sign_dt == 1: first_release_time = particle_release_times.min() else: diff --git a/src/parcels/_core/utils/sgrid.py b/src/parcels/_core/utils/sgrid.py index ff75c1a8f..9fd87e143 100644 --- a/src/parcels/_core/utils/sgrid.py +++ b/src/parcels/_core/utils/sgrid.py @@ -49,6 +49,30 @@ def to_attrs(self) -> dict[str, str | int]: ... def from_attrs(cls, d: dict[str, Hashable]) -> Self: ... +# Note that - for some optional attributes in the SGRID spec - these IDs are not available +# hence this isn't full coverage +_ID_FETCHERS_GRID2DMETADATA = { + "node_dimension1": lambda meta: meta.node_dimensions[0], + "node_dimension2": lambda meta: meta.node_dimensions[1], + "face_dimension1": lambda meta: meta.face_dimensions[0].dim1, + "face_dimension2": lambda meta: meta.face_dimensions[1].dim1, + "type1": lambda meta: meta.face_dimensions[0].padding, + "type2": lambda meta: meta.face_dimensions[1].padding, +} + +_ID_FETCHERS_GRID3DMETADATA = { + "node_dimension1": lambda meta: meta.node_dimensions[0], + "node_dimension2": lambda meta: meta.node_dimensions[1], + "node_dimension3": lambda meta: meta.node_dimensions[2], + "face_dimension1": lambda meta: meta.volume_dimensions[0].dim1, + "face_dimension2": lambda meta: meta.volume_dimensions[1].dim1, + "face_dimension3": lambda meta: meta.volume_dimensions[2].dim1, + "type1": lambda meta: meta.volume_dimensions[0].padding, + "type2": lambda meta: meta.volume_dimensions[1].padding, + "type3": lambda meta: meta.volume_dimensions[2].padding, +} + + class Grid2DMetadata(AttrsSerializable): def __init__( self, @@ -154,6 +178,19 @@ def to_attrs(self) -> dict[str, str | int]: def rename(self, names_dict: dict[str, str]) -> Self: return _metadata_rename(self, names_dict) + def get_value_by_id(self, id: str) -> str: + """In the SGRID specification for 2D grids, different parts of the spec are identified by different "ID"s. + + Easily extract the value for a given ID. + + Example + ------- + # Get padding 2 + >>> get_name_from_id("type2") + "low" + """ + return _ID_FETCHERS_GRID2DMETADATA[id](self) + class Grid3DMetadata(AttrsSerializable): def __init__( @@ -251,6 +288,19 @@ def to_attrs(self) -> dict[str, str | int]: def rename(self, dims_dict: dict[str, str]) -> Self: return _metadata_rename(self, dims_dict) + def get_value_by_id(self, id: str) -> str: + """In the SGRID specification for 3D grids, different parts of the spec are identified by different "ID"s. + + Easily extract the value for a given ID. + + Example + ------- + # Get padding 2 + >>> get_name_from_id("type2") + "low" + """ + return _ID_FETCHERS_GRID3DMETADATA[id](self) + @dataclass class DimDimPadding: diff --git a/src/parcels/convert.py b/src/parcels/convert.py new file mode 100644 index 000000000..554ef8629 --- /dev/null +++ b/src/parcels/convert.py @@ -0,0 +1,269 @@ +""" +This module provide a series of functions model outputs (which might be following their +own conventions) to metadata rich data (i.e., data following SGRID/UGRID as well as CF +conventions). Parcels needs this metadata rich data to discover grid geometries among other +things. + +These functions use knowledge about the model to attach any missing metadata. The functions +emit verbose messaging so that the user is kept in the loop. The returned output is an +Xarray dataset so that users can further provide any missing metadata that was unable to +be determined before they pass it to the FieldSet constructor. +""" + +from __future__ import annotations + +import typing + +import numpy as np +import xarray as xr + +from parcels._core.utils import sgrid +from parcels._logger import logger + +if typing.TYPE_CHECKING: + import uxarray as ux + +_NEMO_CF_STANDARD_NAME_FALLBACKS = { + "UV": [ + ( + "sea_water_x_velocity", + "sea_water_y_velocity", + ), + ], + "W": ["upward_sea_water_velocity", "vertical_sea_water_velocity"], +} + +_NEMO_DIMENSION_COORD_NAMES = ["x", "y", "time", "x", "x_center", "y", "y_center", "depth", "glamf", "gphif"] + +_NEMO_AXIS_VARNAMES = { + "x": "X", + "x_center": "X", + "y": "Y", + "y_center": "Y", + "depth": "Z", + "time": "T", +} + +_NEMO_VARNAMES_MAPPING = { + "time_counter": "time", + "depthw": "depth", + "uo": "U", + "vo": "V", + "wo": "W", +} + + +def _maybe_bring_UV_depths_to_depth(ds): + if "U" in ds.variables and "depthu" in ds.U.coords and "depth" in ds.coords: + ds["U"] = ds["U"].assign_coords(depthu=ds["depth"].values).rename({"depthu": "depth"}) + if "V" in ds.variables and "depthv" in ds.V.coords and "depth" in ds.coords: + ds["V"] = ds["V"].assign_coords(depthv=ds["depth"].values).rename({"depthv": "depth"}) + return ds + + +def _maybe_create_depth_dim(ds): + if "depth" not in ds.dims: + ds = ds.expand_dims({"depth": [0]}) + ds["depth"] = xr.DataArray([0], dims=["depth"]) + return ds + + +def _maybe_rename_coords(ds, axis_varnames): + try: + for axis, [coord] in ds.cf.axes.items(): + ds = ds.rename({coord: axis_varnames[axis]}) + except ValueError as e: + raise ValueError(f"Multiple coordinates found on axis '{axis}'. Check your DataSet.") from e + return ds + + +def _maybe_rename_variables(ds, varnames_mapping): + rename_dict = {old: new for old, new in varnames_mapping.items() if (old in ds.data_vars) or (old in ds.coords)} + if rename_dict: + ds = ds.rename(rename_dict) + return ds + + +def _assign_dims_as_coords(ds, dimension_names): + for axis in dimension_names: + if axis in ds.dims and axis not in ds.coords: + ds = ds.assign_coords({axis: np.arange(ds.sizes[axis])}) + return ds + + +def _drop_unused_dimensions_and_coords(ds, dimension_and_coord_names): + for dim in ds.dims: + if dim not in dimension_and_coord_names: + ds = ds.drop_dims(dim, errors="ignore") + for coord in ds.coords: + if coord not in dimension_and_coord_names: + ds = ds.drop_vars(coord, errors="ignore") + return ds + + +def _set_coords(ds, dimension_names): + for varname in dimension_names: + if varname in ds and varname not in ds.coords: + ds = ds.set_coords([varname]) + return ds + + +def _maybe_remove_depth_from_lonlat(ds): + for coord in ["glamf", "gphif"]: + if coord in ds.coords and "depth" in ds[coord].dims: + ds[coord] = ds[coord].squeeze("depth", drop=True) + return ds + + +def _set_axis_attrs(ds, dim_axis): + for dim, axis in dim_axis.items(): + ds[dim].attrs["axis"] = axis + return ds + + +def _ds_rename_using_standard_names(ds: xr.Dataset | ux.UxDataset, name_dict: dict[str, str]) -> xr.Dataset: + for standard_name, rename_to in name_dict.items(): + name = ds.cf[standard_name].name + ds = ds.rename({name: rename_to}) + logger.info( + f"cf_xarray found variable {name!r} with CF standard name {standard_name!r} in dataset, renamed it to {rename_to!r} for Parcels simulation." + ) + return ds + + +def _discover_U_and_V(ds: xr.Dataset, cf_standard_names_fallbacks) -> xr.Dataset: + # Assumes that the dataset has U and V data + + if "W" not in ds: + for cf_standard_name_W in cf_standard_names_fallbacks["W"]: + if cf_standard_name_W in ds.cf.standard_names: + ds = _ds_rename_using_standard_names(ds, {cf_standard_name_W: "W"}) + break + + if "U" in ds and "V" in ds: + return ds # U and V already present + elif "U" in ds or "V" in ds: + raise ValueError( + "Dataset has only one of the two variables 'U' and 'V'. Please rename the appropriate variable in your dataset to have both 'U' and 'V' for Parcels simulation." + ) + + for cf_standard_name_U, cf_standard_name_V in cf_standard_names_fallbacks["UV"]: + if cf_standard_name_U in ds.cf.standard_names: + if cf_standard_name_V not in ds.cf.standard_names: + raise ValueError( + f"Dataset has variable with CF standard name {cf_standard_name_U!r}, " + f"but not the matching variable with CF standard name {cf_standard_name_V!r}. " + "Please rename the appropriate variables in your dataset to have both 'U' and 'V' for Parcels simulation." + ) + else: + continue + + ds = _ds_rename_using_standard_names(ds, {cf_standard_name_U: "U", cf_standard_name_V: "V"}) + break + return ds + + +def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Dataset): + # TODO: Update docstring + """Create a FieldSet from a xarray.Dataset from NEMO netcdf files. + + Parameters + ---------- + ds : xarray.Dataset + xarray.Dataset as obtained from a set of NEMO netcdf files. + + Returns + ------- + xarray.Dataset + Dataset object following SGRID conventions to be (optionally) modified and passed to a FieldSet constructor. + + Notes + ----- + The NEMO model (https://www.nemo-ocean.eu/) is used by a variety of oceanographic institutions around the world. + Output from these models may differ subtly in terms of variable names and metadata conventions. + This function attempts to standardize these differences to create a Parcels FieldSet. + If you encounter issues with your specific NEMO dataset, please open an issue on the Parcels GitHub repository with details about your dataset. + + """ + fields = fields.copy() + coords = coords[["gphif", "glamf"]] + + 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) + + match name: + case "U": + field_da = field_da.rename({"y": "y_center"}) + case "V": + field_da = field_da.rename({"x": "x_center"}) + case _: + pass + field_da = field_da.reset_coords(drop=True) + fields[name] = field_da + + if "time" in coords.dims: + if coords.sizes["time"] != 1: + raise ValueError("Time dimension in coords must be length 1 (i.e., no time-varying grid).") + coords = coords.isel(time=0).drop("time") + if len(coords.dims) == 3: + for dim, len_ in coords.sizes.items(): + if len_ == 1: + # TODO: log statement about selecting along z dim of 1 + coords = coords.isel({dim: 0}) + if len(coords.dims) != 2: + raise ValueError("Expected coordsinates to be 2 dimensional") + + ds = xr.merge(list(fields.values()) + [coords]) + ds = _maybe_rename_variables(ds, _NEMO_VARNAMES_MAPPING) + ds = _discover_U_and_V(ds, _NEMO_CF_STANDARD_NAME_FALLBACKS) + ds = _maybe_create_depth_dim(ds) + ds = _maybe_bring_UV_depths_to_depth(ds) + ds = _drop_unused_dimensions_and_coords(ds, _NEMO_DIMENSION_COORD_NAMES) + ds = _assign_dims_as_coords(ds, _NEMO_DIMENSION_COORD_NAMES) + ds = _set_coords(ds, _NEMO_DIMENSION_COORD_NAMES) + ds = _maybe_remove_depth_from_lonlat(ds) + ds = _set_axis_attrs(ds, _NEMO_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 "W" in ds.data_vars: + # Negate W to convert from up positive to down positive (as that's the direction of positive z) + ds["W"].data *= -1 + 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=("x", "y"), + node_coordinates=("glamf", "gphif"), + face_dimensions=( + sgrid.DimDimPadding("x_center", "x", sgrid.Padding.LOW), + sgrid.DimDimPadding("y_center", "y", sgrid.Padding.LOW), + ), + vertical_dimensions=(sgrid.DimDimPadding("z_center", "depth", sgrid.Padding.HIGH),), + ).to_attrs(), + ) + + # NEMO models are always in degrees + ds["glamf"].attrs["units"] = "degrees" + ds["gphif"].attrs["units"] = "degrees" + + # Update to use lon and lat for internal naming + ds = sgrid.rename(ds, {"gphif": "lat", "glamf": "lon"}) # TODO: Logging message about rename + return ds diff --git a/src/parcels/interpolators.py b/src/parcels/interpolators.py index 61b58650b..52052ff85 100644 --- a/src/parcels/interpolators.py +++ b/src/parcels/interpolators.py @@ -99,6 +99,19 @@ def _get_corner_data_Agrid( return data.isel(selection_dict).data.reshape(lenT, lenZ, 2, 2, npart) +def _get_offsets_dictionary(grid): + offsets = {} + for axis in ["X", "Y"]: + axis_coords = grid.xgcm_grid.axes[axis].coords.keys() + offsets[axis] = 1 if "right" in axis_coords else 0 + if "Z" in grid.xgcm_grid.axes: + axis_coords = grid.xgcm_grid.axes["Z"].coords.keys() + offsets["Z"] = 1 if "right" in axis_coords else 0 + else: + offsets["Z"] = 0 + return offsets + + def XLinear( particle_positions: dict[str, float | np.ndarray], grid_positions: dict[_XGRID_AXES, dict[str, int | float | np.ndarray]], @@ -186,7 +199,9 @@ def CGrid_Velocity( U = vectorfield.U.data V = vectorfield.V.data grid = vectorfield.grid + offsets = _get_offsets_dictionary(grid) tdim, zdim, ydim, xdim = U.shape[0], U.shape[1], U.shape[2], U.shape[3] + lenT = 2 if np.any(tau > 0) else 1 if grid.lon.ndim == 1: px = np.array([grid.lon[xi], grid.lon[xi + 1], grid.lon[xi + 1], grid.lon[xi]]) @@ -212,82 +227,74 @@ def CGrid_Velocity( py[3], py[0], px[3], px[0], grid._mesh, np.einsum("ij,ji->i", i_u.phi2D_lin(eta, 0.0), py) ) - lenT = 2 if np.any(tau > 0) else 1 - - # Create arrays of corner points for xarray.isel - # TODO C grid may not need all xi and yi cornerpoints, so could speed up here? - - # Time coordinates: 4 points at ti, then 4 points at ti+1 - if lenT == 1: - ti_full = np.repeat(ti, 4) - else: - ti_1 = np.clip(ti + 1, 0, tdim - 1) - ti_full = np.concatenate([np.repeat(ti, 4), np.repeat(ti_1, 4)]) - - # Z coordinates: 4 points at zi, repeated for both time levels - zi_full = np.repeat(zi, lenT * 4) - - # Y coordinates: [yi, yi, yi+1, yi+1] for each spatial point, repeated for time/z - yi_1 = np.clip(yi + 1, 0, ydim - 1) - yi_full = np.tile(np.repeat(np.column_stack([yi, yi_1]), 2), (lenT)) - # # TODO check why in some cases minus needed here!!! - # yi_minus_1 = np.clip(yi - 1, 0, ydim - 1) - # yi = np.tile(np.repeat(np.column_stack([yi_minus_1, yi]), 2), (lenT)) - - # X coordinates: [xi, xi+1, xi, xi+1] for each spatial point, repeated for time/z - xi_1 = np.clip(xi + 1, 0, xdim - 1) - xi_full = np.tile(np.column_stack([xi, xi_1, xi, xi_1]).flatten(), (lenT)) - - for data in [U, V]: - axis_dim = grid.get_axis_dim_mapping(data.dims) - - # Create DataArrays for indexing + def _create_selection_dict(dims, zdir=False): + """Helper function to create DataArrays for indexing.""" + axis_dim = grid.get_axis_dim_mapping(dims) selection_dict = { axis_dim["X"]: xr.DataArray(xi_full, dims=("points")), axis_dim["Y"]: xr.DataArray(yi_full, dims=("points")), } + + # Time coordinates: 2 points at ti, then 2 points at ti+1 + if "time" in dims: + if lenT == 1: + ti_full = np.repeat(ti, 2) + else: + ti_1 = np.clip(ti + 1, 0, tdim - 1) + ti_full = np.concatenate([np.repeat(ti, 2), np.repeat(ti_1, 2)]) + selection_dict["time"] = xr.DataArray(ti_full, dims=("points")) + if "Z" in axis_dim: + if zdir: + # Z coordinates: 1 point at zi and 1 point at zi+1 repeated for lenT time levels + zi_0 = np.clip(zi + offsets["Z"], 0, zdim - 1) + zi_1 = np.clip(zi + offsets["Z"] + 1, 0, zdim - 1) + zi_full = np.tile(np.array([zi_0, zi_1]).flatten(), lenT) + else: + # Z coordinates: 2 points at zi, repeated for lenT time levels + zi_full = np.repeat(zi, lenT * 2) selection_dict[axis_dim["Z"]] = xr.DataArray(zi_full, dims=("points")) - if "time" in data.dims: - selection_dict["time"] = xr.DataArray(ti_full, dims=("points")) - corner_data = data.isel(selection_dict).data.reshape(lenT, len(xsi), 4) + return selection_dict + + def _compute_corner_data(data, selection_dict) -> np.ndarray: + """Helper function to load and reduce corner data over time dimension if needed.""" + corner_data = data.isel(selection_dict).data.reshape(lenT, 2, len(xsi)) if lenT == 2: - tau_full = tau[:, np.newaxis] - corner_data = corner_data[0, :, :] * (1 - tau_full) + corner_data[1, :, :] * tau_full + tau_full = tau[np.newaxis, :] + corner_data = corner_data[0, :] * (1 - tau_full) + corner_data[1, :] * tau_full else: - corner_data = corner_data[0, :, :] - # # See code below for v3 version - # # if self.gridindexingtype == "nemo": - # # U0 = self.U.data[ti, zi, yi + 1, xi] * c4 - # # U1 = self.U.data[ti, zi, yi + 1, xi + 1] * c2 - # # V0 = self.V.data[ti, zi, yi, xi + 1] * c1 - # # V1 = self.V.data[ti, zi, yi + 1, xi + 1] * c3 - # # elif self.gridindexingtype in ["mitgcm", "croco"]: - # # U0 = self.U.data[ti, zi, yi, xi] * c4 - # # U1 = self.U.data[ti, zi, yi, xi + 1] * c2 - # # V0 = self.V.data[ti, zi, yi, xi] * c1 - # # V1 = self.V.data[ti, zi, yi + 1, xi] * c3 - # # TODO Nick can you help use xgcm to fix this implementation? - - # # CROCO and MITgcm grid indexing, - # if data is U: - # U0 = corner_data[:, 0] * c4 - # U1 = corner_data[:, 1] * c2 - # elif data is V: - # V0 = corner_data[:, 0] * c1 - # V1 = corner_data[:, 2] * c3 - # # NEMO grid indexing - if data is U: - U0 = corner_data[:, 2] * c4 - U1 = corner_data[:, 3] * c2 - elif data is V: - V0 = corner_data[:, 1] * c1 - V1 = corner_data[:, 3] * c3 - - U = (1 - xsi) * U0 + xsi * U1 - V = (1 - eta) * V0 + eta * V1 + corner_data = corner_data[0, :] + return corner_data + + # Compute U velocity + yi_o = np.clip(yi + offsets["Y"], 0, ydim - 1) + yi_full = np.tile(np.array([yi_o, yi_o]).flatten(), lenT) + + xi_1 = np.clip(xi + 1, 0, xdim - 1) + xi_full = np.tile(np.array([xi, xi_1]).flatten(), lenT) + + selection_dict = _create_selection_dict(U.dims) + corner_data = _compute_corner_data(U, selection_dict) + + U0 = corner_data[0, :] * c4 + U1 = corner_data[1, :] * c2 + Uvel = (1 - xsi) * U0 + xsi * U1 + + # Compute V velocity + yi_1 = np.clip(yi + 1, 0, ydim - 1) + yi_full = np.tile(np.array([yi, yi_1]).flatten(), lenT) + + xi_o = np.clip(xi + offsets["X"], 0, xdim - 1) + xi_full = np.tile(np.array([xi_o, xi_o]).flatten(), lenT) + + selection_dict = _create_selection_dict(V.dims) + corner_data = _compute_corner_data(V, selection_dict) + + V0 = corner_data[0, :] * c1 + V1 = corner_data[1, :] * c3 + Vvel = (1 - eta) * V0 + eta * V1 if grid._mesh == "spherical": jac = i_u._compute_jacobian_determinant(py, px, eta, xsi) * 1852 * 60.0 @@ -295,16 +302,16 @@ def CGrid_Velocity( jac = i_u._compute_jacobian_determinant(py, px, eta, xsi) u = ( - (-(1 - eta) * U - (1 - xsi) * V) * px[0] - + ((1 - eta) * U - xsi * V) * px[1] - + (eta * U + xsi * V) * px[2] - + (-eta * U + (1 - xsi) * V) * px[3] + (-(1 - eta) * Uvel - (1 - xsi) * Vvel) * px[0] + + ((1 - eta) * Uvel - xsi * Vvel) * px[1] + + (eta * Uvel + xsi * Vvel) * px[2] + + (-eta * Uvel + (1 - xsi) * Vvel) * px[3] ) / jac v = ( - (-(1 - eta) * U - (1 - xsi) * V) * py[0] - + ((1 - eta) * U - xsi * V) * py[1] - + (eta * U + xsi * V) * py[2] - + (-eta * U + (1 - xsi) * V) * py[3] + (-(1 - eta) * Uvel - (1 - xsi) * Vvel) * py[0] + + ((1 - eta) * Uvel - xsi * Vvel) * py[1] + + (eta * Uvel + xsi * Vvel) * py[2] + + (-eta * Uvel + (1 - xsi) * Vvel) * py[3] ) / jac if is_dask_collection(u): u = u.compute() @@ -323,44 +330,18 @@ def CGrid_Velocity( u = np.where(np.abs(dlon / lon) > 1e-4, np.nan, u) if vectorfield.W: - data = vectorfield.W.data - # Time coordinates: 2 points at ti, then 2 points at ti+1 - if lenT == 1: - ti_full = np.repeat(ti, 2) - else: - ti_1 = np.clip(ti + 1, 0, tdim - 1) - ti_full = np.concatenate([np.repeat(ti, 2), np.repeat(ti_1, 2)]) - - # Z coordinates: 1 points at zi, repeated for both time levels - zi_1 = np.clip(zi + 1, 0, zdim - 1) - zi_full = np.tile(np.array([zi, zi_1]).flatten(), lenT) - - # Y coordinates: yi+1 for each spatial point, repeated for time/z - yi_1 = np.clip(yi + 1, 0, ydim - 1) - yi_full = np.tile(yi_1, (lenT) * 2) + W = vectorfield.W.data - # X coordinates: xi+1 for each spatial point, repeated for time/z - xi_1 = np.clip(xi + 1, 0, xdim - 1) - xi_full = np.tile(xi_1, (lenT) * 2) + # Y coordinates: yi+offset for each spatial point, repeated for time + yi_o = np.clip(yi + offsets["Y"], 0, ydim - 1) + yi_full = np.tile(yi_o, (lenT) * 2) - axis_dim = grid.get_axis_dim_mapping(data.dims) + # X coordinates: xi+offset for each spatial point, repeated for time + xi_o = np.clip(xi + offsets["X"], 0, xdim - 1) + xi_full = np.tile(xi_o, (lenT) * 2) - # Create DataArrays for indexing - selection_dict = { - axis_dim["X"]: xr.DataArray(xi_full, dims=("points")), - axis_dim["Y"]: xr.DataArray(yi_full, dims=("points")), - axis_dim["Z"]: xr.DataArray(zi_full, dims=("points")), - } - if "time" in data.dims: - selection_dict["time"] = xr.DataArray(ti_full, dims=("points")) - - corner_data = data.isel(selection_dict).data.reshape(lenT, 2, len(xsi)) - - if lenT == 2: - tau_full = tau[np.newaxis, :] - corner_data = corner_data[0, :, :] * (1 - tau_full) + corner_data[1, :, :] * tau_full - else: - corner_data = corner_data[0, :, :] + selection_dict = _create_selection_dict(W.dims, zdir=True) + corner_data = _compute_corner_data(W, selection_dict) w = corner_data[0, :] * (1 - zeta) + corner_data[1, :] * zeta if is_dask_collection(w): @@ -390,17 +371,19 @@ def CGrid_Tracer( axis_dim = field.grid.get_axis_dim_mapping(field.data.dims) data = field.data + offsets = _get_offsets_dictionary(field.grid) + zi = np.clip(zi + offsets["Z"], 0, data.shape[1] - 1) + yi = np.clip(yi + offsets["Y"], 0, data.shape[2] - 1) + xi = np.clip(xi + offsets["X"], 0, data.shape[3] - 1) + lenT = 2 if np.any(tau > 0) else 1 if lenT == 2: ti_1 = np.clip(ti + 1, 0, data.shape[0] - 1) ti = np.concatenate([np.repeat(ti), np.repeat(ti_1)]) - zi_1 = np.clip(zi + 1, 0, data.shape[1] - 1) - zi = np.concatenate([np.repeat(zi), np.repeat(zi_1)]) - yi_1 = np.clip(yi + 1, 0, data.shape[2] - 1) - yi = np.concatenate([np.repeat(yi), np.repeat(yi_1)]) - xi_1 = np.clip(xi + 1, 0, data.shape[3] - 1) - xi = np.concatenate([np.repeat(xi), np.repeat(xi_1)]) + zi = np.tile(zi, (lenT) * 2) + yi = np.tile(yi, (lenT) * 2) + xi = np.tile(xi, (lenT) * 2) # Create DataArrays for indexing selection_dict = { diff --git a/tests/test_advection.py b/tests/test_advection.py index cc94ccd28..671c78d71 100644 --- a/tests/test_advection.py +++ b/tests/test_advection.py @@ -1,7 +1,6 @@ import numpy as np import pytest import xarray as xr -import xgcm import parcels from parcels import Field, FieldSet, Particle, ParticleFile, ParticleSet, StatusCode, Variable, VectorField, XGrid @@ -25,7 +24,7 @@ AdvectionRK4_3D, AdvectionRK45, ) -from tests.utils import DEFAULT_PARTICLES, round_and_hash_float_array +from tests.utils import DEFAULT_PARTICLES @pytest.mark.parametrize("mesh", ["spherical", "flat"]) @@ -446,30 +445,13 @@ def UpdateP(particles, fieldset): # pragma: no cover def test_nemo_curvilinear_fieldset(): data_folder = parcels.download_example_dataset("NemoCurvilinear_data") - files = data_folder.glob("*.nc4") - ds = xr.open_mfdataset(files, combine="nested", data_vars="minimal", coords="minimal", compat="override") - ds = ( - ds.isel(time_counter=0, drop=True) - .isel(time=0, drop=True) - .isel(z_a=0, drop=True) - .rename({"glamf": "lon", "gphif": "lat", "z": "depth"}) - ) + 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") - xgcm_grid = xgcm.Grid( - ds, - coords={ - "X": {"left": "x"}, - "Y": {"left": "y"}, - }, - periodic=False, - autoparse_metadata=False, - ) - grid = XGrid(xgcm_grid, mesh="spherical") + ds = parcels.convert.nemo_to_sgrid(fields=dict(U=U, V=V), coords=coords) - U = parcels.Field("U", ds["U"], grid, interp_method=XLinear) - V = parcels.Field("V", ds["V"], grid, interp_method=XLinear) - UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity) - fieldset = parcels.FieldSet([U, V, UV]) + fieldset = parcels.FieldSet.from_sgrid_conventions(ds) npart = 20 lonp = 30 * np.ones(npart) @@ -481,87 +463,37 @@ def test_nemo_curvilinear_fieldset(): np.testing.assert_allclose(pset.lat, latp, atol=1e-1) -@pytest.mark.parametrize("kernel", [AdvectionRK4, AdvectionRK4_3D]) +@pytest.mark.parametrize( + "kernel", + [ + AdvectionRK4, + AdvectionRK4_3D, + ], +) def test_nemo_3D_curvilinear_fieldset(kernel): - download_dir = parcels.download_example_dataset("NemoNorthSeaORCA025-N006_data") - ufiles = download_dir.glob("*U.nc") - dsu = xr.open_mfdataset(ufiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) - dsu = dsu.rename({"time_counter": "time", "uo": "U"}) - - vfiles = download_dir.glob("*V.nc") - dsv = xr.open_mfdataset(vfiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) - dsv = dsv.rename({"time_counter": "time", "vo": "V"}) - - wfiles = download_dir.glob("*W.nc") - dsw = xr.open_mfdataset(wfiles, decode_times=False, drop_variables=["nav_lat", "nav_lon"]) - dsw = dsw.rename({"time_counter": "time", "depthw": "depth", "wo": "W"}) - - dsu = dsu.assign_coords(depthu=dsw.depth.values) - dsu = dsu.rename({"depthu": "depth"}) - - dsv = dsv.assign_coords(depthv=dsw.depth.values) - dsv = dsv.rename({"depthv": "depth"}) - - coord_file = f"{download_dir}/coordinates.nc" - dscoord = xr.open_dataset(coord_file, decode_times=False).rename({"glamf": "lon", "gphif": "lat"}) - dscoord = dscoord.isel(time=0, drop=True) - - ds = xr.merge([dsu, dsv, dsw, dscoord]) - ds = ds.drop_vars( - [ - "uos", - "vos", - "nav_lev", - "nav_lon", - "nav_lat", - "tauvo", - "tauuo", - "time_steps", - "gphiu", - "gphiv", - "gphit", - "glamu", - "glamv", - "glamt", - "time_centered_bounds", - "time_counter_bounds", - "time_centered", - ] - ) - ds = ds.drop_vars(["e1f", "e1t", "e1u", "e1v", "e2f", "e2t", "e2u", "e2v"]) - ds["time"] = [np.timedelta64(int(t), "s") + np.datetime64("1900-01-01") for t in ds["time"]] + data_folder = parcels.download_example_dataset("NemoNorthSeaORCA025-N006_data") + U = xr.open_mfdataset(data_folder.glob("*U.nc")) + V = xr.open_mfdataset(data_folder.glob("*V.nc")) + W = xr.open_mfdataset(data_folder.glob("*W.nc")) + coords = xr.open_dataset(data_folder / "coordinates.nc", decode_times=False) - ds["W"] *= -1 # Invert W velocity + ds = parcels.convert.nemo_to_sgrid(fields=dict(U=U["uo"], V=V["vo"], W=W["wo"]), coords=coords) - xgcm_grid = xgcm.Grid( - ds, - coords={ - "X": {"left": "x"}, - "Y": {"left": "y"}, - "Z": {"left": "depth"}, - "T": {"center": "time"}, - }, - periodic=False, - autoparse_metadata=False, - ) - grid = XGrid(xgcm_grid, mesh="spherical") - - U = parcels.Field("U", ds["U"], grid, interp_method=XLinear) - V = parcels.Field("V", ds["V"], grid, interp_method=XLinear) - W = parcels.Field("W", ds["W"], grid, interp_method=XLinear) - UV = parcels.VectorField("UV", U, V, vector_interp_method=CGrid_Velocity) - UVW = parcels.VectorField("UVW", U, V, W, vector_interp_method=CGrid_Velocity) - fieldset = parcels.FieldSet([U, V, W, UV, UVW]) + fieldset = parcels.FieldSet.from_sgrid_conventions(ds) npart = 10 - lons = np.linspace(1.9, 3.4, npart) - lats = np.linspace(52.5, 51.6, npart) - pset = parcels.ParticleSet(fieldset, lon=lons, lat=lats, z=np.ones_like(lons)) + lons_initial = np.linspace(1.9, 3.4, npart) + lats_initial = np.linspace(52.5, 51.6, npart) + z_initial = np.ones_like(lons_initial) + pset = parcels.ParticleSet(fieldset, lon=lons_initial, lat=lats_initial, z=z_initial) pset.execute(kernel, runtime=np.timedelta64(3, "D") + np.timedelta64(18, "h"), dt=np.timedelta64(6, "h")) if kernel == AdvectionRK4: - np.testing.assert_equal(round_and_hash_float_array([p.lon for p in pset], decimals=5), 29977383852960156017546) + np.testing.assert_allclose([p.z for p in pset], z_initial) elif kernel == AdvectionRK4_3D: # TODO check why decimals needs to be so low in RK4_3D (compare to v3) - np.testing.assert_equal(round_and_hash_float_array([p.z for p in pset], decimals=1), 29747210774230389239432) + np.testing.assert_allclose( + [p.z for p in pset], + [0.666162, 0.8667131, 0.92150104, 0.9605109, 0.9577529, 1.0041442, 1.0284728, 1.0033542, 1.2949713, 1.3928112], + ) # fmt:skip diff --git a/tests/test_convert.py b/tests/test_convert.py new file mode 100644 index 000000000..eaf39d810 --- /dev/null +++ b/tests/test_convert.py @@ -0,0 +1,36 @@ +import xarray as xr + +import parcels +import parcels.convert as convert +from parcels._core.utils import sgrid + + +def test_nemo_to_sgrid(): + 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) + + assert ds["grid"].attrs == { + "cf_role": "grid_topology", + "topology_dimension": 2, + "node_dimensions": "x y", + "face_dimensions": "x_center:x (padding:low) y_center:y (padding:low)", + "node_coordinates": "lon lat", + "vertical_dimensions": "z_center:depth (padding:high)", + } + + meta = sgrid.parse_grid_attrs(ds["grid"].attrs) + + # Assuming that node_dimension1 and node_dimension2 correspond to X and Y respectively + # check that U and V are properly defined on the staggered grid + assert { + meta.get_value_by_id("node_dimension1"), # X edge + meta.get_value_by_id("face_dimension2"), # Y center + }.issubset(set(ds["U"].dims)) + assert { + meta.get_value_by_id("face_dimension1"), # X center + meta.get_value_by_id("node_dimension2"), # Y edge + }.issubset(set(ds["V"].dims)) diff --git a/tests/utils/test_sgrid.py b/tests/utils/test_sgrid.py index e7de520ca..c3f3d666a 100644 --- a/tests/utils/test_sgrid.py +++ b/tests/utils/test_sgrid.py @@ -52,6 +52,30 @@ def create_example_grid3dmetadata(with_node_coordinates: bool): grid3dmetadata = create_example_grid3dmetadata(with_node_coordinates=True) +@pytest.mark.parametrize( + ("sgrid_metadata", "id", "value"), + [ + (grid2dmetadata, "node_dimension1", "node_dimension1"), + (grid2dmetadata, "node_dimension2", "node_dimension2"), + (grid2dmetadata, "face_dimension1", "face_dimension1"), + (grid2dmetadata, "face_dimension2", "face_dimension2"), + # (grid2dmetadata, "vertical_dimensions_dim1", "vertical_dimensions_dim1"), #! ID's NOT IMPLEMENTED IN SGRID SPEC + # (grid2dmetadata, "vertical_dimensions_dim2", "vertical_dimensions_dim2"), + (grid3dmetadata, "node_dimension1", "node_dimension1"), + (grid3dmetadata, "node_dimension2", "node_dimension2"), + (grid3dmetadata, "node_dimension3", "node_dimension3"), + (grid3dmetadata, "face_dimension1", "face_dimension1"), + (grid3dmetadata, "face_dimension2", "face_dimension2"), + (grid3dmetadata, "face_dimension3", "face_dimension3"), + (grid3dmetadata, "type1", sgrid.Padding.LOW), + (grid3dmetadata, "type2", sgrid.Padding.LOW), + (grid3dmetadata, "type3", sgrid.Padding.LOW), + ], +) +def test_get_value_by_id(sgrid_metadata: sgrid.Grid2DMetadata | sgrid.Grid3DMetadata, id, value): + assert sgrid_metadata.get_value_by_id(id) == value + + def dummy_sgrid_ds(grid: sgrid.Grid2DMetadata | sgrid.Grid3DMetadata) -> xr.Dataset: if isinstance(grid, sgrid.Grid2DMetadata): return dummy_sgrid_2d_ds(grid)