diff --git a/parcels/basegrid.py b/parcels/basegrid.py index e1a866e225..2765dd59ec 100644 --- a/parcels/basegrid.py +++ b/parcels/basegrid.py @@ -4,6 +4,8 @@ from enum import IntEnum from typing import TYPE_CHECKING +import numpy as np + if TYPE_CHECKING: import numpy as np @@ -67,7 +69,6 @@ def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int, """ ... - @abstractmethod def ravel_index(self, axis_indices: dict[str, int]) -> int: """ Convert a dictionary of axis indices to a single encoded index (ei). @@ -101,9 +102,10 @@ def ravel_index(self, axis_indices: dict[str, int]) -> int: NotImplementedError Raised if the method is not implemented for the current grid type. """ - ... + dims = np.array([self.get_axis_dim(axis) for axis in self.axes], dtype=int) + indices = np.array([axis_indices[axis] for axis in self.axes], dtype=int) + return _ravel(dims, indices) - @abstractmethod def unravel_index(self, ei: int) -> dict[str, int]: """ Convert a single encoded index (ei) back to a dictionary of axis indices. @@ -134,4 +136,106 @@ def unravel_index(self, ei: int) -> dict[str, int]: NotImplementedError Raised if the method is not implemented for the current grid type. """ + dims = np.array([self.get_axis_dim(axis) for axis in self.axes], dtype=int) + indices = _unravel(dims, ei) + return dict(zip(self.axes, indices, strict=True)) + + @property + @abstractmethod + def axes(self) -> list[str]: + """ + Return a list of axis names that are part of this grid. + + This list must at least be of length 1, and `get_axis_dim` should + return a valid integer for each axis name in the list. + + Returns + ------- + list[str] + List of axis names, e.g. ["Z", "Y", "X"] for a 3D structured grid or ["Z", "FACE"] for an unstructured grid. + """ + ... + + @abstractmethod + def get_axis_dim(self, axis: str) -> int: + """ + Return the dimensionality (number of cells/faces) along a specific axis. + + Parameters + ---------- + axis : str + The name of the axis to get the dimensionality for. Must be one of the values returned by self.axes. + + Returns + ------- + int + The number of cells/edges along the specified axis. + + Raises + ------ + ValueError + If the specified axis is not part of this grid. + """ ... + + +def _unravel(dims, ei): + """ + Converts a flattened (raveled) index back to multi-dimensional indices. + + Args: + dims (1d-array-like): The dimensions along each axis + ei (int): The flattened index to convert + + Returns + ------- + array-like: Indices along each axis corresponding to the given flattened index + + Example: + >>> dims = [2, 3, 4] + >>> ei = 9 + >>> unravel(dims, ei) + array([0, 2, 1]) + # Calculation: + # i0 = 9 // (3*4) = 9 // 12 = 0 + # remainder = 9 % 12 = 9 + # i1 = 9 // 4 = 2 + # i2 = 9 % 4 = 1 + """ + strides = np.cumprod(dims[::-1])[::-1] + + indices = np.empty(len(dims), dtype=int) + + for i in range(len(dims) - 1): + indices[i] = ei // strides[i + 1] + ei = ei % strides[i + 1] + + indices[-1] = ei + return indices + + +def _ravel(dims, indices): + """ + Converts indices to a flattened (raveled) index. + + Args: + dims (1d-array-like): The dimensions along each axis + indices (array-like): Indices along each axis to convert + + Returns + ------- + int: The flattened index corresponding to the given indices + + Example: + >>> dims = [2, 3, 4] + >>> indices = [0, 2, 1] + >>> ravel(dims, indices) + 9 + # Calculation: 0 * (3 * 4) + 2 * (4) + 1 = 0 + 8 + 1 = 9 + """ + strides = np.cumprod(dims[::-1])[::-1] + ei = 0 + for i in range(len(dims) - 1): + ei += indices[i] * strides[i + 1] + + return ei + indices[-1] diff --git a/parcels/field.py b/parcels/field.py index 3268f11c4e..a93eb2cc00 100644 --- a/parcels/field.py +++ b/parcels/field.py @@ -10,7 +10,6 @@ import xarray as xr from parcels._core.utils.time import TimeInterval -from parcels._core.utils.unstructured import get_vertical_location_from_dims from parcels._reprs import default_repr from parcels._typing import ( Mesh, @@ -158,7 +157,7 @@ def __init__( self.grid = grid try: - self.time_interval = get_time_interval(data) + self.time_interval = _get_time_interval(data) except ValueError as e: e.add_note( f"Error getting time interval for field {name!r}. Are you sure that the time dimension on the xarray dataset is stored as timedelta, datetime or cftime datetime objects?" @@ -206,41 +205,6 @@ def units(self, value): raise ValueError(f"Units must be a UnitConverter object, got {type(value)}") self._units = value - @property - def lat(self): - if type(self.data) is ux.UxDataArray: - if self.data.attrs["location"] == "node": - return self.grid.node_lat - elif self.data.attrs["location"] == "face": - return self.grid.face_lat - elif self.data.attrs["location"] == "edge": - return self.grid.edge_lat - else: - return self.grid.lat - - @property - def lon(self): - if type(self.data) is ux.UxDataArray: - if self.data.attrs["location"] == "node": - return self.grid.node_lon - elif self.data.attrs["location"] == "face": - return self.grid.face_lon - elif self.data.attrs["location"] == "edge": - return self.grid.edge_lon - else: - return self.grid.lon - - @property - def depth(self): - if type(self.data) is ux.UxDataArray: - vertical_location = get_vertical_location_from_dims(self.data.dims) - if vertical_location == "center": - return self.grid.nz1 - elif vertical_location == "face": - return self.grid.nz - else: - return self.grid.depth - @property def xdim(self): if type(self.data) is xr.DataArray: @@ -316,13 +280,6 @@ def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True): else: return value - def _rescale_and_set_minmax(self, data): - data[np.isnan(data)] = 0 - return data - - def __getattr__(self, key: str): - return getattr(self.data, key) - def __getitem__(self, key): self._check_velocitysampling() try: @@ -333,9 +290,6 @@ def __getitem__(self, key): except tuple(AllParcelsErrorCodes.keys()) as error: return _deal_with_errors(error, key, vector_type=None) - def __contains__(self, key: str): - return key in self.data - class VectorField: """VectorField class that holds vector field data needed to execute particles.""" @@ -350,9 +304,9 @@ def __init__( self.grid = U.grid if W is None: - assert_same_time_interval((U, V)) + _assert_same_time_interval((U, V)) else: - assert_same_time_interval((U, V, W)) + _assert_same_time_interval((U, V, W)) self.time_interval = U.time_interval @@ -384,14 +338,6 @@ def vector_interp_method(self, method: Callable): _assert_same_function_signature(method, ref=ZeroInterpolator) self._vector_interp_method = method - # @staticmethod - # TODO : def _check_grid_dimensions(grid1, grid2): - # return ( - # np.allclose(grid1.lon, grid2.lon) - # and np.allclose(grid1.lat, grid2.lat) - # and np.allclose(grid1.depth, grid2.depth) - # and np.allclose(grid1.time, grid2.time) - # ) def eval(self, time: datetime, z, y, x, particle=None, applyConversion=True): """Interpolate field values in space and time. @@ -498,14 +444,14 @@ def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux ) -def get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | None: +def _get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | None: if len(data.time) == 1: return None return TimeInterval(data.time.values[0], data.time.values[-1]) -def assert_same_time_interval(fields: list[Field]) -> None: +def _assert_same_time_interval(fields: list[Field]) -> None: if len(fields) == 0: return diff --git a/parcels/fieldfilebuffer.py b/parcels/fieldfilebuffer.py deleted file mode 100644 index c1a0e7dd12..0000000000 --- a/parcels/fieldfilebuffer.py +++ /dev/null @@ -1,184 +0,0 @@ -import datetime -import warnings -from pathlib import Path - -import numpy as np -import xarray as xr - -from parcels._typing import InterpMethodOption, PathLike -from parcels.tools.warnings import FileWarning - - -class NetcdfFileBuffer: - def __init__( - self, - filename, - dimensions, - indices, - *, - interp_method: InterpMethodOption = "linear", - gridindexingtype="nemo", - ): - self.filename: PathLike | list[PathLike] = filename - self.dimensions = dimensions # Dict with dimension keys for file data - self.indices = indices - self.dataset = None - self.interp_method = interp_method - self.gridindexingtype = gridindexingtype - - def __enter__(self): - self.dataset = open_xarray_dataset(self.filename) - return self - - def __exit__(self, type, value, traceback): - self.close() - - def close(self): - if self.dataset is not None: - self.dataset.close() - self.dataset = None - - @property - def xdim(self): - lon = self.dataset[self.dimensions["lon"]] - xdim = lon.size if len(lon.shape) == 1 else lon.shape[-1] - if self.gridindexingtype in ["croco"]: - xdim -= 1 - return xdim - - @property - def ydim(self): - lat = self.dataset[self.dimensions["lat"]] - ydim = lat.size if len(lat.shape) == 1 else lat.shape[-2] - if self.gridindexingtype in ["croco"]: - ydim -= 1 - return ydim - - @property - def zdim(self): - if "depth" not in self.dimensions: - return 1 - depth = self.dataset[self.dimensions["depth"]] - zdim = depth.size if len(depth.shape) == 1 else depth.shape[-3] - if self.gridindexingtype in ["croco"]: - zdim -= 1 - return zdim - - @property - def latlon(self): - lon = self.dataset[self.dimensions["lon"]] - lat = self.dataset[self.dimensions["lat"]] - if self.gridindexingtype not in ["croco"]: - if len(lon.shape) == 3: # some lon, lat have a time dimension 1 - lon = lon[0, :, :] - lat = lat[0, :, :] - elif len(lon.shape) == 4: # some lon, lat have a time and depth dimension 1 - lon = lon[0, 0, :, :] - lat = lat[0, 0, :, :] - - if len(lon.shape) > 1: - if is_rectilinear(lon, lat): - lon = lon[0, :] - lat = lat[:, 0] - return lat, lon - - @property - def depth(self): - self.indices["depth"] = range(self.zdim) - if "depth" in self.dimensions: - return self.dataset[self.dimensions["depth"]] - return np.zeros(1) - - @property - def data_full_zdim(self): - return self.zdim - - def _check_extend_depth(self, data, dim): - return ( - self.indices["depth"][-1] == self.data_full_zdim - 1 - and data.shape[dim] == self.data_full_zdim - 1 - and self.interp_method in ["bgrid_velocity", "bgrid_w_velocity", "bgrid_tracer"] - ) - - def _apply_indices(self, data, ti): - if len(data.shape) == 1: - if self.indices["depth"] is not None: - data = data[self.indices["depth"]] - elif len(data.shape) == 3: - if self._check_extend_depth(data, 0): - data = data[self.indices["depth"][:-1], :, :] - elif len(self.indices["depth"]) > 1: - data = data[self.indices["depth"], :, :] - else: - data = data[ti, :, :] - else: - if self._check_extend_depth(data, 1): - data = data[ti, self.indices["depth"][:-1], :, :] - else: - data = data[ti, self.indices["depth"], :, :] - return data - - @property - def data(self): - return self.data_access() - - def data_access(self): - data = self.dataset[self.name] - ti = range(data.shape[0]) - return np.array(self._apply_indices(data, ti)) - - @property - def time(self): - return self.time_access() - - def time_access(self): - if "time" not in self.dimensions: - return np.array([None]) - - time_da = self.dataset[self.dimensions["time"]] - time = ( - np.array([time_da[self.dimensions["time"]].data]) - if len(time_da.shape) == 0 - else np.array(time_da[self.dimensions["time"]]) - ) - if isinstance(time[0], datetime.datetime): - raise NotImplementedError( - "Parcels currently only parses dates ranging from 1678 AD to 2262 AD, which are stored by xarray as np.datetime64. If you need a wider date range, please open an Issue on the parcels github page." - ) - return time - - -def open_xarray_dataset(filename: Path | str) -> xr.Dataset: - try: - # Unfortunately we need to do if-else here, cause the lock-parameter is either False or a Lock-object - # (which we would rather want to have being auto-managed). - # If 'lock' is not specified, the Lock-object is auto-created and managed by xarray internally. - ds = xr.open_mfdataset(filename, decode_cf=True) - ds["decoded"] = True - except: - warnings.warn( # TODO: Is this warning necessary? What cases does this except block get triggered - is it to do with the bare except??? - f"File {filename} could not be decoded properly by xarray (version {xr.__version__}). " - "It will be opened with no decoding. Filling values might be wrongly parsed.", - FileWarning, - stacklevel=2, - ) - - ds = xr.open_mfdataset(filename, decode_cf=False) - ds["decoded"] = False - return ds - - -def is_rectilinear(lon_subset, lat_subset) -> bool: - """Test if all columns and rows are the same for lon and lat (in which case grid is rectilinear). - - lon_subset and lat_subset are 2D numpy arrays - """ - for xi in range(1, lon_subset.shape[0]): - if not np.allclose(lon_subset[0, :], lon_subset[xi, :]): - return False - - for yi in range(1, lat_subset.shape[1]): - if not np.allclose(lat_subset[:, 0], lat_subset[:, yi]): - return False - - return True diff --git a/parcels/fieldset.py b/parcels/fieldset.py index fbb3055668..c423000fc4 100644 --- a/parcels/fieldset.py +++ b/parcels/fieldset.py @@ -82,34 +82,6 @@ def time_interval(self): return None return functools.reduce(lambda x, y: x.intersection(y), time_intervals) - def dimrange(self, dim): - """Returns maximum value of a dimension (lon, lat, depth or time) - on 'left' side and minimum value on 'right' side for all grids - in a gridset. Useful for finding e.g. longitude range that - overlaps on all grids in a gridset. - """ - maxleft, minright = (-np.inf, np.inf) - dim2ds = { - "depth": ["nz1", "nz"], - "lat": ["node_lat", "face_lat", "edge_lat"], - "lon": ["node_lon", "face_lon", "edge_lon"], - "time": ["time"], - } - for ds in self.datasets: - for field in ds.data_vars: - for d in dim2ds[dim]: # check all possible dimensions - if d in ds[field].dims: - if dim == "depth": - maxleft = max(maxleft, ds[field][d].min().data) - minright = min(minright, ds[field][d].max().data) - else: - maxleft = max(maxleft, ds[field][d].data[0]) - minright = min(minright, ds[field][d].data[-1]) - maxleft = 0 if maxleft == -np.inf else maxleft # if all len(dim) == 1 - minright = 0 if minright == np.inf else minright # if all len(dim) == 1 - - return maxleft, minright - def add_field(self, field: Field, name: str | None = None): """Add a :class:`parcels.field.Field` object to the FieldSet. @@ -204,33 +176,6 @@ def gridset(self) -> list[BaseGrid]: grids.append(field.grid) return grids - # def computeTimeChunk(self, time=0.0, dt=1): - # """Load a chunk of three data time steps into the FieldSet. - # This is used when FieldSet uses data imported from netcdf, - # with default option deferred_load. The loaded time steps are at or immediatly before time - # and the two time steps immediately following time if dt is positive (and inversely for negative dt) - - # Parameters - # ---------- - # time : - # Time around which the FieldSet data are to be loaded. - # Time is provided as a double, relatively to Fieldset.time_origin. - # Default is 0. - # dt : - # time step of the integration scheme, needed to set the direction of time chunk loading. - # Default is 1. - # """ - # nextTime = np.inf if dt > 0 else -np.inf - - # if abs(nextTime) == np.inf or np.isnan(nextTime): # Second happens when dt=0 - # return nextTime - # else: - # nSteps = int((nextTime - time) / dt) - # if nSteps == 0: - # return nextTime - # else: - # return time + nSteps * dt - class CalendarError(Exception): # TODO: Move to a parcels errors module """Exception raised when the calendar of a field is not compatible with the rest of the Fields. The user should ensure that they only add fields to a FieldSet that have compatible CFtime calendars.""" diff --git a/parcels/uxgrid.py b/parcels/uxgrid.py index 641213b7a0..e95eb84b53 100644 --- a/parcels/uxgrid.py +++ b/parcels/uxgrid.py @@ -54,6 +54,19 @@ def depth(self): return np.zeros(1) return self.z.values + @property + def axes(self) -> list[_UXGRID_AXES]: + return ["Z", "FACE"] + + def get_axis_dim(self, axis: _UXGRID_AXES) -> int: + if axis not in self.axes: + raise ValueError(f"Axis {axis!r} is not part of this grid. Available axes: {self.axes}") + + if axis == "Z": + return len(self.z.values) + elif axis == "FACE": + return self.uxgrid.n_face + def search(self, z, y, x, ei=None): tol = 1e-10 @@ -101,11 +114,3 @@ def _get_barycentric_coordinates(self, y, x, fi): bcoord = np.asarray(_barycentric_coordinates(nodes, coord)) err = abs(np.dot(bcoord, nodes[:, 0]) - coord[0]) + abs(np.dot(bcoord, nodes[:, 1]) - coord[1]) return bcoord, err - - def ravel_index(self, axis_indices: dict[_UXGRID_AXES, int]): - return axis_indices["FACE"] + self.uxgrid.n_face * axis_indices["Z"] - - def unravel_index(self, ei) -> dict[_UXGRID_AXES, int]: - zi = ei // self.uxgrid.n_face - fi = ei % self.uxgrid.n_face - return {"Z": zi, "FACE": fi} diff --git a/parcels/xgrid.py b/parcels/xgrid.py index 21d3101851..c3c110485c 100644 --- a/parcels/xgrid.py +++ b/parcels/xgrid.py @@ -1,4 +1,5 @@ from collections.abc import Hashable, Mapping +from functools import cached_property from typing import Literal, cast import numpy as np @@ -17,13 +18,11 @@ _XGCM_AXES = Mapping[_XGCM_AXIS_DIRECTION, xgcm.Axis] -def get_cell_edge_count_along_dim(axis: xgcm.Axis | None) -> int: - if axis is None: - return 1 +def get_cell_count_along_dim(axis: xgcm.Axis) -> int: first_coord = list(axis.coords.items())[0] _, coord_var = first_coord - return axis._ds[coord_var].size + return axis._ds[coord_var].size - 1 def get_time(axis: xgcm.Axis) -> npt.NDArray: @@ -112,21 +111,23 @@ def _datetimes(self): def time(self): return self._datetimes.astype(np.float64) / 1e9 - @property - def xdim(self): - return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("X")) + @cached_property + def xdim(self) -> int: + return self.get_axis_dim("X") - @property - def ydim(self): - return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("Y")) + @cached_property + def ydim(self) -> int: + return self.get_axis_dim("Y") - @property - def zdim(self): - return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("Z")) + @cached_property + def zdim(self) -> int: + return self.get_axis_dim("Z") - @property - def tdim(self): - return get_cell_edge_count_along_dim(self.xgcm_grid.axes.get("T")) + def get_axis_dim(self, axis: _XGRID_AXES) -> int: + if axis not in self.axes: + raise ValueError(f"Axis {axis!r} is not part of this grid. Available axes: {self.axes}") + + return get_cell_count_along_dim(self.xgcm_grid.axes[axis]) @property def _z4d(self) -> Literal[0, 1]: @@ -184,20 +185,6 @@ def search(self, z, y, x, ei=None): raise NotImplementedError("Searching in >2D lon/lat arrays is not implemented yet.") - def ravel_index(self, axis_indices: dict[_XGRID_AXES, int]) -> int: - xi = axis_indices.get("X", 0) - yi = axis_indices.get("Y", 0) - zi = axis_indices.get("Z", 0) - return xi + self.xdim * yi + self.xdim * self.ydim * zi - - def unravel_index(self, ei) -> dict[_XGRID_AXES, int]: - zi = ei // (self.xdim * self.ydim) - ei = ei % (self.xdim * self.ydim) - - yi = ei // self.xdim - xi = ei % self.xdim - return {"Z": zi, "Y": yi, "X": xi} - def get_axis_from_dim_name(axes: _XGCM_AXES, dim: str) -> _XGCM_AXIS_DIRECTION | None: """For a given dimension name in a grid, returns the direction axis it is on.""" diff --git a/pyproject.toml b/pyproject.toml index daee2826ce..8fde067a20 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dependencies = [ "pytest", "scipy", "xarray", + "uxarray", "pooch", ] @@ -68,6 +69,7 @@ pymbolic = "*" scipy = ">=0.16.0" tqdm = "*" xarray = ">=0.10.8" +uxarray = ">=2025.3.0" cftime = ">=1.3.1" dask = ">=2.0" scikit-learn = "*" diff --git a/tests/v4/test_basegrid.py b/tests/v4/test_basegrid.py new file mode 100644 index 0000000000..86a4f12ce2 --- /dev/null +++ b/tests/v4/test_basegrid.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import itertools + +import numpy as np +import pytest + +from parcels.basegrid import BaseGrid + + +class TestGrid(BaseGrid): + def __init__(self, axis_dim: dict[str, int]): + self.axis_dim = axis_dim + + def search(self, z: float, y: float, x: float, ei=None) -> dict[str, tuple[int, float | np.ndarray]]: + pass + + @property + def axes(self) -> list[str]: + return list(self.axis_dim.keys()) + + def get_axis_dim(self, axis: str) -> int: + return self.axis_dim[axis] + + +@pytest.mark.parametrize( + "grid", + [ + TestGrid({"Z": 10, "Y": 20, "X": 30}), + TestGrid({"Z": 5, "Y": 15}), + TestGrid({"Z": 8}), + TestGrid({"Z": 12, "FACE": 25}), + ], +) +def test_basegrid_ravel_unravel_index(grid): + axes = grid.axes + dimensionalities = (grid.get_axis_dim(axis) for axis in axes) + all_possible_axis_indices = itertools.product(*[range(dim) for dim in dimensionalities]) + + encountered_eis = [] + + for axis_indices_numeric in all_possible_axis_indices: + axis_indices = dict(zip(axes, axis_indices_numeric, strict=True)) + + ei = grid.ravel_index(axis_indices) + axis_indices_test = grid.unravel_index(ei) + assert axis_indices_test == axis_indices + encountered_eis.append(ei) + + encountered_eis = sorted(encountered_eis) + assert len(set(encountered_eis)) == len(encountered_eis), "Raveled indices are not unique." + assert np.allclose(np.diff(np.array(encountered_eis)), 1), "Raveled indices are not consecutive integers." + assert encountered_eis[0] == 0, "Raveled indices do not start at 0." diff --git a/tests/v4/test_uxarray_fieldset.py b/tests/v4/test_uxarray_fieldset.py index f1ed7262fe..3399740116 100644 --- a/tests/v4/test_uxarray_fieldset.py +++ b/tests/v4/test_uxarray_fieldset.py @@ -77,16 +77,16 @@ def uvw_fesom_channel(ds_fesom_channel) -> VectorField: def test_fesom_fieldset(ds_fesom_channel, uv_fesom_channel): fieldset = FieldSet([uv_fesom_channel, uv_fesom_channel.U, uv_fesom_channel.V]) # Check that the fieldset has the expected properties - assert (fieldset.U == ds_fesom_channel.U).all() - assert (fieldset.V == ds_fesom_channel.V).all() + assert (fieldset.U.data == ds_fesom_channel.U).all() + assert (fieldset.V.data == ds_fesom_channel.V).all() def test_fesom_in_particleset(ds_fesom_channel, uv_fesom_channel): fieldset = FieldSet([uv_fesom_channel, uv_fesom_channel.U, uv_fesom_channel.V]) # Check that the fieldset has the expected properties - assert (fieldset.U == ds_fesom_channel.U).all() - assert (fieldset.V == ds_fesom_channel.V).all() + assert (fieldset.U.data == ds_fesom_channel.U).all() + assert (fieldset.V.data == ds_fesom_channel.V).all() pset = ParticleSet(fieldset, pclass=Particle) assert pset.fieldset == fieldset @@ -94,8 +94,8 @@ def test_fesom_in_particleset(ds_fesom_channel, uv_fesom_channel): def test_set_interp_methods(ds_fesom_channel, uv_fesom_channel): fieldset = FieldSet([uv_fesom_channel, uv_fesom_channel.U, uv_fesom_channel.V]) # Check that the fieldset has the expected properties - assert (fieldset.U == ds_fesom_channel.U).all() - assert (fieldset.V == ds_fesom_channel.V).all() + assert (fieldset.U.data == ds_fesom_channel.U).all() + assert (fieldset.V.data == ds_fesom_channel.V).all() # Set the interpolation method for each field fieldset.U.interp_method = UXPiecewiseConstantFace diff --git a/tests/v4/test_uxgrid.py b/tests/v4/test_uxgrid.py new file mode 100644 index 0000000000..9c675b278b --- /dev/null +++ b/tests/v4/test_uxgrid.py @@ -0,0 +1,23 @@ +import pytest + +from parcels._datasets.unstructured.generic import datasets as uxdatasets +from parcels.uxgrid import UxGrid + + +@pytest.mark.parametrize("uxds", [pytest.param(uxds, id=key) for key, uxds in uxdatasets.items()]) +def test_uxgrid_init_on_generic_datasets(uxds): + UxGrid(uxds.uxgrid, z=uxds.coords["nz"]) + + +@pytest.mark.parametrize("uxds", [uxdatasets["stommel_gyre_delaunay"]]) +def test_uxgrid_axes(uxds): + grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"]) + assert grid.axes == ["Z", "FACE"] + + +@pytest.mark.parametrize("uxds", [uxdatasets["stommel_gyre_delaunay"]]) +def test_xgrid_get_axis_dim(uxds): + grid = UxGrid(uxds.uxgrid, z=uxds.coords["nz"]) + + assert grid.get_axis_dim("FACE") == 721 + assert grid.get_axis_dim("Z") == 2 diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py index 301a16c131..2af1e1f019 100644 --- a/tests/v4/test_xgrid.py +++ b/tests/v4/test_xgrid.py @@ -6,7 +6,7 @@ from numpy.testing import assert_allclose from parcels import xgcm -from parcels._datasets.structured.generic import T, X, Y, Z, datasets +from parcels._datasets.structured.generic import X, Y, Z, datasets from parcels.xgrid import XGrid, _search_1d_array GridTestCase = namedtuple("GridTestCase", ["Grid", "attr", "expected"]) @@ -16,10 +16,9 @@ GridTestCase(datasets["ds_2d_left"], "lat", datasets["ds_2d_left"].YG.values), GridTestCase(datasets["ds_2d_left"], "depth", datasets["ds_2d_left"].ZG.values), GridTestCase(datasets["ds_2d_left"], "time", datasets["ds_2d_left"].time.values.astype(np.float64) / 1e9), - GridTestCase(datasets["ds_2d_left"], "xdim", X), - GridTestCase(datasets["ds_2d_left"], "ydim", Y), - GridTestCase(datasets["ds_2d_left"], "zdim", Z), - GridTestCase(datasets["ds_2d_left"], "tdim", T), + GridTestCase(datasets["ds_2d_left"], "xdim", X - 1), + GridTestCase(datasets["ds_2d_left"], "ydim", Y - 1), + GridTestCase(datasets["ds_2d_left"], "zdim", Z - 1), ] @@ -45,9 +44,18 @@ def test_xgrid_init_on_generic_datasets(ds): XGrid(xgcm.Grid(ds, periodic=False)) -def test_xgrid_axes(): - # Tests that the xgrid.axes property correctly identifies the axes and ordering - ... +@pytest.mark.parametrize("ds", [datasets["ds_2d_left"]]) +def test_xgrid_axes(ds): + grid = XGrid(xgcm.Grid(ds, periodic=False)) + assert grid.axes == ["Z", "Y", "X"] + + +@pytest.mark.parametrize("ds", [datasets["ds_2d_left"]]) +def test_xgrid_get_axis_dim(ds): + grid = XGrid(xgcm.Grid(ds, periodic=False)) + assert grid.get_axis_dim("Z") == Z - 1 + assert grid.get_axis_dim("Y") == Y - 1 + assert grid.get_axis_dim("X") == X - 1 def test_invalid_xgrid_field_array(): @@ -85,30 +93,6 @@ def test_invalid_lon_lat(): XGrid(xgcm.Grid(ds, periodic=False)) -def test_xgrid_ravel_unravel_index(): - ds = datasets["ds_2d_left"] - grid = XGrid(xgcm.Grid(ds, periodic=False)) - - xdim = grid.xdim - ydim = grid.ydim - zdim = grid.zdim - - encountered_eis = [] - for xi in range(xdim): - for yi in range(ydim): - for zi in range(zdim): - axis_indices = {"Z": zi, "Y": yi, "X": xi} - ei = grid.ravel_index(axis_indices) - axis_indices_test = grid.unravel_index(ei) - assert axis_indices_test == axis_indices - encountered_eis.append(ei) - - encountered_eis = sorted(encountered_eis) - assert len(set(encountered_eis)) == len(encountered_eis), "Raveled indices are not unique." - assert np.allclose(np.diff(np.array(encountered_eis)), 1), "Raveled indices are not consecutive integers." - assert encountered_eis[0] == 0, "Raveled indices do not start at 0." - - @pytest.mark.parametrize( "ds", [