Merge branch 'v4-dev' into feature/uxarray_xarray_fields

Author: fluidnumericsJoe

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/structured/grid_datasets.py

@@ -0,0 +1,169 @@
+"""Datasets focussing on grid geometry"""
+
+import numpy as np
+import xarray as xr
+
+N = 30
+T = 10
+
+
+def rotated_curvilinear_grid():
+    XG = np.arange(N)
+    YG = np.arange(2 * N)
+    LON, LAT = np.meshgrid(XG, YG)
+
+    angle = -np.pi / 24
+    rotation = np.array([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]])
+
+    # rotate the LON and LAT grids
+    LON, LAT = np.einsum("ji, mni -> jmn", rotation, np.dstack([LON, LAT]))
+
+    return xr.Dataset(
+        {
+            "data_g": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "data_c": (["ZC", "YC", "XC"], np.random.rand(3 * N, 2 * N, N)),
+        },
+        coords={
+            "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
+            "YG": (["YG"], YG, {"axis": "Y", "c_grid_axis_shift": -0.5}),
+            "XC": (["XC"], XG + 0.5, {"axis": "X"}),
+            "YC": (["YC"], YG + 0.5, {"axis": "Y"}),
+            "ZG": (
+                ["ZG"],
+                np.arange(3 * N),
+                {"axis": "Z", "c_grid_axis_shift": -0.5},
+            ),
+            "ZC": (
+                ["ZC"],
+                np.arange(3 * N) + 0.5,
+                {"axis": "Z"},
+            ),
+            "depth": (["ZG"], np.arange(3 * N), {"axis": "Z"}),
+            "time": (["time"], np.arange(T), {"axis": "T"}),
+            "lon": (
+                ["YG", "XG"],
+                LON,
+                {"axis": "X", "c_grid_axis_shift": -0.5},  # ? Needed?
+            ),
+            "lat": (
+                ["YG", "XG"],
+                LAT,
+                {"axis": "Y", "c_grid_axis_shift": -0.5},  # ? Needed?
+            ),
+        },
+    )
+
+
+def _cartesion_to_polar(x, y):
+    r = np.sqrt(x**2 + y**2)
+    theta = np.arctan2(y, x)
+    return r, theta
+
+
+def _polar_to_cartesian(r, theta):
+    x = r * np.cos(theta)
+    y = r * np.sin(theta)
+    return x, y
+
+
+def unrolled_cone_curvilinear_grid():
+    # Not a great unrolled cone, but this is good enough for testing
+    # you can use matplotlib pcolormesh to plot
+    XG = np.arange(N)
+    YG = np.arange(2 * N) * 0.25
+
+    pivot = -10, 0
+    LON, LAT = np.meshgrid(XG, YG)
+
+    new_lon_lat = []
+
+    min_lon = np.min(XG)
+    for lon, lat in zip(LON.flatten(), LAT.flatten(), strict=True):
+        r, _ = _cartesion_to_polar(lon - pivot[0], lat - pivot[1])
+        _, theta = _cartesion_to_polar(min_lon - pivot[0], lat - pivot[1])
+        theta *= 1.2
+        r *= 1.2
+        lon, lat = _polar_to_cartesian(r, theta)
+        new_lon_lat.append((lon + pivot[0], lat + pivot[1]))
+
+    new_lon, new_lat = zip(*new_lon_lat, strict=True)
+    LON, LAT = np.array(new_lon).reshape(LON.shape), np.array(new_lat).reshape(LAT.shape)
+
+    return xr.Dataset(
+        {
+            "data_g": (["ZG", "YG", "XG"], np.random.rand(3 * N, 2 * N, N)),
+            "data_c": (["ZC", "YC", "XC"], np.random.rand(3 * N, 2 * N, N)),
+        },
+        coords={
+            "XG": (["XG"], XG, {"axis": "X", "c_grid_axis_shift": -0.5}),
+            "YG": (["YG"], YG, {"axis": "Y", "c_grid_axis_shift": -0.5}),
+            "XC": (["XC"], XG + 0.5, {"axis": "X"}),
+            "YC": (["YC"], YG + 0.5, {"axis": "Y"}),
+            "ZG": (
+                ["ZG"],
+                np.arange(3 * N),
+                {"axis": "Z", "c_grid_axis_shift": -0.5},
+            ),
+            "ZC": (
+                ["ZC"],
+                np.arange(3 * N) + 0.5,
+                {"axis": "Z"},
+            ),
+            "depth": (["ZG"], np.arange(3 * N), {"axis": "Z"}),
+            "time": (["time"], np.arange(T), {"axis": "T"}),
+            "lon": (
+                ["YG", "XG"],
+                LON,
+                {"axis": "X", "c_grid_axis_shift": -0.5},  # ? Needed?
+            ),
+            "lat": (
+                ["YG", "XG"],
+                LAT,
+                {"axis": "Y", "c_grid_axis_shift": -0.5},  # ? Needed?
+            ),
+        },
+    )
+
+
+datasets = {
+    "2d_left_rotated": rotated_curvilinear_grid(),
+    "ds_2d_left": xr.Dataset(
+        {
+            "data_g": (["time", "ZG", "YG", "XG"], np.random.rand(T, 3 * N, 2 * N, N)),
+            "data_c": (["time", "ZC", "YC", "XC"], np.random.rand(T, 3 * N, 2 * N, N)),
+        },
+        coords={
+            "XG": (
+                ["XG"],
+                2 * np.pi / N * np.arange(0, N),
+                {"axis": "X", "c_grid_axis_shift": -0.5},
+            ),
+            "XC": (["XC"], 2 * np.pi / N * (np.arange(0, N) + 0.5), {"axis": "X"}),
+            "YG": (
+                ["YG"],
+                2 * np.pi / (2 * N) * np.arange(0, 2 * N),
+                {"axis": "Y", "c_grid_axis_shift": -0.5},
+            ),
+            "YC": (
+                ["YC"],
+                2 * np.pi / (2 * N) * (np.arange(0, 2 * N) + 0.5),
+                {"axis": "Y"},
+            ),
+            "ZG": (
+                ["ZG"],
+                np.arange(3 * N),
+                {"axis": "Z", "c_grid_axis_shift": -0.5},
+            ),
+            "ZC": (
+                ["ZC"],
+                np.arange(3 * N) + 0.5,
+                {"axis": "Z"},
+            ),
+            "lon": (["XG"], 2 * np.pi / N * np.arange(0, N)),
+            "lat": (["YG"], 2 * np.pi / (2 * N) * np.arange(0, 2 * N)),
+            "depth": (["ZG"], np.arange(3 * N)),
+            "time": (["time"], np.arange(T), {"axis": "T"}),
+        },
+    ),
+    "2d_left_unrolled_cone": unrolled_cone_curvilinear_grid(),
+}

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)

parcels/grid.py

@@ -10,7 +10,6 @@
     "CurvilinearSGrid",
     "CurvilinearZGrid",
     "Grid",
-    "GridCode",
     "GridType",
     "RectilinearSGrid",
     "RectilinearZGrid",
@@ -24,11 +23,6 @@ class GridType(IntEnum):
     CurvilinearSGrid = 3
 
 
-# GridCode has been renamed to GridType for consistency.
-# TODO: Remove alias in Parcels v4
-GridCode = GridType
-
-
 class Grid:
     """Grid class that defines a (spatial and temporal) grid on which Fields are defined."""
 
@@ -40,7 +34,6 @@ def __init__(
         time_origin: TimeConverter | None,
         mesh: Mesh,
     ):
-        self._ti = -1
         lon = np.array(lon)
         lat = np.array(lat)
         time = np.zeros(1, dtype=np.float64) if time is None else time
@@ -112,7 +105,6 @@ def create_grid(
         time,
         time_origin,
         mesh: Mesh,
-        **kwargs,
     ):
         lon = np.array(lon)
         lat = np.array(lat)
@@ -122,14 +114,14 @@ def create_grid(
 
         if len(lon.shape) <= 1:
             if depth is None or len(depth.shape) <= 1:
-                return RectilinearZGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+                return RectilinearZGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh)
             else:
-                return RectilinearSGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+                return RectilinearSGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh)
         else:
             if depth is None or len(depth.shape) <= 1:
-                return CurvilinearZGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+                return CurvilinearZGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh)
             else:
-                return CurvilinearSGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+                return CurvilinearSGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh)
 
     def _check_zonal_periodic(self):
         if self.zonal_periodic or self.mesh == "flat" or self.lon.size == 1:

parcels/particleset.py

@@ -13,7 +13,7 @@
 from parcels._compat import MPI
 from parcels.application_kernels.advection import AdvectionRK4
 from parcels.field import Field
-from parcels.grid import CurvilinearGrid, GridType
+from parcels.grid import GridType
 from parcels.interaction.interactionkernel import InteractionKernel
 from parcels.interaction.neighborsearch import (
     BruteFlatNeighborSearch,
@@ -430,7 +430,7 @@ def populate_indices(self):
         may be quite expensive.
         """
         for i, grid in enumerate(self.fieldset.gridset.grids):
-            if not isinstance(grid, CurvilinearGrid):
+            if grid._gtype not in [GridType.CurvilinearZGrid, GridType.CurvilinearSGrid]:
                 continue
 
             tree_data = np.stack((grid.lon.flat, grid.lat.flat), axis=-1)