From 9292255453bcbe3912ab09ee63a1f71f868cde53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B7=AF=E4=BB=8E=E4=BB=8A=E5=A4=9C=E7=99=BD?= <1174273035@qq.com> Date: Thu, 22 Oct 2026 00:03:42 +0800 Subject: [PATCH 1/3] Fix structured Dataset mapping to preserve Grid face order --- test/io/test_structured_face_order.py | 125 ++++++++++++++++++++++++++ uxarray/core/dataset.py | 7 +- uxarray/core/utils.py | 20 ++--- 3 files changed, 140 insertions(+), 12 deletions(-) create mode 100644 test/io/test_structured_face_order.py diff --git a/test/io/test_structured_face_order.py b/test/io/test_structured_face_order.py new file mode 100644 index 000000000..6d0dbaf99 --- /dev/null +++ b/test/io/test_structured_face_order.py @@ -0,0 +1,125 @@ +"""Synthetic regressions for structured data-to-face alignment (GH #1828).""" + +import numpy as np +import pytest +import xarray as xr + +import uxarray as ux +from uxarray.errors import DimensionError + + +def _structured_dataset(reverse_variables=False, time_order=None, lazy=False): + """Return equal fields with deliberately different spatial dimension orders.""" + ds = xr.Dataset( + { + "yfirst": (("y", "x"), [[0.0, 20.0, 40.0], [100.0, 120.0, 140.0]]), + "xfirst": (("x", "y"), [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), + }, + coords={ + "x": ("x", [0.0, 20.0, 40.0], {"standard_name": "longitude"}), + "y": ("y", [-10.0, 50.0], {"standard_name": "latitude"}), + }, + attrs={"synthetic": "GH1828"}, + ) + if reverse_variables: + ds = ds[["xfirst", "yfirst"]] + if time_order is not None: + ds = xr.concat([ds, ds + 1000.0], dim=xr.IndexVariable("time", [0, 1])) + ds = ds.transpose(*time_order) + if lazy: + ds = ds.chunk({"x": 1, "y": 1}) + return ds + + +def _assert_face_alignment(mapped, original): + """Check each value against geometry, not just another flattened dataset.""" + grid = mapped.uxgrid + assert grid._source_dims_dict == {"n_face": ("x", "y")} + # The six non-polar faces have four nodes. Read their locations to establish + # the physical row-major face order independently of the mapping helper. + nodes = grid.face_node_connectivity.values + np.testing.assert_allclose( + grid.node_lon.values[nodes].mean(axis=1), np.tile(original.x.values, 2) + ) + np.testing.assert_allclose( + grid.node_lat.values[nodes].mean(axis=1), np.repeat(original.y.values, 3) + ) + for name in original.data_vars: + other_dims = [dim for dim in original[name].dims if dim not in ("x", "y")] + field = original[name].transpose(*other_dims, "y", "x") + expected = field.values.reshape(*field.shape[:-2], 6) + actual = mapped[name].transpose(*other_dims, "n_face") + np.testing.assert_array_equal(actual.values, expected) + np.testing.assert_allclose( + actual.integrate().values, + np.sum(expected * grid.face_areas.values, axis=-1), + ) + + +@pytest.mark.parametrize("constructor", ["from_xarray", "from_structured"]) +@pytest.mark.parametrize("reverse_variables", [False, True]) +@pytest.mark.parametrize("time_order", [None, ("time", "y", "x"), ("x", "time", "y")]) +@pytest.mark.parametrize("lazy", [False, True]) +def test_structured_face_order(constructor, reverse_variables, time_order, lazy): + """Variable insertion order, axis order, and chunks cannot change face values.""" + ds = _structured_dataset(reverse_variables, time_order, lazy) + original = ds.copy(deep=True) + grid = ux.Grid.from_structured(lon=ds.x, lat=ds.y) + if constructor == "from_xarray": + mapped = ux.UxDataset.from_xarray(ds, uxgrid=grid) + assert mapped.uxgrid is grid + else: + mapped = ux.UxDataset.from_structured(ds) + if lazy: + assert all(mapped[name].chunks is not None for name in ds.data_vars) + xr.testing.assert_identical(ds, original) + _assert_face_alignment(mapped, original) + + +def test_structured_variable_order_preserves_face_index(): + """Equal datasets must retain equal n_face indexes after conversion.""" + ds0 = _structured_dataset() + ds1 = _structured_dataset(reverse_variables=True) + assert ds0.equals(ds1) + grid = ux.Grid.from_structured(lon=ds0.x, lat=ds0.y) + mapped0 = ux.UxDataset.from_xarray(ds0, uxgrid=grid) + mapped1 = ux.UxDataset.from_xarray(ds1, uxgrid=grid) + assert mapped0.indexes["n_face"].equals(mapped1.indexes["n_face"]) + for name in ds0.data_vars: + xr.testing.assert_identical(mapped0[name], mapped1[name]) + + +@pytest.mark.parametrize("transpose", [False, True]) +def test_structured_netcdf_face_selection(tmp_path, transpose): + """Read a member file, attach the same Grid, and select the correct face.""" + ds = _structured_dataset() + if transpose: + ds = ds.transpose("x", "y") + grid = ux.Grid.from_structured(lon=ds.x, lat=ds.y) + path = tmp_path / "synthetic-member.nc" + ds.to_netcdf(path) + with xr.open_dataset(path) as reopened: + mapped = ux.UxDataset.from_xarray(reopened, uxgrid=grid) + _assert_face_alignment(mapped, reopened) + selected = mapped["yfirst"].isel(n_face=[1, 4]) + np.testing.assert_array_equal(selected.values, [20.0, 120.0]) + assert selected.uxgrid.n_face == 2 + + +@pytest.mark.parametrize("missing", ["x", "y", "both"]) +def test_structured_missing_source_dimension_raises(missing): + """Do not silently produce an unflattened result for an incompatible Grid.""" + ds = _structured_dataset() + grid = ux.Grid.from_structured(lon=ds.x, lat=ds.y) + ds = ds.isel({dim: 0 for dim in ("x", "y") if missing in (dim, "both")}) + with pytest.raises(DimensionError, match="source spatial dimensions"): + ux.UxDataset.from_xarray(ds, uxgrid=grid) + + +def test_structured_coordinate_name_mapping_mismatch_raises(): + """Keep bad source metadata visible instead of silently retaining x/y.""" + ds = _structured_dataset().rename_vars({"x": "longitude", "y": "latitude"}) + # Grid.from_dataset currently records coordinate names rather than dimension + # names for this input. Correcting that metadata is separate from GH1828. + with pytest.raises(DimensionError, match="source spatial dimensions"): + ux.UxDataset.from_structured(ds) diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index da3e73102..04173b5db 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -322,7 +322,12 @@ def from_xarray(cls, ds: xr.Dataset, uxgrid: Grid = None, ugrid_dims: dict = Non Parameters ---------- ds: xr.Dataset - An Xarray dataset containing data residing on an unstructured grid + An Xarray dataset containing data residing on an unstructured grid, + or structured data paired with a ``Grid`` that records its source + longitude and latitude dimensions. Structured dimensions are stacked + onto ``n_face`` in the Grid's latitude-major, longitude-minor face + order, independently of variable or dimension order. Non-spatial + dimensions are retained. uxgrid: Grid, optional ``Grid`` object representing an unstructured grid. If a grid is not provided, the source ds will be parsed to see if a ``Grid`` can be constructed. diff --git a/uxarray/core/utils.py b/uxarray/core/utils.py index 84c0a2dc6..1b449325b 100644 --- a/uxarray/core/utils.py +++ b/uxarray/core/utils.py @@ -61,18 +61,16 @@ def _map_dims_to_ugrid( "nCell": "n_face")""" if grid.source_grid_spec == "Structured": - # Case for structured grids, flatten bottom two sptial dimensions - lon_name, lat_name = _source_dims_dict["n_face"] - - for var_name in ds.data_vars: - if lon_name in ds[var_name].dims and lat_name in ds[var_name].dims: - if ds[var_name].dims[-1] == lon_name: - dim_ordered = [lat_name, lon_name] - else: - dim_ordered = [lon_name, lat_name] - - ds = ds.stack(n_face=dim_ordered) + if lon_name not in ds.dims or lat_name not in ds.dims: + raise DimensionError( + "Structured data must contain the Grid's source spatial " + f"dimensions {lon_name!r} and {lat_name!r}; got {tuple(ds.dims)!r}." + ) + # Structured Grid faces are latitude-major, longitude-minor. The source + # mapping stores (longitude, latitude), not the stacking order. Match + # the geometry regardless of each variable's axis or insertion order. + ds = ds.stack(n_face=(lat_name, lon_name)) elif grid.source_grid_spec == "GEOS-CS": # stack dimensions to flatten them to map to nodes or faces From 1e03ce32fa8abc48b617e01c0e2ed2de034c4643 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B7=AF=E4=BB=8E=E4=BB=8A=E5=A4=9C=E7=99=BD?= <1174273035@qq.com> Date: Thu, 22 Oct 2026 22:44:34 +0800 Subject: [PATCH 2/3] docs: illustrate structured ugrid_dims mapping --- uxarray/core/dataset.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/uxarray/core/dataset.py b/uxarray/core/dataset.py index 04173b5db..491b407c1 100644 --- a/uxarray/core/dataset.py +++ b/uxarray/core/dataset.py @@ -332,7 +332,10 @@ def from_xarray(cls, ds: xr.Dataset, uxgrid: Grid = None, ugrid_dims: dict = Non ``Grid`` object representing an unstructured grid. If a grid is not provided, the source ds will be parsed to see if a ``Grid`` can be constructed. ugrid_dims: dict, optional - A dictionary mapping dataset dimensions to UGRID dimensions. + A dictionary describing the mapping between dataset and UGRID dimensions. + For structured grids, for example, ``{"n_face": ("x", "y")}`` uses + the UGRID dimension as the key and the dataset's longitude and + latitude dimensions as the tuple value, in that order. Returns ------- From ccbeb01c195f22f7f5dcd777f03e68e5207a701d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B7=AF=E4=BB=8E=E4=BB=8A=E5=A4=9C=E7=99=BD?= <1174273035@qq.com> Date: Fri, 23 Oct 2026 00:30:36 +0800 Subject: [PATCH 3/3] test: cover explicit structured dimension mappings Cover renamed axes with rectangular and square grids and transposed storage. Clarify that missing dimensions refer to the selected mapping, without changing the guard or stacking behavior. --- test/io/test_structured_face_order.py | 35 +++++++++++++++++++++++++-- uxarray/core/utils.py | 2 +- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/test/io/test_structured_face_order.py b/test/io/test_structured_face_order.py index 6d0dbaf99..9088e2be6 100644 --- a/test/io/test_structured_face_order.py +++ b/test/io/test_structured_face_order.py @@ -89,6 +89,37 @@ def test_structured_variable_order_preserves_face_index(): xr.testing.assert_identical(mapped0[name], mapped1[name]) +@pytest.mark.parametrize("square", [False, True]) +@pytest.mark.parametrize("transpose", [False, True]) +def test_structured_explicit_mapping_with_renamed_dimensions(square, transpose): + """Explicit mappings may use names unlike the Grid's source dimensions.""" + lon = xr.DataArray([0.0, 20.0, 40.0], dims="x") + lat = xr.DataArray([-10.0, 20.0, 50.0] if square else [-10.0, 50.0], dims="y") + grid = ux.Grid.from_structured(lon=lon, lat=lat) + ds = xr.Dataset({"field": 2.0 * lat + lon}).rename({"x": "column", "y": "row"}) + if transpose: + ds = ds.transpose("column", "row") + original = ds.copy(deep=True) + grid_data = grid._ds.copy(deep=True) + source_dims = grid._source_dims_dict.copy() + mapping = {"n_face": ("column", "row")} + + mapped = ux.UxDataset.from_xarray(ds, uxgrid=grid, ugrid_dims=mapping) + + assert mapped.uxgrid is grid + assert grid.source_grid_spec == "Structured" + assert grid._source_dims_dict == source_dims == {"n_face": ("x", "y")} + assert mapping == {"n_face": ("column", "row")} + xr.testing.assert_identical(grid._ds, grid_data) + xr.testing.assert_identical(ds, original) + # Recover each face's physical center independently of dimension stacking. + nodes = grid.face_node_connectivity.values + face_lon = grid.node_lon.values[nodes].mean(axis=1) + face_lat = grid.node_lat.values[nodes].mean(axis=1) + assert mapped.field.dims == ("n_face",) + np.testing.assert_allclose(mapped.field.values, 2.0 * face_lat + face_lon) + + @pytest.mark.parametrize("transpose", [False, True]) def test_structured_netcdf_face_selection(tmp_path, transpose): """Read a member file, attach the same Grid, and select the correct face.""" @@ -112,7 +143,7 @@ def test_structured_missing_source_dimension_raises(missing): ds = _structured_dataset() grid = ux.Grid.from_structured(lon=ds.x, lat=ds.y) ds = ds.isel({dim: 0 for dim in ("x", "y") if missing in (dim, "both")}) - with pytest.raises(DimensionError, match="source spatial dimensions"): + with pytest.raises(DimensionError, match="selected mapping's spatial dimensions"): ux.UxDataset.from_xarray(ds, uxgrid=grid) @@ -121,5 +152,5 @@ def test_structured_coordinate_name_mapping_mismatch_raises(): ds = _structured_dataset().rename_vars({"x": "longitude", "y": "latitude"}) # Grid.from_dataset currently records coordinate names rather than dimension # names for this input. Correcting that metadata is separate from GH1828. - with pytest.raises(DimensionError, match="source spatial dimensions"): + with pytest.raises(DimensionError, match="selected mapping's spatial dimensions"): ux.UxDataset.from_structured(ds) diff --git a/uxarray/core/utils.py b/uxarray/core/utils.py index 1b449325b..b3a78138c 100644 --- a/uxarray/core/utils.py +++ b/uxarray/core/utils.py @@ -64,7 +64,7 @@ def _map_dims_to_ugrid( lon_name, lat_name = _source_dims_dict["n_face"] if lon_name not in ds.dims or lat_name not in ds.dims: raise DimensionError( - "Structured data must contain the Grid's source spatial " + "Structured data must contain the selected mapping's spatial " f"dimensions {lon_name!r} and {lat_name!r}; got {tuple(ds.dims)!r}." ) # Structured Grid faces are latitude-major, longitude-minor. The source