Implementing convert.nemo_to_sgrid()

docs/user_guide/examples/tutorial_nemo_curvilinear.ipynb

@@ -64,47 +64,17 @@
    "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",
+    "# Drop some auxiliary dimensions - otherwise Parcels will complain\n",
+    "ds_fields = ds_fields.isel(time=0, z_a=0, z=0, drop=True)\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",
-    ")\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",
-    "U.units = parcels.GeographicPolar()\n",
-    "V.units = (\n",
-    "    parcels.GeographicPolar()\n",
-    ")  # U and V need GeographicPolar for  C-Grid interpolation to work correctly\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_nemo(ds_fields)"
    ]
   },
   {
@@ -122,7 +92,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",
@@ -152,8 +122,12 @@
    "outputs": [],
    "source": [
     "u, v = fieldset.UV.eval(\n",
-    "    np.array([0]), np.array([0]), np.array([20]), np.array([50]), apply_conversion=False\n",
-    ")  # do not convert m/s to deg/s\n",
+    "    time=np.array([0]),\n",
+    "    z=np.array([0]),\n",
+    "    y=np.array([40]),\n",
+    "    x=np.array([50]),\n",
+    "    apply_conversion=False,  # do not convert m/s to deg/s\n",
+    ")\n",
     "print(f\"(u, v) = ({u[0]:.3f}, {v[0]:.3f})\")\n",
     "assert np.isclose(u, 1.0, atol=1e-3)"
    ]

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()"
    ]
   }
  ],

docs/user_guide/index.md

@@ -30,12 +30,12 @@ TODO: Add links to Reference API throughout
 :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
 ```
 
 <!-- examples/documentation_indexing.ipynb -->
-<!-- examples/tutorial_nemo_3D.ipynb -->
 <!-- examples/tutorial_croco_3D.ipynb -->
 <!-- examples/tutorial_timevaryingdepthdimensions.ipynb -->