diff --git a/ci/scripts/download-all-tutorial-datasets.py b/ci/scripts/download-all-tutorial-datasets.py index 3676012bb..e27c44464 100644 --- a/ci/scripts/download-all-tutorial-datasets.py +++ b/ci/scripts/download-all-tutorial-datasets.py @@ -1,4 +1,4 @@ import parcels.tutorial for name in parcels.tutorial.list_datasets(): - parcels.tutorial.open_dataset(name) + parcels.tutorial.get_dataset_files(name) diff --git a/docs/user_guide/examples/tutorial_fesom.ipynb b/docs/user_guide/examples/tutorial_fesom.ipynb index 964457571..6c40a8079 100644 --- a/docs/user_guide/examples/tutorial_fesom.ipynb +++ b/docs/user_guide/examples/tutorial_fesom.ipynb @@ -44,9 +44,9 @@ "source": [ "## Get the FESOM tutorial dataset\n", "\n", - "We use a small periodic-channel snapshot from a FESOM2 simulation that ships with Parcels' tutorial data registry. As in the [quickstart](../getting_started/tutorial_quickstart.md), `parcels.tutorial.open_dataset` downloads the files into a local cache on first use; subsequent calls just return the cached copy.\n", + "We use a small periodic-channel snapshot from a FESOM2 simulation that ships with Parcels' tutorial data registry. \n", "\n", - "`uxarray` expects file paths rather than an in-memory dataset, so we trigger the downloads and then point `ux.open_mfdataset` at the cached files:" + "`uxarray` expects file paths rather than an in-memory dataset, so we trigger the downloads with `parcels.tutorial.get_dataset_files()` and then point `ux.open_mfdataset` at the cached files:" ] }, { @@ -55,17 +55,8 @@ "metadata": {}, "outputs": [], "source": [ - "for name in [\n", - " \"FESOM_periodic_channel/fesom_channel\", # grid description\n", - " \"FESOM_periodic_channel/u.fesom_channel\", # zonal velocity (face-registered)\n", - " \"FESOM_periodic_channel/v.fesom_channel\", # meridional velocity (face-registered)\n", - " \"FESOM_periodic_channel/w.fesom_channel\", # vertical velocity (node-registered)\n", - "]:\n", - " parcels.tutorial.open_dataset(name)\n", - "\n", - "from parcels._datasets.remote import _DATA_HOME\n", - "\n", - "data_dir = Path(_DATA_HOME) / \"data\" / \"FESOM_periodic_channel\"\n", + "files = parcels.tutorial.get_dataset_files(\"FESOM_periodic_channel/fesom_channel\")\n", + "data_dir = Path(files[0]).parent\n", "\n", "grid_path = str(data_dir / \"fesom_channel.nc\")\n", "data_paths = [\n", @@ -147,7 +138,7 @@ "outputs": [], "source": [ "fieldset = parcels.FieldSet.from_ugrid_conventions(ds, mesh=\"spherical\")\n", - "\n", + "fieldset = fieldset.to_windowed_arrays()\n", "fieldset.describe()" ] }, diff --git a/docs/user_guide/examples/tutorial_swash.ipynb b/docs/user_guide/examples/tutorial_swash.ipynb new file mode 100644 index 000000000..c23e59834 --- /dev/null +++ b/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 +} diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index 06a578d79..703f72cb4 100644 --- a/docs/user_guide/index.md +++ b/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 diff --git a/pixi.toml b/pixi.toml index 71035f952..9bb6c4272 100644 --- a/pixi.toml +++ b/pixi.toml @@ -158,6 +158,7 @@ mypy = "*" lxml = "*" # in CI types-tqdm = "*" pandas-stubs = "*" +scipy-stubs = "*" [feature.typing.tasks] typing = { cmd = "mypy src/parcels --install-types", description = "Run static type checking with mypy." } diff --git a/src/parcels/_datasets/remote.py b/src/parcels/_datasets/remote.py index 5fc93868a..47fa99e67 100644 --- a/src/parcels/_datasets/remote.py +++ b/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: ... @@ -140,11 +149,16 @@ def __init__(self, pup: pooch.Pooch, path_relative_to_pup: str, pre_decode_cf_ca # Function to apply to the dataset before the decoding the CF variables self.pup = pup - self.pre_decode_cf_callable: None | Callable[[xr.Dataset], xr.Dataset] = pre_decode_cf_callable + self.pre_decode_cf_callable: Callable[[xr.Dataset], xr.Dataset] | None = pre_decode_cf_callable 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,12 @@ 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]: + for file in self.pup.registry: + if fnmatch.fnmatch(file, self.path_relative_to_root): + self.pup.fetch(file) + return sorted(glob.glob(f"{self.pup.path}/{self.path_relative_to_root}")) + 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) @@ -205,13 +225,21 @@ class _Purpose(enum.Enum): _TPurpose = Literal["testing", "tutorial"] + +class _Scope(enum.Enum): + DOWNLOAD_ONLY = "download_only" + OPEN = "open" + + +_TScope = Literal["download_only", "open"] + # The first here is a human readable key used to open datasets, with an object to open the datasets # fmt: off -_DATASET_KEYS_AND_CONFIGS: dict[str, tuple[_ParcelsDataset, _Purpose]] = dict([ +_DATASET_KEYS_AND_CONFIGS: dict[str, tuple[_ParcelsDataset, _Purpose, _Scope]] = dict([ # ("MovingEddies_data/P", (_V3Dataset(_ODIE,"data/MovingEddies_data/moving_eddiesP.nc"), _Purpose.TUTORIAL)), # ("MovingEddies_data/U", (_V3Dataset(_ODIE,"data/MovingEddies_data/moving_eddiesU.nc"), _Purpose.TUTORIAL)), # ("MovingEddies_data/V", (_V3Dataset(_ODIE,"data/MovingEddies_data/moving_eddiesV.nc"), _Purpose.TUTORIAL)), - ("MITgcm_example_data/mitgcm_UV_surface_zonally_reentrant", (_V3Dataset(_ODIE,"data/MITgcm_example_data/mitgcm_UV_surface_zonally_reentrant.nc"), _Purpose.TUTORIAL)), + ("MITgcm_example_data/mitgcm_UV_surface_zonally_reentrant", (_V3Dataset(_ODIE,"data/MITgcm_example_data/mitgcm_UV_surface_zonally_reentrant.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), # ("OFAM_example_data/U", (_V3Dataset(_ODIE,"data/OFAM_example_data/OFAM_simple_U.nc"), _Purpose.TUTORIAL)), # ("OFAM_example_data/V", (_V3Dataset(_ODIE,"data/OFAM_example_data/OFAM_simple_V.nc"), _Purpose.TUTORIAL)), # ("Peninsula_data/U", (_V3Dataset(_ODIE,"data/Peninsula_data/peninsulaU.nc"), _Purpose.TUTORIAL)), @@ -219,39 +247,39 @@ class _Purpose(enum.Enum): # ("Peninsula_data/P", (_V3Dataset(_ODIE,"data/Peninsula_data/peninsulaP.nc"), _Purpose.TUTORIAL)), # ("Peninsula_data/T", (_V3Dataset(_ODIE,"data/Peninsula_data/peninsulaT.nc"), _Purpose.TUTORIAL)), # ("GlobCurrent_example_data/data", (_V3Dataset(_ODIE,"data/GlobCurrent_example_data/*000000-GLOBCURRENT-L4-CUReul_hs-ALT_SUM-v02.0-fv01.0.nc", pre_decode_cf_callable=patch_dataset_v4_compat), _Purpose.TUTORIAL)), - ("CopernicusMarine_data_for_Argo_tutorial/data", (_V3Dataset(_ODIE,"data/CopernicusMarine_data_for_Argo_tutorial/cmems_mod_glo_phy-*.nc"), _Purpose.TUTORIAL)), + ("CopernicusMarine_data_for_Argo_tutorial/data", (_V3Dataset(_ODIE,"data/CopernicusMarine_data_for_Argo_tutorial/cmems_mod_glo_phy-*.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), # ("DecayingMovingEddy_data/U", (_V3Dataset(_ODIE,"data/DecayingMovingEddy_data/decaying_moving_eddyU.nc"), _Purpose.TUTORIAL)), # ("DecayingMovingEddy_data/V", (_V3Dataset(_ODIE,"data/DecayingMovingEddy_data/decaying_moving_eddyV.nc"), _Purpose.TUTORIAL)), - ("FESOM_periodic_channel/fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/fesom_channel.nc"), _Purpose.TUTORIAL)), - ("FESOM_periodic_channel/u.fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/u.fesom_channel.nc"), _Purpose.TUTORIAL)), - ("FESOM_periodic_channel/v.fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/v.fesom_channel.nc"), _Purpose.TUTORIAL)), - ("FESOM_periodic_channel/w.fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/w.fesom_channel.nc"), _Purpose.TUTORIAL)), - ("SCHISM_LakeOntario/out2d", (_V3Dataset(_ODIE,"data/SCHISM_LakeOntario/out2d.schism_lake_ontario.nc"), _Purpose.TUTORIAL)), - ("SCHISM_LakeOntario/horizontalVelX", (_V3Dataset(_ODIE,"data/SCHISM_LakeOntario/horizontalVelX.schism_lake_ontario.nc"), _Purpose.TUTORIAL)), - ("SCHISM_LakeOntario/horizontalVelY", (_V3Dataset(_ODIE,"data/SCHISM_LakeOntario/horizontalVelY.schism_lake_ontario.nc"), _Purpose.TUTORIAL)), - ("NemoCurvilinear_data_zonal/U", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/U_purely_zonal-ORCA025_grid_U.nc4"), _Purpose.TUTORIAL)), - ("NemoCurvilinear_data_zonal/V", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/V_purely_zonal-ORCA025_grid_V.nc4"), _Purpose.TUTORIAL)), - ("NemoCurvilinear_data_zonal/mesh_mask", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/mesh_mask.nc4", _preprocess_drop_time_from_mesh2), _Purpose.TUTORIAL)), - ("NemoNorthSeaORCA025-N006_data/U", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05U.nc"), _Purpose.TUTORIAL)), - ("NemoNorthSeaORCA025-N006_data/V", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05V.nc"), _Purpose.TUTORIAL)), - ("NemoNorthSeaORCA025-N006_data/W", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05W.nc"), _Purpose.TUTORIAL)), - ("NemoNorthSeaORCA025-N006_data/mesh_mask", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/coordinates.nc", _preprocess_drop_time_from_mesh1), _Purpose.TUTORIAL)), + ("FESOM_periodic_channel/fesom_channel", (_V3Dataset(_ODIE,"data/FESOM_periodic_channel/*fesom_channel.nc"), _Purpose.TUTORIAL, _Scope.DOWNLOAD_ONLY)), + ("SCHISM_LakeOntario/out2d", (_V3Dataset(_ODIE,"data/SCHISM_LakeOntario/out2d.schism_lake_ontario.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("SCHISM_LakeOntario/horizontalVelX", (_V3Dataset(_ODIE,"data/SCHISM_LakeOntario/horizontalVelX.schism_lake_ontario.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("SCHISM_LakeOntario/horizontalVelY", (_V3Dataset(_ODIE,"data/SCHISM_LakeOntario/horizontalVelY.schism_lake_ontario.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("NemoCurvilinear_data_zonal/U", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/U_purely_zonal-ORCA025_grid_U.nc4"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("NemoCurvilinear_data_zonal/V", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/V_purely_zonal-ORCA025_grid_V.nc4"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("NemoCurvilinear_data_zonal/mesh_mask", (_V3Dataset(_ODIE,"data/NemoCurvilinear_data/mesh_mask.nc4", _preprocess_drop_time_from_mesh2), _Purpose.TUTORIAL, _Scope.OPEN)), + ("NemoNorthSeaORCA025-N006_data/U", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05U.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("NemoNorthSeaORCA025-N006_data/V", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05V.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("NemoNorthSeaORCA025-N006_data/W", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/ORCA025-N06_200001*05W.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("NemoNorthSeaORCA025-N006_data/mesh_mask", (_V3Dataset(_ODIE,"data/NemoNorthSeaORCA025-N006_data/coordinates.nc", _preprocess_drop_time_from_mesh1), _Purpose.TUTORIAL, _Scope.OPEN)), # "POPSouthernOcean_data/t.x1_SAMOC_flux.16900*.nc", # TODO v4: In v3 but should not be in v4 https://github.com/Parcels-code/Parcels/issues/2571#issuecomment-4214476973 # ("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)), + ("CROCOidealized_data/data", (_V3Dataset(_ODIE,"data/CROCOidealized_data/CROCO_idealized.nc"), _Purpose.TUTORIAL, _Scope.OPEN)), + ("SWASH_data/data", (_V3Dataset(_ODIE,"data-matlab/F1*.mat"), _Purpose.TUTORIAL, _Scope.DOWNLOAD_ONLY)), ] + [ - ("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)), - ("Benchmarks_MOi_data_metadata-only/U", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/U.zip"), _Purpose.TESTING)), - ("Benchmarks_MOi_data_metadata-only/V", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/V.zip"), _Purpose.TESTING)), - ("Benchmarks_MOi_data_metadata-only/W", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/W.zip"), _Purpose.TESTING)), - ("Benchmarks_MOi_data_metadata-only/mesh", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/mesh.zip"), _Purpose.TESTING)), + ("Benchmarks_FESOM2-baroclinic-gyre/data", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/data.zip', zarr_format=2), _Purpose.TESTING, _Scope.OPEN)), + ("Benchmarks_FESOM2-baroclinic-gyre/grid", (_ZarrZipDataset(_ODIE, 'data-zarr/Benchmarks_FESOM2-baroclinic-gyre/grid.zip', zarr_format=2), _Purpose.TESTING, _Scope.OPEN)), + ("Benchmarks_MOi_data_metadata-only/U", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/U.zip"), _Purpose.TESTING, _Scope.OPEN)), + ("Benchmarks_MOi_data_metadata-only/V", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/V.zip"), _Purpose.TESTING, _Scope.OPEN)), + ("Benchmarks_MOi_data_metadata-only/W", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/W.zip"), _Purpose.TESTING, _Scope.OPEN)), + ("Benchmarks_MOi_data_metadata-only/mesh", (_ZarrZipDataset(_ODIE, "data-zarr/Benchmarks_MOi_data_metadata-only/mesh.zip"), _Purpose.TESTING, _Scope.OPEN)), ]) # fmt: on -def list_remote_datasets(purpose: _TPurpose | Literal["any"] = "any") -> list[str]: +def list_remote_datasets( + purpose: _TPurpose | Literal["any"] = "any", scope: _TScope | Literal["any"] = "any" +) -> list[str]: """List the available remote datasets. Use :func:`open_dataset` to download and open one of the datasets. @@ -262,17 +290,30 @@ def list_remote_datasets(purpose: _TPurpose | Literal["any"] = "any") -> list[st Filter datasets by purpose. Use ``'any'`` (default) to return all datasets, ``'tutorial'`` for tutorial datasets, or ``'testing'`` for datasets used in tests. + scope : {'any', 'download_only', 'open'}, optional + Filter datasets by scope. Use ``'any'`` (default) to return all datasets, + ``'open'`` to return datasets that can be downloaded and opened, or + ``'download_only'`` for datasets that can only be downloaded. Returns ------- datasets : list of str The names of the available datasets matching the given purpose. """ - if purpose == "any": - return list(_DATASET_KEYS_AND_CONFIGS.keys()) - purpose_enum = _Purpose(purpose) - return [k for (k, (_, p)) in _DATASET_KEYS_AND_CONFIGS.items() if p == purpose_enum] + def _test_purpose(p, purpose): + if purpose == "any": + return True + return p == _Purpose(purpose) + + def _test_scope(s, scope): + if scope == "any": + return True + return s == _Scope(scope) + + return [ + k for (k, (_, p, s)) in _DATASET_KEYS_AND_CONFIGS.items() if _test_purpose(p, purpose) and _test_scope(s, scope) + ] def open_remote_dataset(name: str, purpose: _TPurpose | Literal["any"] = "any"): @@ -298,6 +339,37 @@ def open_remote_dataset(name: str, purpose: _TPurpose | Literal["any"] = "any"): raise ValueError( f"Dataset {name!r} not found. Available datasets are: " + ", ".join(list_remote_datasets(purpose=purpose)) ) - + if name not in list_remote_datasets(scope="open"): + raise ValueError( + f"Dataset {name!r} does not have Scope 'open', so is not intended to be opened. Use get_remote_dataset() to download the files instead." + ) 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() diff --git a/src/parcels/convert.py b/src/parcels/convert.py index b9eeffa98..40f022d4a 100644 --- a/src/parcels/convert.py +++ b/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 @@ -576,6 +578,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"} diff --git a/src/parcels/tutorial.py b/src/parcels/tutorial.py index 49a4a8b6b..7a5466edf 100644 --- a/src/parcels/tutorial.py +++ b/src/parcels/tutorial.py @@ -1,3 +1,4 @@ +from parcels._datasets.remote import get_remote_dataset as _get_remote_dataset from parcels._datasets.remote import list_remote_datasets as _list_remote_datasets from parcels._datasets.remote import open_remote_dataset as _open_remote_dataset @@ -34,3 +35,22 @@ def open_dataset(name: str): The requested dataset. """ return _open_remote_dataset(name, purpose="tutorial") + + +def get_dataset_files(name: str): + """Download the files of a tutorial 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`. + + Returns + ------- + list of str + The list of dataset files. + """ + return _get_remote_dataset(name, purpose="tutorial") diff --git a/tests/datasets/test_remote.py b/tests/datasets/test_remote.py index d4ad168cd..99c21806c 100644 --- a/tests/datasets/test_remote.py +++ b/tests/datasets/test_remote.py @@ -24,12 +24,26 @@ def test_open_dataset_non_existing(): @pytest.mark.flaky -@pytest.mark.parametrize("name", remote.list_remote_datasets()) +@pytest.mark.parametrize("name", remote.list_remote_datasets(scope="open")) def test_open_dataset(name): ds = remote.open_remote_dataset(name) assert isinstance(ds, xr.Dataset) +@pytest.mark.flaky +@pytest.mark.parametrize("name", remote.list_remote_datasets(scope="download_only")) +def test_download_dataset(name): + files = remote.get_remote_dataset(name) + assert all(isinstance(f, str) for f in files) + + +@pytest.mark.flaky +@pytest.mark.parametrize("name", ["SWASH_data/data"]) +def test_error_on_download_scope(name): + with pytest.raises(ValueError, match="does not have Scope 'open'"): + remote.open_remote_dataset(name) + + @pytest.mark.parametrize("name", remote.list_remote_datasets()) def test_dataset_keys(name): assert not name.endswith((".zarr", ".zip", ".nc")), "Dataset name should not have suffix" diff --git a/tests/test_convert.py b/tests/test_convert.py index 2d4fbc3b3..2e8430f2e 100644 --- a/tests/test_convert.py +++ b/tests/test_convert.py @@ -182,6 +182,13 @@ def test_convert_copernicusmarine_nodepth(caplog): assert "No depth dimension found in dataset. Added a singleton depth dimension." in caplog.text +def test_convert_swash(): + data_file, coord_file = parcels.tutorial.get_dataset_files("SWASH_data/data") + + ds_fset = convert.swash_to_sgrid(data_file=data_file, coord_file=coord_file) + FieldSet.from_sgrid_conventions(ds_fset) + + def test_convert_fesom_to_ugrid(): grid_file = open_remote_dataset("Benchmarks_FESOM2-baroclinic-gyre/grid") data_files = open_remote_dataset("Benchmarks_FESOM2-baroclinic-gyre/data") diff --git a/tests/test_uxarray_fieldset.py b/tests/test_uxarray_fieldset.py index a39239bd5..5933057ae 100644 --- a/tests/test_uxarray_fieldset.py +++ b/tests/test_uxarray_fieldset.py @@ -4,7 +4,6 @@ import pytest import uxarray as ux -import parcels._datasets.remote as _parcels_remote import parcels.tutorial from parcels import ( FieldSet, @@ -19,15 +18,13 @@ @pytest.fixture def ds_fesom_channel() -> ux.UxDataset: - # Download FESOM files via the new tutorial API - parcels.tutorial.open_dataset("FESOM_periodic_channel/fesom_channel") - # uxarray requires file paths; access the downloaded files from the pooch cache - _fesom_dir = Path(_parcels_remote._DATA_HOME) / "data" / "FESOM_periodic_channel" - grid_path = str(_fesom_dir / "fesom_channel.nc") + files = parcels.tutorial.get_dataset_files("FESOM_periodic_channel/fesom_channel") + fesom_dir = Path(files[0]).parent + grid_path = str(fesom_dir / "fesom_channel.nc") data_path = [ - str(_fesom_dir / "u.fesom_channel.nc"), - str(_fesom_dir / "v.fesom_channel.nc"), - str(_fesom_dir / "w.fesom_channel.nc"), + str(fesom_dir / "u.fesom_channel.nc"), + str(fesom_dir / "v.fesom_channel.nc"), + str(fesom_dir / "w.fesom_channel.nc"), ] ds = ux.open_mfdataset(grid_path, data_path).rename_vars({"u": "U", "v": "V", "w": "W"}) ds = convert.fesom_to_ugrid(ds)