Creating convert.mitgcm_to_sgrid

docs/user_guide/examples/tutorial_mitgcm.ipynb

@@ -0,0 +1,155 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# 🖥️ MITgcm tutorial"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
+   "source": [
+    "This tutorial will show how to run a 3D simulation with output from the MITgcm model. The key thing about MITgcm is that its output is on a C-grid, which means that the U and V velocity components are not defined at the same location as the particle positions. Parcels has built-in functionality to deal with C-grid data. See also the [Nemo tutorial](tutorial_nemo.ipynb) for another example of C-grid data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import xarray as xr\n",
+    "\n",
+    "import parcels\n",
+    "\n",
+    "data_folder = parcels.download_example_dataset(\"MITgcm_example_data\")\n",
+    "ds_fields = xr.open_dataset(data_folder / \"mitgcm_UV_surface_zonally_reentrant.nc\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3",
+   "metadata": {},
+   "source": [
+    "We can use a combination of `parcels.convert.mitgcm_to_sgrid` and `FieldSet.from_sgrid_conventions` to read in the data. See below for an example.\n",
+    "\n",
+    "```{note}\n",
+    "It is very important that you provide the corner nodes as coordinates when converting MITgcm data to S-grid conventions. These corner nodes are typically called `XG` and `YG` in MITgcm output. Failing to do so will lead to incorrect interpolation of the velocity fields.\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "coords = ds_fields[[\"XG\", \"YG\", \"Zl\", \"time\"]]\n",
+    "ds_fset = parcels.convert.mitgcm_to_sgrid(\n",
+    "    fields={\"U\": ds_fields.UVEL, \"V\": ds_fields.VVEL}, coords=coords\n",
+    ")\n",
+    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5",
+   "metadata": {},
+   "source": [
+    "We can then run a simulation on this zonally re-entrant dataset as follows:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {
+    "tags": [
+     "hide-output"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "perc = np.arange(5, 100, 10)\n",
+    "lon_vals = np.percentile(fieldset.UV.grid.lon, perc)\n",
+    "lat_vals = np.percentile(fieldset.UV.grid.lat, perc)\n",
+    "\n",
+    "X, Y = np.meshgrid(\n",
+    "    lon_vals,\n",
+    "    lat_vals,\n",
+    ")\n",
+    "Z = np.zeros_like(X)\n",
+    "\n",
+    "\n",
+    "def DeleteParticle(particles, fieldset):\n",
+    "    any_error = particles.state >= 50\n",
+    "    particles[any_error].state = parcels.StatusCode.Delete\n",
+    "\n",
+    "\n",
+    "pset = parcels.ParticleSet(fieldset=fieldset, lon=X, lat=Y, z=Z)\n",
+    "\n",
+    "outputfile = parcels.ParticleFile(\n",
+    "    store=\"mitgcm_particles.zarr\",\n",
+    "    outputdt=np.timedelta64(5000, \"s\"),\n",
+    ")\n",
+    "\n",
+    "pset.execute(\n",
+    "    [parcels.kernels.AdvectionRK2, DeleteParticle],\n",
+    "    runtime=np.timedelta64(10, \"D\"),\n",
+    "    dt=np.timedelta64(1, \"h\"),\n",
+    "    output_file=outputfile,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7",
+   "metadata": {},
+   "source": [
+    "And make a simple plot of the particle trajectories:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds = xr.open_zarr(\"mitgcm_particles.zarr\")\n",
+    "\n",
+    "plt.plot(ds.lon.T, ds.lat.T, \".-\")\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "docs",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.14.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/user_guide/index.md

@@ -26,6 +26,7 @@ The tutorials written for Parcels v3 are currently being updated for Parcels v4.
 :titlesonly:
 examples/explanation_grids.md
 examples/tutorial_nemo.ipynb
+examples/tutorial_mitgcm.ipynb
 examples/tutorial_velocityconversion.ipynb
 examples/tutorial_nestedgrids.ipynb
 examples/tutorial_manipulating_field_data.ipynb

src/parcels/_core/utils/time.py

@@ -162,8 +162,15 @@ def timedelta_to_float(dt: float | timedelta | np.timedelta64) -> float:
         return dt.total_seconds()
     if isinstance(dt, np.timedelta64):
         return float(dt / np.timedelta64(1, "s"))
-    if hasattr(dt, "dtype") and np.issubdtype(dt.dtype, np.timedelta64):  # in case of array
-        return (dt / np.timedelta64(1, "s")).astype(float)
+    if hasattr(dt, "dtype"):
+        if np.issubdtype(dt.dtype, np.timedelta64):  # in case of array
+            return (dt / np.timedelta64(1, "s")).astype(float)
+        elif np.issubdtype(dt.dtype, np.object_):  # in case of array of timedeltas
+            try:
+                f = np.vectorize(lambda x: x.total_seconds())
+                return f(dt)
+            except Exception as e:
+                raise ValueError(f"Expected a timedelta-like object, got {dt!r}.") from e
     return float(dt)
 
 

src/parcels/convert.py

@@ -42,6 +42,27 @@
     "wo": "W",
 }
 
+_MITGCM_AXIS_VARNAMES = {
+    "XC": "X",
+    "XG": "X",
+    "Xp1": "X",
+    "lon": "X",
+    "YC": "Y",
+    "YG": "Y",
+    "Yp1": "Y",
+    "lat": "Y",
+    "Zu": "Z",
+    "Zl": "Z",
+    "Zp1": "Z",
+    "time": "T",
+}
+
+_MITGCM_VARNAMES_MAPPING = {
+    "XG": "lon",
+    "YG": "lat",
+    "Zl": "depth",
+}
+
 _COPERNICUS_MARINE_AXIS_VARNAMES = {
     "X": "lon",
     "Y": "lat",
@@ -114,7 +135,8 @@ def _maybe_remove_depth_from_lonlat(ds):
 
 def _set_axis_attrs(ds, dim_axis):
     for dim, axis in dim_axis.items():
-        ds[dim].attrs["axis"] = axis
+        if dim in ds:
+            ds[dim].attrs["axis"] = axis
     return ds
 
 
@@ -128,6 +150,16 @@ def _ds_rename_using_standard_names(ds: xr.Dataset | ux.UxDataset, name_dict: di
     return ds
 
 
+def _maybe_swap_depth_direction(ds: xr.Dataset) -> xr.Dataset:
+    if ds["depth"].size > 1:
+        if ds["depth"][0] > ds["depth"][-1]:
+            logger.info(
+                "Depth dimension appears to be decreasing upward (i.e., from positive to negative values). Swapping depth dimension to be increasing upward for Parcels simulation."
+            )
+            ds = ds.reindex(depth=ds["depth"][::-1])
+    return ds
+
+
 # TODO is this function still needed, now that we require users to provide field names explicitly?
 def _discover_U_and_V(ds: xr.Dataset, cf_standard_names_fallbacks) -> xr.Dataset:
     # Assumes that the dataset has U and V data
@@ -266,6 +298,66 @@ def nemo_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Da
     return ds
 
 
+def mitgcm_to_sgrid(*, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Dataset) -> xr.Dataset:
+    """Create an sgrid-compliant xarray.Dataset from a dataset of MITgcm netcdf files.
+
+    Parameters
+    ----------
+    fields : dict[str, xr.Dataset | xr.DataArray]
+        Dictionary of xarray.DataArray objects as obtained from a set of MITgcm netcdf files.
+    coords : xarray.Dataset, optional
+        xarray.Dataset containing coordinate variables.
+
+    Returns
+    -------
+    xarray.Dataset
+        Dataset object following SGRID conventions to be (optionally) modified and passed to a FieldSet constructor.
+
+    Notes
+    -----
+    See the MITgcm tutorial for more information on how to use MITgcm model outputs in Parcels
+
+    """
+    fields = fields.copy()
+
+    for name, field_da in fields.items():
+        if isinstance(field_da, xr.Dataset):
+            field_da = field_da[name]
+            # TODO: logging message, warn if multiple fields are in this dataset
+        else:
+            field_da = field_da.rename(name)
+        fields[name] = field_da
+
+    ds = xr.merge(list(fields.values()) + [coords])
+    ds.attrs.clear()  # Clear global attributes from the merging
+
+    ds = _maybe_rename_variables(ds, _MITGCM_VARNAMES_MAPPING)
+    ds = _set_axis_attrs(ds, _MITGCM_AXIS_VARNAMES)
+    ds = _maybe_swap_depth_direction(ds)
+
+    if "grid" in ds.cf.cf_roles:
+        raise ValueError(
+            "Dataset already has a 'grid' variable (according to cf_roles). Didn't expect there to be grid metadata on copernicusmarine datasets - please open an issue with more information about your dataset."
+        )
+
+    ds["grid"] = xr.DataArray(
+        0,
+        attrs=sgrid.Grid2DMetadata(
+            cf_role="grid_topology",
+            topology_dimension=2,
+            node_dimensions=("lon", "lat"),
+            node_coordinates=("lon", "lat"),
+            face_dimensions=(
+                sgrid.DimDimPadding("XC", "lon", sgrid.Padding.HIGH),
+                sgrid.DimDimPadding("YC", "lat", sgrid.Padding.HIGH),
+            ),
+            vertical_dimensions=(sgrid.DimDimPadding("depth", "depth", sgrid.Padding.HIGH),),
+        ).to_attrs(),
+    )
+
+    return ds
+
+
 def copernicusmarine_to_sgrid(
     *, fields: dict[str, xr.Dataset | xr.DataArray], coords: xr.Dataset | None = None
 ) -> xr.Dataset:

tests/test_advection.py

@@ -3,7 +3,18 @@
 import xarray as xr
 
 import parcels
-from parcels import Field, FieldSet, Particle, ParticleFile, ParticleSet, StatusCode, Variable, VectorField, XGrid
+from parcels import (
+    Field,
+    FieldSet,
+    Particle,
+    ParticleFile,
+    ParticleSet,
+    StatusCode,
+    Variable,
+    VectorField,
+    XGrid,
+    convert,
+)
 from parcels._core.utils.time import timedelta_to_float
 from parcels._datasets.structured.generated import (
     decaying_moving_eddy_dataset,
@@ -491,3 +502,33 @@ def test_nemo_3D_curvilinear_fieldset(kernel):
             [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
+
+
+def test_mitgcm():
+    data_folder = parcels.download_example_dataset("MITgcm_example_data")
+    ds_fields = xr.open_dataset(data_folder / "mitgcm_UV_surface_zonally_reentrant.nc")
+
+    coords = ds_fields[["XG", "YG", "Zl", "time"]]
+    ds_fset = convert.mitgcm_to_sgrid(fields={"U": ds_fields.UVEL, "V": ds_fields.VVEL}, coords=coords)
+    fieldset = FieldSet.from_sgrid_conventions(ds_fset)
+
+    npart = 10
+    lon = [24e3] * npart
+    lat = np.linspace(22e3, 1950e3, npart)
+
+    pset = parcels.ParticleSet(fieldset, lon=lon, lat=lat)
+    pset.execute(AdvectionRK4, runtime=np.timedelta64(5, "D"), dt=np.timedelta64(30, "m"))
+
+    lon_v3 = [
+        25334.3084714,
+        82824.04760837,
+        136410.63322281,
+        98325.83708985,
+        83152.54325753,
+        89321.35275493,
+        237376.5840757,
+        56860.97672692,
+        153947.52685014,
+        28349.16658616,
+    ]
+    np.testing.assert_allclose(pset.lon, lon_v3, atol=10)

tests/test_convert.py

@@ -6,6 +6,7 @@
 from parcels import FieldSet
 from parcels._core.utils import sgrid
 from parcels._datasets.structured.circulation_models import datasets as datasets_circulation_models
+from parcels.interpolators._xinterpolators import _get_offsets_dictionary
 
 
 def test_nemo_to_sgrid():
@@ -39,6 +40,18 @@ def test_nemo_to_sgrid():
     }.issubset(set(ds["V"].dims))
 
 
+def test_convert_mitgcm_offsets():
+    data_folder = parcels.download_example_dataset("MITgcm_example_data")
+    ds_fields = xr.open_dataset(data_folder / "mitgcm_UV_surface_zonally_reentrant.nc")
+    coords = ds_fields[["XG", "YG", "Zl", "time"]]
+    ds_fset = convert.mitgcm_to_sgrid(fields={"U": ds_fields.UVEL, "V": ds_fields.VVEL}, coords=coords)
+    fieldset = FieldSet.from_sgrid_conventions(ds_fset)
+    offsets = _get_offsets_dictionary(fieldset.UV.grid)
+    assert offsets["X"] == 0
+    assert offsets["Y"] == 0
+    assert offsets["Z"] == 0
+
+
 _COPERNICUS_DATASETS = [
     datasets_circulation_models["ds_copernicusmarine"],
     datasets_circulation_models["ds_copernicusmarine_waves"],