Add parcels._datasets subpackage and developer documentation

parcels/_datasets/__init__.py

@@ -0,0 +1,46 @@
+"""
+Datasets compatible with Parcels.
+
+This subpackage uses xarray to generate *idealised* structured and unstructured hydrodynamical datasets that are compatible with Parcels. The goals are three-fold:
+
+1. To provide users with documentation for the types of datasets they can expect Parcels to work with.
+2. To supply our tutorials with hydrodynamical datasets.
+3. To offer developers datasets for use in test cases.
+
+Note that this subpackage is part of the private API for Parcels. Users should not rely directly on the functions defined within this module. Instead, if you want to generate your own datasets, copy the functions from this module into your own code.
+
+Developers, note that you should only add functions that create idealised datasets to this subpackage if they are (a) quick to generate, and (b) only use dependencies already shipped with Parcels. No data files should be added to this subpackage. Real world data files should be added to the `OceanParcels/parcels-data` repository on GitHub.
+
+Parcels Dataset Philosophy
+-------------------------
+
+When adding datasets, there may be a tension between wanting to add a specific dataset or wanting to add machinery to generate completely parameterised datasets (e.g., with different grid resolutions, with different ranges, with different datetimes etc.). There are trade-offs to both approaches:
+
+Working with specific hardcoded datasets:
+
+* Pros
+    * the example is stable and self-contained
+    * easy to see exactly what the dataset is, there is little to no dependency on other functions defined in the same module
+      * datasets don't "break" due to changes in other functions (e.g., grid edges becoming out of sync with grid centres)
+* Cons
+    * inflexible for use in tests where you want to test a large range of datasets, or you want to test a specific resolution
+
+Working with generated datasets is the opposite of all the above.
+
+Most of the time we only want a single dataset. For example, for use in a tutorial, or for testing a specific feature of Parcels - such as (in the case of structured grids) checking that the grid from a certain (ocean) circulation model is correctly parsed, or checking that indexing is correctly picked up. As such, one should often opt for hardcoded datasets. These are more stable and easier to see exactly what the dataset is. We may have specific examples that become the default "go to" dataset for testing when we don't care about the details of the dataset.
+
+Sometimes we may want to test Parcels against a whole range of datasets varying in a certain way - to ensure Parcels works as expected. For these, we should add machinery to create generated datasets.
+
+Structure
+--------
+
+This subpackage is broken down into structured and unstructured parts. Each of these have common submodules:
+
+* ``circulation_model`` -> hardcoded datasets with the intention of mimicking datasets from a certain (ocean) circulation model
+* ``generic`` -> hardcoded datasets that are generic, and not tied to a certain (ocean) circulation model. Instead these focus on the fundamental properties of the dataset
+* ``generated`` -> functions to generate datasets with varying properties
+* ``utils`` -> any utility functions necessary related to either generating or validating datasets
+
+There may be extra submodules than the ones listed above.
+
+"""

parcels/_datasets/structured/__init__.py

@@ -0,0 +1 @@
+"""Structured datasets."""

parcels/_datasets/unstructured/__init__.py

@@ -0,0 +1 @@
+"""Unstructured datasets."""

parcels/_datasets/utils.py

@@ -0,0 +1,69 @@
+from typing import Any
+
+import numpy as np
+import xarray as xr
+
+from parcels._compat import add_note
+
+_SUPPORTED_ATTR_TYPES = int | float | str | np.ndarray
+
+
+def _print_mismatched_keys(d1: dict[Any, Any], d2: dict[Any, Any]) -> None:
+    k1 = set(d1.keys())
+    k2 = set(d2.keys())
+    if len(k1 ^ k2) == 0:
+        return
+    print("Mismatched keys:")
+    print(f"L: {k1 - k2!r}")
+    print(f"R: {k2 - k1!r}")
+
+
+def assert_common_attrs_equal(
+    xr_attrs_1: dict[str, _SUPPORTED_ATTR_TYPES], xr_attrs_2: dict[str, _SUPPORTED_ATTR_TYPES], *, verbose: bool = True
+) -> None:
+    d1, d2 = xr_attrs_1, xr_attrs_2
+
+    common_keys = set(d1.keys()) & set(d2.keys())
+    if verbose:
+        _print_mismatched_keys(d1, d2)
+
+    for key in common_keys:
+        try:
+            if isinstance(d1[key], np.ndarray):
+                np.testing.assert_array_equal(d1[key], d2[key])
+            else:
+                assert d1[key] == d2[key], f"{d1[key]} != {d2[key]}"
+        except AssertionError as e:
+            add_note(e, f"error on key {key!r}")
+            raise
+
+
+def assert_common_variables_common_attrs_equal(ds1: xr.Dataset, ds2: xr.Dataset, *, verbose: bool = True) -> None:
+    if verbose:
+        print("Checking dataset attrs...")
+
+    assert_common_attrs_equal(ds1.attrs, ds2.attrs, verbose=verbose)
+
+    ds1_vars = set(ds1.variables)
+    ds2_vars = set(ds2.variables)
+
+    common_variables = ds1_vars & ds2_vars
+    if len(ds1_vars ^ ds2_vars) > 0 and verbose:
+        print("Mismatched variables:")
+        print(f"L: {ds1_vars - ds2_vars}")
+        print(f"R: {ds2_vars - ds1_vars}")
+
+    for var in common_variables:
+        if verbose:
+            print(f"Checking {var!r} attrs")
+        assert_common_attrs_equal(ds1[var].attrs, ds2[var].attrs, verbose=verbose)
+
+
+def dataset_repr_diff(ds1: xr.Dataset, ds2: xr.Dataset) -> str:
+    """Return a text diff of two datasets."""
+    repr1 = repr(ds1)
+    repr2 = repr(ds2)
+    import difflib
+
+    diff = difflib.ndiff(repr1.splitlines(keepends=True), repr2.splitlines(keepends=True))
+    return "".join(diff)