diff --git a/parcels/field.py b/parcels/field.py index dba61dbd6a..1a1fbb7175 100644 --- a/parcels/field.py +++ b/parcels/field.py @@ -29,7 +29,7 @@ _raise_field_out_of_bound_error, ) from parcels.uxgrid import UxGrid -from parcels.xgrid import XGrid +from parcels.xgrid import XGrid, _transpose_xfield_data_to_tzyx from ._index_search import _search_time_index @@ -146,6 +146,9 @@ def __init__( _assert_compatible_combination(data, grid) + if isinstance(grid, XGrid): + data = _transpose_xfield_data_to_tzyx(data, grid.xgcm_grid) + self.name = name self.data = data self.grid = grid @@ -186,8 +189,9 @@ def __init__( else: raise ValueError("Unsupported mesh type in data array attributes. Choose either: 'spherical' or 'flat'") - if "time" not in self.data.dims: - raise ValueError("Field is missing a 'time' dimension. ") + if self.data.shape[0] > 1: + if "time" not in self.data.coords: + raise ValueError("Field data is missing a 'time' coordinate.") @property def units(self): @@ -439,7 +443,7 @@ def _assert_compatible_combination(data: xr.DataArray | ux.UxDataArray, grid: ux def _get_time_interval(data: xr.DataArray | ux.UxDataArray) -> TimeInterval | None: - if len(data.time) == 1: + if data.shape[0] == 1: return None return TimeInterval(data.time.values[0], data.time.values[-1]) diff --git a/parcels/fieldset.py b/parcels/fieldset.py index c423000fc4..5937c7f573 100644 --- a/parcels/fieldset.py +++ b/parcels/fieldset.py @@ -132,7 +132,6 @@ def add_constant_field(self, name: str, value, mesh: Mesh = "flat"): """ da = xr.DataArray( data=np.full((1, 1, 1, 1), value), - dims=["time", "ZG", "YG", "XG"], ) grid = XGrid(xgcm.Grid(da)) self.add_field( diff --git a/parcels/xgrid.py b/parcels/xgrid.py index 2cfda507e7..81391ee818 100644 --- a/parcels/xgrid.py +++ b/parcels/xgrid.py @@ -1,4 +1,4 @@ -from collections.abc import Hashable, Mapping +from collections.abc import Hashable, Mapping, Sequence from functools import cached_property from typing import Literal, cast @@ -10,13 +10,17 @@ from parcels._index_search import _search_indices_curvilinear_2d from parcels.basegrid import BaseGrid -_XGRID_AXES_ORDERING = "ZYX" _XGRID_AXES = Literal["X", "Y", "Z"] +_XGRID_AXES_ORDERING: Sequence[_XGRID_AXES] = "ZYX" _XGCM_AXIS_DIRECTION = Literal["X", "Y", "Z", "T"] _XGCM_AXIS_POSITION = Literal["center", "left", "right", "inner", "outer"] _XGCM_AXES = Mapping[_XGCM_AXIS_DIRECTION, xgcm.Axis] +_FIELD_DATA_ORDERING: Sequence[_XGCM_AXIS_DIRECTION] = "TZYX" + +_DEFAULT_XGCM_KWARGS = {"periodic": False} + def get_cell_count_along_dim(axis: xgcm.Axis) -> int: first_coord = list(axis.coords.items())[0] @@ -34,6 +38,48 @@ def _get_xgrid_axes(grid: xgcm.Grid) -> list[_XGRID_AXES]: return sorted(spatial_axes, key=_XGRID_AXES_ORDERING.index) +def _drop_field_data(ds: xr.Dataset) -> xr.Dataset: + """ + Removes DataArrays from the dataset that are associated with field data so that + when passed to the XGCM grid, the object only functions as an in memory representation + of the grid. + """ + return ds.drop_vars(ds.data_vars) + + +def _transpose_xfield_data_to_tzyx(da: xr.DataArray, xgcm_grid: xgcm.Grid) -> xr.DataArray: + """ + Transpose a DataArray of any shape into a 4D array of order TZYX. Uses xgcm to determine + the axes, and inserts mock dimensions of size 1 for any axes not present in the DataArray. + """ + ax_dims = [(get_axis_from_dim_name(xgcm_grid.axes, dim), dim) for dim in da.dims] + + if all(ax_dim[0] is None for ax_dim in ax_dims): + # Assuming its a 1D constant field (hence has no axes) + assert da.shape == (1, 1, 1, 1) + return da.rename({old_dim: f"mock{axis}" for old_dim, axis in zip(da.dims, _FIELD_DATA_ORDERING, strict=True)}) + + # All dimensions must be associated with an axis in the grid + if any(ax_dim[0] is None for ax_dim in ax_dims): + raise ValueError( + f"DataArray {da.name!r} with dims {da.dims} has dimensions that are not associated with a direction on the provided grid." + ) + + axes_not_in_field = set(_FIELD_DATA_ORDERING) - set(ax_dim[0] for ax_dim in ax_dims) + + mock_dims_to_create = {} + for ax in axes_not_in_field: + mock_dims_to_create[f"mock{ax}"] = 1 + ax_dims.append((ax, f"mock{ax}")) + + if mock_dims_to_create: + da = da.expand_dims(mock_dims_to_create, create_index_for_new_dim=False) + + ax_dims = sorted(ax_dims, key=lambda x: _FIELD_DATA_ORDERING.index(x[0])) + + return da.transpose(*[ax_dim[1] for ax_dim in ax_dims]) + + class XGrid(BaseGrid): """ Class to represent a structured grid in Parcels. Wraps a xgcm-like Grid object (we use a trimmed down version of the xgcm.Grid class that is vendored with Parcels). @@ -53,6 +99,18 @@ def __init__(self, grid: xgcm.Grid, mesh="flat"): if len(set(grid.axes) & {"X", "Y", "Z"}) > 0: # Only if spatial grid is >0D (see #2054 for further development) assert_valid_lat_lon(ds["lat"], ds["lon"], grid.axes) + @classmethod + def from_dataset(cls, ds: xr.Dataset, mesh="flat", xgcm_kwargs=None): + """WARNING: unstable API, subject to change in future versions.""" # TODO v4: make private or remove warning on v4 release + if xgcm_kwargs is None: + xgcm_kwargs = {} + + xgcm_kwargs = {**_DEFAULT_XGCM_KWARGS, **xgcm_kwargs} + + ds = _drop_field_data(ds) + grid = xgcm.Grid(ds, **xgcm_kwargs) + return cls(grid, mesh=mesh) + @property def axes(self) -> list[_XGRID_AXES]: return _get_xgrid_axes(self.xgcm_grid) diff --git a/tests/utils.py b/tests/utils.py index e0fd5f34fe..a76b7e8a6e 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,12 +1,19 @@ """General helper functions and utilies for test suite.""" +from __future__ import annotations + from pathlib import Path +from typing import TYPE_CHECKING import numpy as np import xarray as xr import parcels from parcels import FieldSet +from parcels.xgrid import _FIELD_DATA_ORDERING, get_axis_from_dim_name + +if TYPE_CHECKING: + from parcels.xgrid import XGrid PROJECT_ROOT = Path(__file__).resolve().parents[1] TEST_ROOT = PROJECT_ROOT / "tests" @@ -116,3 +123,13 @@ def create_fieldset_zeros_simple(xdim=40, ydim=100, withtime=False): def assert_empty_folder(path: Path): assert [p.name for p in path.iterdir()] == [] + + +def assert_valid_field_data(data: xr.DataArray, grid: XGrid): + assert len(data.shape) == 4, f"Field data should have 4 dimensions (time, depth, lat, lon), got dims {data.dims}" + + for ax_expected, dim in zip(_FIELD_DATA_ORDERING, data.dims, strict=True): + ax_actual = get_axis_from_dim_name(grid.xgcm_grid.axes, dim) + if ax_actual is None: + continue # None is ok + assert ax_actual == ax_expected, f"Expected axis {ax_expected} for dimension '{dim}', got {ax_actual}" diff --git a/tests/v4/test_datasets.py b/tests/v4/test_datasets.py index e5201e58b7..d2d21040bb 100644 --- a/tests/v4/test_datasets.py +++ b/tests/v4/test_datasets.py @@ -1,11 +1,11 @@ +from parcels import xgcm from parcels._datasets.structured.generic import datasets -from parcels.xgcm import Grid def test_left_indexed_dataset(): """Checks that 'ds_2d_left' is right indexed on all variables.""" ds = datasets["ds_2d_left"] - grid = Grid(ds) + grid = xgcm.Grid(ds) for _axis_name, axis in grid.axes.items(): for pos, _dim_name in axis.coords.items(): @@ -15,7 +15,7 @@ def test_left_indexed_dataset(): def test_right_indexed_dataset(): """Checks that 'ds_2d_right' is right indexed on all variables.""" ds = datasets["ds_2d_right"] - grid = Grid(ds) + grid = xgcm.Grid(ds) for _axis_name, axis in grid.axes.items(): for pos, _dim_name in axis.coords.items(): assert pos in ["center", "right"] diff --git a/tests/v4/test_field.py b/tests/v4/test_field.py index a555a672f2..6f1068896a 100644 --- a/tests/v4/test_field.py +++ b/tests/v4/test_field.py @@ -5,7 +5,7 @@ import uxarray as ux import xarray as xr -from parcels import Field, UXPiecewiseConstantFace, UXPiecewiseLinearNode, VectorField, xgcm +from parcels import Field, UXPiecewiseConstantFace, UXPiecewiseLinearNode, VectorField from parcels._datasets.structured.generic import T as T_structured from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.unstructured.generic import datasets as datasets_unstructured @@ -15,7 +15,7 @@ def test_field_init_param_types(): data = datasets_structured["ds_2d_left"] - grid = XGrid(xgcm.Grid(data)) + grid = XGrid.from_dataset(data) with pytest.raises(ValueError, match="Expected `name` to be a string"): Field(name=123, data=data["data_g"], grid=grid) @@ -32,7 +32,7 @@ def test_field_init_param_types(): @pytest.mark.parametrize( "data,grid", [ - pytest.param(ux.UxDataArray(), XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])), id="uxdata-grid"), + pytest.param(ux.UxDataArray(), XGrid.from_dataset(datasets_structured["ds_2d_left"]), id="uxdata-grid"), pytest.param( xr.DataArray(), UxGrid( @@ -57,7 +57,7 @@ def test_field_incompatible_combination(data, grid): [ pytest.param( datasets_structured["ds_2d_left"]["data_g"], - XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])), + XGrid.from_dataset(datasets_structured["ds_2d_left"]), id="ds_2d_left", ), # TODO: Perhaps this test should be expanded to cover more datasets? ], @@ -80,10 +80,10 @@ def test_field_init_fail_on_float_time_dim(): (users are expected to use timedelta64 or datetime). """ ds = datasets_structured["ds_2d_left"].copy() - ds["time"] = np.arange(0, T_structured, dtype="float64") + ds["time"] = (ds["time"].dims, np.arange(0, T_structured, dtype="float64"), ds["time"].attrs) data = ds["data_g"] - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) with pytest.raises( ValueError, match="Error getting time interval.*. Are you sure that the time dimension on the xarray dataset is stored as timedelta, datetime or cftime datetime objects\?", @@ -100,7 +100,7 @@ def test_field_init_fail_on_float_time_dim(): [ pytest.param( datasets_structured["ds_2d_left"]["data_g"], - XGrid(xgcm.Grid(datasets_structured["ds_2d_left"])), + XGrid.from_dataset(datasets_structured["ds_2d_left"]), id="ds_2d_left", ), ], @@ -119,7 +119,7 @@ def test_vectorfield_init_different_time_intervals(): def test_field_invalid_interpolator(): ds = datasets_structured["ds_2d_left"] - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid): return 0.0 @@ -131,7 +131,7 @@ def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, inval def test_vectorfield_invalid_interpolator(): ds = datasets_structured["ds_2d_left"] - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) def invalid_interpolator_wrong_signature(self, ti, position, tau, t, z, y, invalid): return 0.0 diff --git a/tests/v4/test_fieldset.py b/tests/v4/test_fieldset.py index f623b7f1c4..c3871018be 100644 --- a/tests/v4/test_fieldset.py +++ b/tests/v4/test_fieldset.py @@ -5,7 +5,6 @@ import pytest import xarray as xr -from parcels import xgcm from parcels._datasets.structured.circulation_models import ( datasets as datasets_circulation_models, # noqa: F401 ) # just making sure the import works. Will eventually be used in tests @@ -14,6 +13,7 @@ from parcels.field import Field, VectorField from parcels.fieldset import CalendarError, FieldSet, _datetime_to_msg from parcels.xgrid import XGrid +from tests import utils ds = datasets_structured["ds_2d_left"] @@ -21,7 +21,7 @@ @pytest.fixture def fieldset() -> FieldSet: """Fixture to create a FieldSet object for testing.""" - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) U = Field("U", ds["U (A grid)"], grid, mesh_type="flat") V = Field("V", ds["V (A grid)"], grid, mesh_type="flat") UV = VectorField("UV", U, V) @@ -55,7 +55,7 @@ def test_fieldset_add_constant_field(fieldset): def test_fieldset_add_field(fieldset): - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) field = Field("test_field", ds["U (A grid)"], grid, mesh_type="flat") fieldset.add_field(field) assert fieldset.test_field == field @@ -68,7 +68,7 @@ def test_fieldset_add_field_wrong_type(fieldset): def test_fieldset_add_field_already_exists(fieldset): - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) field = Field("test_field", ds["U (A grid)"], grid, mesh_type="flat") fieldset.add_field(field, "test_field") with pytest.raises(ValueError, match="FieldSet already has a Field with name 'test_field'"): @@ -85,16 +85,32 @@ def test_fieldset_gridset(fieldset): assert len(fieldset.gridset) == 2 +@pytest.mark.parametrize("ds", [pytest.param(ds, id=k) for k, ds in datasets_structured.items()]) +def test_fieldset_from_structured_generic_datasets(ds): + grid = XGrid.from_dataset(ds) + fields = [] + for var in ds.data_vars: + fields.append(Field(var, ds[var], grid, mesh_type="flat")) + + fieldset = FieldSet(fields) + + assert len(fieldset.fields) == len(ds.data_vars) + for field in fieldset.fields.values(): + utils.assert_valid_field_data(field.data, field.grid) + + assert len(fieldset.gridset) == 1 + + def test_fieldset_gridset_multiple_grids(): ... def test_fieldset_time_interval(): - grid1 = XGrid(xgcm.Grid(ds)) + grid1 = XGrid.from_dataset(ds) field1 = Field("field1", ds["U (A grid)"], grid1, mesh_type="flat") ds2 = ds.copy() - ds2["time"] = ds2["time"] + np.timedelta64(timedelta(days=1)) - grid2 = XGrid(xgcm.Grid(ds2)) + ds2["time"] = (ds2["time"].dims, ds2["time"].data + np.timedelta64(timedelta(days=1)), ds2["time"].attrs) + grid2 = XGrid.from_dataset(ds2) field2 = Field("field2", ds2["U (A grid)"], grid2, mesh_type="flat") fieldset = FieldSet([field1, field2]) @@ -114,16 +130,24 @@ def test_fieldset_time_interval_constant_fields(): def test_fieldset_init_incompatible_calendars(): ds1 = ds.copy() - ds1["time"] = xr.date_range("2000", "2001", T_structured, calendar="365_day", use_cftime=True) + ds1["time"] = ( + ds1["time"].dims, + xr.date_range("2000", "2001", T_structured, calendar="365_day", use_cftime=True), + ds1["time"].attrs, + ) - grid = XGrid(xgcm.Grid(ds1)) + grid = XGrid.from_dataset(ds1) U = Field("U", ds1["U (A grid)"], grid, mesh_type="flat") V = Field("V", ds1["V (A grid)"], grid, mesh_type="flat") UV = VectorField("UV", U, V) ds2 = ds.copy() - ds2["time"] = xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True) - grid2 = XGrid(xgcm.Grid(ds2)) + ds2["time"] = ( + ds2["time"].dims, + xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True), + ds2["time"].attrs, + ) + grid2 = XGrid.from_dataset(ds2) incompatible_calendar = Field("test", ds2["data_g"], grid2, mesh_type="flat") with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): @@ -132,16 +156,24 @@ def test_fieldset_init_incompatible_calendars(): def test_fieldset_add_field_incompatible_calendars(fieldset): ds_test = ds.copy() - ds_test["time"] = xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True) - grid = XGrid(xgcm.Grid(ds_test)) + ds_test["time"] = ( + ds_test["time"].dims, + xr.date_range("2000", "2001", T_structured, calendar="360_day", use_cftime=True), + ds_test["time"].attrs, + ) + grid = XGrid.from_dataset(ds_test) field = Field("test_field", ds_test["data_g"], grid, mesh_type="flat") with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): fieldset.add_field(field, "test_field") ds_test = ds.copy() - ds_test["time"] = np.linspace(0, 100, T_structured, dtype="timedelta64[s]") - grid = XGrid(xgcm.Grid(ds_test)) + ds_test["time"] = ( + ds_test["time"].dims, + np.linspace(0, 100, T_structured, dtype="timedelta64[s]"), + ds_test["time"].attrs, + ) + grid = XGrid.from_dataset(ds_test) field = Field("test_field", ds_test["data_g"], grid, mesh_type="flat") with pytest.raises(CalendarError, match="Expected field '.*' to have calendar compatible with datetime object"): diff --git a/tests/v4/test_index_search.py b/tests/v4/test_index_search.py index 7f9290f12c..5d5c602b94 100644 --- a/tests/v4/test_index_search.py +++ b/tests/v4/test_index_search.py @@ -1,7 +1,6 @@ import numpy as np import pytest -from parcels import xgcm from parcels._datasets.structured.generic import datasets from parcels._index_search import _search_indices_curvilinear_2d from parcels.field import Field @@ -13,7 +12,7 @@ @pytest.fixture def field_cone(): ds = datasets["2d_left_unrolled_cone"] - grid = XGrid(xgcm.Grid(ds, periodic=False)) + grid = XGrid.from_dataset(ds) field = Field( name="test_field", data=ds["data_g"], diff --git a/tests/v4/test_kernel.py b/tests/v4/test_kernel.py index 1125d925cc..cb521bc787 100644 --- a/tests/v4/test_kernel.py +++ b/tests/v4/test_kernel.py @@ -6,7 +6,6 @@ Field, FieldSet, ParticleSet, - xgcm, ) from parcels._datasets.structured.generic import datasets as datasets_structured from parcels.xgrid import XGrid @@ -16,7 +15,7 @@ @pytest.fixture def fieldset() -> FieldSet: ds = datasets_structured["ds_2d_left"] - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) U = Field("U", ds["U (A grid)"], grid, mesh_type="flat") V = Field("V", ds["V (A grid)"], grid, mesh_type="flat") return FieldSet([U, V]) diff --git a/tests/v4/test_particleset.py b/tests/v4/test_particleset.py index 7a9dd7be64..ace5e8dba6 100644 --- a/tests/v4/test_particleset.py +++ b/tests/v4/test_particleset.py @@ -13,7 +13,6 @@ ParticleSet, ParticleSetWarning, Variable, - xgcm, ) from parcels._datasets.structured.generic import datasets as datasets_structured from parcels.xgrid import XGrid @@ -23,7 +22,7 @@ @pytest.fixture def fieldset() -> FieldSet: ds = datasets_structured["ds_2d_left"] - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) U = Field("U", ds["U (A grid)"], grid, mesh_type="flat") V = Field("V", ds["V (A grid)"], grid, mesh_type="flat") return FieldSet([U, V]) diff --git a/tests/v4/test_particleset_execute.py b/tests/v4/test_particleset_execute.py index f63210add0..b1994d4d4e 100644 --- a/tests/v4/test_particleset_execute.py +++ b/tests/v4/test_particleset_execute.py @@ -10,7 +10,6 @@ StatusCode, UXPiecewiseConstantFace, VectorField, - xgcm, ) from parcels._datasets.structured.generic import datasets as datasets_structured from parcels._datasets.unstructured.generic import datasets as datasets_unstructured @@ -22,7 +21,7 @@ @pytest.fixture def fieldset() -> FieldSet: ds = datasets_structured["ds_2d_left"] - grid = XGrid(xgcm.Grid(ds)) + grid = XGrid.from_dataset(ds) U = Field("U", ds["U (A grid)"], grid, mesh_type="flat") V = Field("V", ds["V (A grid)"], grid, mesh_type="flat") return FieldSet([U, V]) diff --git a/tests/v4/test_xgrid.py b/tests/v4/test_xgrid.py index 210ca6b742..30bcc7c1aa 100644 --- a/tests/v4/test_xgrid.py +++ b/tests/v4/test_xgrid.py @@ -1,3 +1,4 @@ +import itertools from collections import namedtuple import numpy as np @@ -5,11 +6,11 @@ import xarray as xr from numpy.testing import assert_allclose -from parcels import xgcm from parcels._datasets.structured.generic import X, Y, Z, datasets -from parcels.xgrid import XGrid, _search_1d_array +from parcels.xgrid import XGrid, _search_1d_array, _transpose_xfield_data_to_tzyx +from tests import utils -GridTestCase = namedtuple("GridTestCase", ["Grid", "attr", "expected"]) +GridTestCase = namedtuple("GridTestCase", ["ds", "attr", "expected"]) test_cases = [ GridTestCase(datasets["ds_2d_left"], "lon", datasets["ds_2d_left"].XG.values), @@ -34,25 +35,39 @@ def assert_equal(actual, expected): @pytest.mark.parametrize("ds, attr, expected", test_cases) def test_xgrid_properties_ground_truth(ds, attr, expected): - grid = XGrid(xgcm.Grid(ds, periodic=False)) + grid = XGrid.from_dataset(ds) actual = getattr(grid, attr) assert_equal(actual, expected) @pytest.mark.parametrize("ds", [pytest.param(ds, id=key) for key, ds in datasets.items()]) -def test_xgrid_init_on_generic_datasets(ds): - XGrid(xgcm.Grid(ds, periodic=False)) +def test_xgrid_from_dataset_on_generic_datasets(ds): + XGrid.from_dataset(ds) @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]]) def test_xgrid_axes(ds): - grid = XGrid(xgcm.Grid(ds, periodic=False)) + grid = XGrid.from_dataset(ds) assert grid.axes == ["Z", "Y", "X"] +@pytest.mark.parametrize("ds", [datasets["ds_2d_left"]]) +def test_transpose_xfield_data_to_tzyx(ds): + da = ds["data_g"] + grid = XGrid.from_dataset(ds) + + all_combinations = (itertools.combinations(da.dims, n) for n in range(len(da.dims))) + all_combinations = itertools.chain(*all_combinations) + for subset_dims in all_combinations: + isel = {dim: 0 for dim in subset_dims} + da_subset = da.isel(isel, drop=True) + da_test = _transpose_xfield_data_to_tzyx(da_subset, grid.xgcm_grid) + utils.assert_valid_field_data(da_test, grid) + + @pytest.mark.parametrize("ds", [datasets["ds_2d_left"]]) def test_xgrid_get_axis_dim(ds): - grid = XGrid(xgcm.Grid(ds, periodic=False)) + grid = XGrid.from_dataset(ds) assert grid.get_axis_dim("Z") == Z - 1 assert grid.get_axis_dim("Y") == Y - 1 assert grid.get_axis_dim("X") == X - 1 @@ -72,7 +87,7 @@ def test_invalid_lon_lat(): ValueError, match=".*is defined on the center of the grid, but must be defined on the F points\.", ): - XGrid(xgcm.Grid(ds, periodic=False)) + XGrid.from_dataset(ds) ds = datasets["ds_2d_left"].copy() ds["lon"], _ = xr.broadcast(ds["YG"], ds["XG"]) @@ -80,7 +95,7 @@ def test_invalid_lon_lat(): ValueError, match=".*have different dimensionalities\.", ): - XGrid(xgcm.Grid(ds, periodic=False)) + XGrid.from_dataset(ds) ds = datasets["ds_2d_left"].copy() ds["lon"], ds["lat"] = xr.broadcast(ds["YG"], ds["XG"]) @@ -90,7 +105,7 @@ def test_invalid_lon_lat(): ValueError, match=".*must be defined on the X and Y axes and transposed to have dimensions in order of Y, X\.", ): - XGrid(xgcm.Grid(ds, periodic=False)) + XGrid.from_dataset(ds) @pytest.mark.parametrize( @@ -101,7 +116,7 @@ def test_invalid_lon_lat(): ], ) # for key, ds in datasets.items()]) def test_xgrid_search_cpoints(ds): - grid = XGrid(xgcm.Grid(ds, periodic=False)) + grid = XGrid.from_dataset(ds) lat_array, lon_array = get_2d_fpoint_mesh(grid) lat_array, lon_array = corner_to_cell_center_points(lat_array, lon_array) @@ -148,10 +163,10 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi): @pytest.mark.parametrize( - "grid, da_name, expected", + "ds, da_name, expected", [ pytest.param( - XGrid(xgcm.Grid(datasets["ds_2d_left"], periodic=False)), + datasets["ds_2d_left"], "U (C grid)", { "XG": (np.int64(0), np.float64(0.0)), @@ -161,7 +176,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi): id="MITgcm indexing style U (C grid)", ), pytest.param( - XGrid(xgcm.Grid(datasets["ds_2d_left"], periodic=False)), + datasets["ds_2d_left"], "V (C grid)", { "XC": (np.int64(-1), np.float64(0.5)), @@ -171,7 +186,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi): id="MITgcm indexing style V (C grid)", ), pytest.param( - XGrid(xgcm.Grid(datasets["ds_2d_right"], periodic=False)), + datasets["ds_2d_right"], "U (C grid)", { "XG": (np.int64(0), np.float64(0.0)), @@ -181,7 +196,7 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi): id="NEMO indexing style U (C grid)", ), pytest.param( - XGrid(xgcm.Grid(datasets["ds_2d_right"], periodic=False)), + datasets["ds_2d_right"], "V (C grid)", { "XC": (np.int64(0), np.float64(0.5)), @@ -192,10 +207,11 @@ def test_search_1d_array(array, x, expected_xi, expected_xsi): ), ], ) -def test_xgrid_localize_zero_position(grid, da_name, expected): +def test_xgrid_localize_zero_position(ds, da_name, expected): """Test localize function using left and right datasets.""" + grid = XGrid.from_dataset(ds) + da = ds[da_name] position = grid.search(0, 0, 0) - da = grid.xgcm_grid._ds[da_name] local_position = grid.localize(position, da.dims) assert local_position == expected, f"Expected {expected}, got {local_position}"