Swash support

docs/user_guide/examples/tutorial_swash.ipynb

@@ -0,0 +1,147 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "source": [
+    "# 🖥️ SWASH tutorial"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {},
+   "source": [
+    "This tutorial shows how to load in native Matlab files from the SWASH model and convert them to a Parcels-compatible `FieldSet`. The tutorial also shows how to run a simple particle tracking simulation using the SWASH data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "import parcels\n",
+    "import parcels.tutorial"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data_file, coord_file = parcels.tutorial.get_dataset_files(\"SWASH_data/data\")\n",
+    "print(data_file, coord_file)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4",
+   "metadata": {},
+   "source": [
+    "Unlike the other tutorials, this tutorial uses the original SWASH output files in Matlab format. If you want to use your own SWASH output files, you should use the `GRD.mat` and `ALL.mat` files in the `parcels.convert.swash_to_sgrid()` function below."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds = parcels.convert.swash_to_sgrid(data_file=data_file, coord_file=coord_file)\n",
+    "fieldset = parcels.FieldSet.from_sgrid_conventions(ds)\n",
+    "fieldset.describe()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6",
+   "metadata": {},
+   "source": [
+    "Now, we can use this `FieldSet` to run a simulation (note it's very short because the dataset provided in this tutorial is only 20 seconds long)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "npart = 10  # number of particles to be released\n",
+    "y = np.linspace(1, 16, npart)\n",
+    "x = np.repeat(15, npart)\n",
+    "\n",
+    "layerI = 0  # at which water depth, the particles are released\n",
+    "z = np.repeat(ds.depth.values[layerI], npart)\n",
+    "pset = parcels.ParticleSet(fieldset, x=x, y=y, z=z)\n",
+    "\n",
+    "output_file = parcels.ParticleFile(\n",
+    "    \"output-swash.parquet\", outputdt=np.timedelta64(5, \"s\"), mode=\"w\"\n",
+    ")\n",
+    "pset.execute(\n",
+    "    parcels.kernels.AdvectionRK2,\n",
+    "    runtime=np.timedelta64(20, \"s\"),\n",
+    "    dt=np.timedelta64(1, \"s\"),\n",
+    "    output_file=output_file,\n",
+    "    verbose_progress=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8",
+   "metadata": {},
+   "source": [
+    "And then plot the results of the simulation, along with the water level at a given time step. The starting positions of the particles are also shown in white."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = parcels.read_particlefile(\"output-swash.parquet\")\n",
+    "\n",
+    "fig, ax = plt.subplots(figsize=(8, 4))\n",
+    "waterlevel = ds.isel(time=2).watlev.plot(x=\"lon\", y=\"lat\", cmap=\"magma\", ax=ax)\n",
+    "for traj in df.partition_by(\"particle_id\"):\n",
+    "    ax.plot(traj[\"x\"][0], traj[\"y\"][0], \"wo\", markersize=5)\n",
+    "    ax.plot(traj[\"x\"], traj[\"y\"], color=\"k\")\n",
+    "ax.set_xlim([14, 16])\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Parcels:docs (3.14.6)",
+   "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.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/user_guide/index.md

@@ -40,6 +40,7 @@ examples/explanation_grids.md
 examples/tutorial_nemo.ipynb
 examples/tutorial_croco_3D.ipynb
 examples/tutorial_mitgcm.ipynb
+examples/tutorial_swash.ipynb
 examples/tutorial_fesom.ipynb
 examples/tutorial_schism.ipynb
 examples/tutorial_velocityconversion.ipynb

src/parcels/_datasets/remote.py

@@ -1,5 +1,7 @@
 import abc
 import enum
+import fnmatch
+import glob
 import os
 from collections.abc import Callable
 from pathlib import Path
@@ -107,6 +109,10 @@ def _get_data_home() -> Path:
         "data/SWASH_data/field_0065552.nc",
         "data/SWASH_data/field_0065557.nc",
     ]
+    + [
+        "data-matlab/F1ALL.mat",
+        "data-matlab/F1GRD.mat",
+    ]
     # + [f"data/WOA_data/woa18_decav_t{m:02d}_04.nc" for m in range(1, 13)]
     + ["data/CROCOidealized_data/CROCO_idealized.nc"]
     # These datasets are from v4 of Parcels where we're opting for Zipped zarr datasets
@@ -130,6 +136,9 @@ def _get_data_home() -> Path:
 
 
 class _ParcelsDataset(abc.ABC):
+    @abc.abstractmethod
+    def get_dataset_files(self) -> list[str]: ...
+
     @abc.abstractmethod
     def open_dataset(self) -> xr.Dataset: ...
 
@@ -145,6 +154,11 @@ def __init__(self, pup: pooch.Pooch, path_relative_to_pup: str, pre_decode_cf_ca
         first, second, *_ = path_relative_to_pup.split("/")
         self.v3_dataset_name = f"{first}/{second}"  # e.g., data/my_dataset
 
+    def get_dataset_files(self) -> list[str]:
+        self.download_relevant_files()
+        matches = sorted(glob.glob(f"{self.pup.path}/{self.path_relative_to_root}"))
+        return matches
+
     def open_dataset(self) -> xr.Dataset:
         self.download_relevant_files()
         with xr.set_options(use_new_combine_kwarg_defaults=True):
@@ -165,7 +179,7 @@ def open_dataset(self) -> xr.Dataset:
 
     def download_relevant_files(self) -> None:
         for file in self.pup.registry:
-            if self.v3_dataset_name in file:
+            if fnmatch.fnmatch(file, self.path_relative_to_root):
                 self.pup.fetch(file)
         return
 
@@ -176,6 +190,11 @@ def __init__(self, pup, path_relative_to_pup, zarr_format: Literal[2, 3] = 3):
         self.path_relative_to_root = path_relative_to_pup
         self.zarr_format = zarr_format
 
+    def get_dataset_files(self) -> list[str]:
+        raise NotImplementedError(
+            "get_dataset_files is not supported for zipped zarr datasets. Use open_dataset instead."
+        )
+
     def open_dataset(self) -> xr.Dataset:
         self.pup.fetch(self.path_relative_to_root)
         return xr.open_zarr(ZipStore(Path(self.pup.path) / self.path_relative_to_root), zarr_format=self.zarr_format)
@@ -240,6 +259,7 @@ class _Purpose(enum.Enum):
     # ("SWASH_data/data", (_V3Dataset(_ODIE,"data/SWASH_data/field_00655*.nc"), _Purpose.TUTORIAL)),
     # ("WOA_data/data", (_V3Dataset(_ODIE,"data/WOA_data/woa18_decav_t*_04.nc", _preprocess_set_cf_calendar_360_day), _Purpose.TUTORIAL)),
     ("CROCOidealized_data/data", (_V3Dataset(_ODIE,"data/CROCOidealized_data/CROCO_idealized.nc"), _Purpose.TUTORIAL)),
+    ("SWASH_data/data", (_V3Dataset(_ODIE,"data-matlab/F1*.mat"), _Purpose.TUTORIAL)),
 ] + [
     ("Benchmarks_FESOM2-baroclinic-gyre/data", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/data.zip', zarr_format=2), _Purpose.TESTING)),
     ("Benchmarks_FESOM2-baroclinic-gyre/grid", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/grid.zip', zarr_format=2),_Purpose.TESTING)),
@@ -301,3 +321,31 @@ def open_remote_dataset(name: str, purpose: _TPurpose | Literal["any"] = "any"):
 
     dataset_config = _DATASET_KEYS_AND_CONFIGS[name][0]
     return dataset_config.open_dataset()
+
+
+def get_remote_dataset(name: str, purpose: _TPurpose | Literal["any"] = "any") -> list[str]:
+    """Download the files of a remote dataset.
+
+    Use :func:`list_datasets` to see the available dataset names.
+
+    Parameters
+    ----------
+    name : str
+        Name of the dataset to open. Must be one of the keys returned by
+        :func:`list_datasets`.
+    purpose : {'any', 'testing', 'tutorial'}, optional
+        Purpose filter used to populate the error message when ``name`` is not
+        found. Defaults to ``'any'``.
+
+    Returns
+    -------
+    list of str
+        The list of dataset files.
+    """
+    if name not in list_remote_datasets(purpose=purpose):
+        raise ValueError(
+            f"Dataset {name!r} not found. Available datasets are: " + ", ".join(list_remote_datasets(purpose=purpose))
+        )
+
+    dataset_config = _DATASET_KEYS_AND_CONFIGS[name][0]
+    return dataset_config.get_dataset_files()

src/parcels/convert.py

@@ -13,11 +13,13 @@
 from __future__ import annotations
 
 import enum
+import re
 import typing
 import warnings
 from typing import cast
 
 import numpy as np
+import scipy.io as sio
 import xarray as xr
 
 import parcels._sgrid as sgrid
@@ -567,6 +569,112 @@ def copernicusmarine_to_sgrid(
     return ds
 
 
+def swash_to_sgrid(data_file: str, coord_file: str) -> xr.Dataset:
+    """Create an sgrid-compliant xarray.Dataset from a dataset of SWASH netcdf files.
+
+    Parameters
+    ----------
+    data_file : str
+        Path to the SWASH data file (MATLAB binary format).
+    coord_file : str
+        Path to the SWASH coordinate file (MATLAB binary format).
+
+    Returns
+    -------
+    xarray.Dataset
+        Dataset object following SGRID conventions to be (optionally) modified and passed to a FieldSet constructor.
+    """
+    warnings.warn(
+        "The swash_to_sgrid function is experimental and may not work for all SWASH datasets. "
+        "Furthermore, we are not entirely confident that the SGrid layout for SWASH is implemented correctly. "
+        "Please report any issues to the Parcels GitHub repository.",
+        UserWarning,
+        stacklevel=2,
+    )
+    coord = sio.loadmat(coord_file)
+    lon = coord["Xp"]
+    lat = coord["Yp"]
+    XC = np.arange(lon.shape[1])
+    YC = np.arange(lat.shape[0])
+    bot = coord["Botlev"]
+
+    mat = sio.loadmat(data_file)
+    keys = [k for k in mat.keys() if not k.startswith("__")]
+
+    time_keys = sorted(
+        set((int(m.group(1)), int(m.group(2))) for k in keys for m in [re.search(r"_(\d{6})_(\d{3})$", k)] if m)
+    )
+    times = np.array([t[0] * 1000 + t[1] for t in time_keys]).astype("timedelta64[ms]")
+
+    nz = len(set([k.split("_")[1] for k in keys if "Vksi" in k]))
+    depth_centers = np.linspace(1.0 / (2 * nz), 1.0 - 1.0 / (2 * nz), nz)
+    depth_interfaces = np.linspace(0, 1, nz + 1)
+
+    nt, ny, nx = len(times), len(YC), len(XC)
+    watlev = np.full((nt, ny, nx), np.nan, dtype=np.float32)
+    vksi = np.full((nt, nz, ny, nx), np.nan, dtype=np.float32)
+    veta = np.full((nt, nz, ny, nx), np.nan, dtype=np.float32)
+    w = np.full((nt, nz + 1, ny, nx), np.nan, dtype=np.float32)
+
+    for ti, (ts_int, ts_dec) in enumerate(time_keys):
+        ts_str = f"{ts_int:06d}_{ts_dec:03d}"
+        for k in keys:
+            if re.match(rf"Watlev_{ts_str}$", k):
+                watlev[ti, :, :] = mat[k]
+            m = re.match(rf"Vksi_k(\d+)_{ts_str}$", k)
+            if m:
+                vksi[ti, int(m.group(1)) - 1, :, :] = mat[k]
+            m = re.match(rf"Veta_k(\d+)_{ts_str}$", k)
+            if m:
+                veta[ti, int(m.group(1)) - 1, :, :] = mat[k]
+            m = re.match(rf"w(\d+)_{ts_str}$", k)
+            if m and int(m.group(1)) < nz:
+                w[ti, int(m.group(1)), :, :] = mat[k]
+
+    # TODO double-check that the C-grid definition for SWASH here is correct
+    ds = xr.Dataset(
+        {
+            "watlev": (["time", "YG", "XG"], watlev),
+            "U": (["time", "depth", "YC", "XG"], vksi),
+            "V": (["time", "depth", "YG", "XC"], veta),
+            "W": (["time", "depth_f", "YC", "XC"], w),
+            "botlev": (["YG", "XG"], bot),
+        },
+        coords={
+            "time": (["time"], times, {"axis": "T", "units": "ms"}),
+            "depth": (["depth"], depth_centers, {"axis": "Z", "units": "normalised", "positive": "down"}),
+            "depth_f": (["depth_f"], depth_interfaces, {"axis": "Z", "units": "normalised", "positive": "down"}),
+            "YG": (["YG"], YC + 0.5, {"axis": "Y", "c_grid_axis_shift": +0.5}),
+            "YC": (["YC"], YC, {"axis": "Y"}),
+            "XG": (["XG"], XC + 0.5, {"axis": "X", "c_grid_axis_shift": +0.5}),
+            "XC": (["XC"], XC, {"axis": "X"}),
+            "lat": (["YG", "XG"], lat, {"axis": "Y", "units": "m", "c_grid_axis_shift": +0.5}),
+            "lon": (["YG", "XG"], lon, {"axis": "X", "units": "m", "c_grid_axis_shift": +0.5}),
+        },
+    )
+    header = mat["__header__"]
+    if isinstance(header, bytes):
+        header = header.decode("utf-8")
+    ds.attrs.update(header=header, version=mat["__version__"], globals=mat["__globals__"])
+
+    ds["grid"] = xr.DataArray(
+        0,
+        attrs=sgrid.SGrid2DMetadata(
+            cf_role="grid_topology",
+            topology_dimension=2,
+            node_dimensions=("XC", "YC"),
+            node_coordinates=("lon", "lat"),
+            face_dimensions=(
+                sgrid.FaceNodePadding("XC", "XG", sgrid.Padding.HIGH),
+                sgrid.FaceNodePadding("YC", "YG", sgrid.Padding.HIGH),
+            ),
+            vertical_dimensions=(sgrid.FaceNodePadding("depth", "depth_f", sgrid.Padding.LOW),),
+        ).to_attrs(),
+    )
+
+    return ds
+
+
 # Known vertical dimension mappings by model
 _FESOM2_VERTICAL_DIMS = {"interface": "nz", "center": "nz1"}
 _ICON_VERTICAL_DIMS = {"interface": "depth_2", "center": "depth"}