From 2d5c796f9d57cf79d876bdc8bd9a407119cf0112 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, 16 Oct 2026 02:30:39 +0800 Subject: [PATCH 1/7] Support reusing an existing Grid when opening datasets --- benchmarks/grid_reuse.py | 49 +++++ test/core/test_grid_reuse.py | 355 +++++++++++++++++++++++++++++++++++ uxarray/core/api.py | 109 +++++++++-- 3 files changed, 499 insertions(+), 14 deletions(-) create mode 100644 benchmarks/grid_reuse.py create mode 100644 test/core/test_grid_reuse.py diff --git a/benchmarks/grid_reuse.py b/benchmarks/grid_reuse.py new file mode 100644 index 000000000..3351b5593 --- /dev/null +++ b/benchmarks/grid_reuse.py @@ -0,0 +1,49 @@ +"""Opening synthetic ensemble data with a rebuilt or reusable grid.""" + +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np +import xarray as xr + +import uxarray as ux + + +class OpenDatasetGridReuse: + param_names = ["grid_input", "data_input"] + params = [["file", "dataset", "grid"], ["file", "dataset"]] + + def setup(self, grid_input, data_input): + """Prepare synthetic inputs outside the timed member-opening path.""" + self.directory = TemporaryDirectory() + self.grid = ux.Grid.from_topology( + node_lon=np.array([0.0, 1.0, 1.0, 0.0]), + node_lat=np.array([0.0, 0.0, 1.0, 1.0]), + face_node_connectivity=np.array([[0, 1, 2], [0, 2, 3]]), + ) + self.grid_ds = self.grid.to_xarray() + self.data_ds = xr.Dataset( + {"temperature": (("time", "n_face"), np.arange(20).reshape(10, 2))} + ) + directory = Path(self.directory.name) + grid_path = directory / "grid.nc" + data_path = directory / "member.nc" + self.grid_ds.to_netcdf(grid_path) + self.data_ds.to_netcdf(data_path) + self.grid_input = { + "file": grid_path, + "dataset": self.grid_ds, + "grid": self.grid, + }[grid_input] + self.data_input = data_path if data_input == "file" else self.data_ds + + def teardown(self, grid_input, data_input): + """Remove the generated files after measurement.""" + self.directory.cleanup() + + def time_open_dataset(self, grid_input, data_input): + """Time opening one member, retaining the Grid when already available.""" + result = ux.open_dataset(self.grid_input, self.data_input) + result.close() + if grid_input == "file": + result.uxgrid._ds.close() diff --git a/test/core/test_grid_reuse.py b/test/core/test_grid_reuse.py new file mode 100644 index 000000000..22d0a8b9a --- /dev/null +++ b/test/core/test_grid_reuse.py @@ -0,0 +1,355 @@ +"""Regression tests for opening ensemble members with a shared Grid.""" + +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest import TestCase +from unittest.mock import patch + +import dask.array as da +import numpy as np +import xarray as xr + +import uxarray as ux + + +class TestGridReuse(TestCase): + def setUp(self): + """Create a synthetic tetrahedral grid and two independent members.""" + self.directory = TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.grid_path = Path(self.directory.name) / "grid.nc" + # Equal node and face counts ensure source names, rather than sizes, + # distinguish node-centered and face-centered variables after reuse. + self.grid_ds = xr.Dataset( + { + "mesh": xr.DataArray( + 0, + attrs={ + "cf_role": "mesh_topology", + "topology_dimension": 2, + "node_coordinates": "lon lat", + "face_node_connectivity": "face_nodes", + "node_dimension": "vertices", + "face_dimension": "cells", + }, + ), + "lon": ("vertices", [0.0, 120.0, -120.0, 0.0]), + "lat": ("vertices", [0.0, 0.0, 0.0, 90.0]), + "face_nodes": ( + ("cells", "corners"), + [[0, 1, 2], [0, 3, 1], [1, 3, 2], [2, 3, 0]], + {"start_index": 0}, + ), + } + ) + self.grid_ds.to_netcdf(self.grid_path) + self.grid = ux.open_grid(self.grid_path) + self.addCleanup(self.grid._ds.close) + self.source_dims = self.grid._source_dims_dict.copy() + self.grid_snapshot = self.grid._ds.copy(deep=True) + self.members = [ + xr.Dataset( + {"temperature": (("time", "cells"), np.arange(12).reshape(3, 4) + i)}, + coords={"time": [0, 1, 2]}, + ) + for i in range(2) + ] + self.paths = [] + for i, member in enumerate(self.members): + path = Path(self.directory.name) / f"member_{i}.nc" + member.to_netcdf(path) + self.paths.append(path) + + def assert_grid_unchanged(self): + """Opening data must not alter the shared grid or its source mapping.""" + self.assertEqual(self.grid._source_dims_dict, self.source_dims) + xr.testing.assert_identical(self.grid._ds, self.grid_snapshot) + + def test_open_dataset_reuses_grid_for_file_and_dataset_members(self): + """A single Grid is attached by identity to independent member data.""" + for sources in (self.paths, self.members): + for source, expected in zip(sources, self.members): + with self.subTest(source=type(source).__name__): + result = ux.open_dataset(self.grid, source) + self.addCleanup(result.close) + self.assertIs(result.uxgrid, self.grid) + self.assertEqual(result.temperature.dims, ("time", "n_face")) + np.testing.assert_array_equal( + result.temperature, expected.temperature + ) + self.assertEqual( + result.source_datasets, + None if isinstance(source, xr.Dataset) else str(source), + ) + self.assert_grid_unchanged() + + def test_members_with_different_grid_dimensions_preserve_mapping(self): + """Opening face-only data must not lose the next member's node mapping.""" + node_data = xr.Dataset({"velocity": ("vertices", np.arange(4))}) + for member, name, dimension in ( + (self.members[0], "temperature", "n_face"), + (node_data, "velocity", "n_node"), + (self.members[1], "temperature", "n_face"), + ): + result = ux.open_dataset(self.grid, member) + self.assertEqual(result[name].dims[-1], dimension) + self.assert_grid_unchanged() + + def test_chunks_preserve_grid_and_caller_options(self): + """Chunk data with source or UGRID names without rechunking the Grid.""" + for source in (self.paths[0], self.members[0]): + for chunk_grid in (True, False): + for chunks in ( + {"n_face": 2, "time": 1}, + {"cells": 2, "time": 1}, + {"cells": 4, "n_face": 2, "time": 1}, + ): + with self.subTest( + source=type(source).__name__, + chunk_grid=chunk_grid, + chunks=chunks, + ): + original_chunks = chunks.copy() + grid_kwargs = {} + result = ux.open_dataset( + self.grid, + source, + chunks=chunks, + chunk_grid=chunk_grid, + grid_kwargs=grid_kwargs, + ) + self.addCleanup(result.close) + self.assertIs(result.uxgrid, self.grid) + self.assertIsInstance(result.temperature.data, da.Array) + self.assertEqual(result.temperature.chunks, ((1, 1, 1), (2, 2))) + np.testing.assert_array_equal( + result.temperature.compute(), self.members[0].temperature + ) + self.assertEqual(chunks, original_chunks) + self.assertEqual(grid_kwargs, {}) + self.assert_grid_unchanged() + + def test_scalar_and_auto_chunks(self): + """All non-dict chunk forms reach the member data unchanged.""" + for source in (self.paths[0], self.members[0]): + for chunks in (-1, "auto", {}, None): + with self.subTest(source=type(source).__name__, chunks=chunks): + result = ux.open_dataset(self.grid, source, chunks=chunks) + self.addCleanup(result.close) + self.assertEqual( + isinstance(result.temperature.data, da.Array), + chunks is not None, + ) + self.assert_grid_unchanged() + + def test_prechunked_grid_is_not_rechunked(self): + """Each member retains the same already-chunked grid arrays.""" + self.grid.chunk(n_node=2, n_face=2) + connectivity = self.grid.face_node_connectivity.data + result = ux.open_dataset(self.grid, self.paths[0], chunks={"n_face": 1}) + self.addCleanup(result.close) + self.assertIs(result.uxgrid, self.grid) + self.assertIs(self.grid.face_node_connectivity.data, connectivity) + self.assertEqual(result.temperature.chunks[-1], (1, 1, 1, 1)) + self.assertEqual(self.grid.face_node_connectivity.chunks[0], (2, 2)) + + def test_existing_grid_does_not_open_or_reconstruct_grid(self): + """An in-memory Grid/Dataset pair must not invoke file backends.""" + with ( + patch("uxarray.core.api.open_grid") as open_grid, + patch("xarray.open_dataset") as open_dataset, + ): + result = ux.open_dataset(self.grid, self.members[0], chunks={"n_face": 2}) + self.assertIs(result.uxgrid, self.grid) + open_grid.assert_not_called() + open_dataset.assert_not_called() + + def test_open_mfdataset_reuses_grid(self): + """The multi-file API shares the same reuse and chunking contract.""" + chunks = {"n_face": 2, "time": 1} + with patch("uxarray.core.api.open_grid") as open_grid: + result = ux.open_mfdataset( + self.grid, + self.paths, + chunks=chunks, + combine="nested", + concat_dim="member", + ) + self.addCleanup(result.close) + open_grid.assert_not_called() + self.assertIs(result.uxgrid, self.grid) + self.assertEqual(result.temperature.dims, ("member", "time", "n_face")) + self.assertEqual(result.temperature.chunks, ((1, 1), (1, 1, 1), (2, 2))) + np.testing.assert_array_equal( + result.temperature.compute(), + np.stack([member.temperature.values for member in self.members]), + ) + self.assertEqual(chunks, {"n_face": 2, "time": 1}) + self.assert_grid_unchanged() + + def test_grid_construction_options_are_rejected(self): + """Options requiring grid construction should not be silently ignored.""" + for opener, data in ( + (ux.open_dataset, self.members[0]), + (ux.open_mfdataset, self.paths), + ): + for options, message in ( + ({"use_dual": True}, "use_dual"), + ({"grid_kwargs": {"drop_variables": "face_nodes"}}, "grid_kwargs"), + ): + with self.subTest(opener=opener.__name__, options=options): + with self.assertRaisesRegex(ValueError, message): + opener(self.grid, data, **options) + self.assert_grid_unchanged() + + def test_grid_requires_separate_data(self): + """A Grid alone does not provide member data.""" + with self.assertRaisesRegex(TypeError, "filename_or_obj"): + ux.open_dataset(self.grid) + + def test_canonical_dimensions_are_not_inferred_from_equal_sizes(self): + """Canonical face and node dimensions remain distinct across members.""" + members = [ + member.rename({"cells": "n_face"}).assign( + velocity=("n_node", np.arange(4) + i) + ) + for i, member in enumerate(self.members) + ] + paths = [] + for i, member in enumerate(members): + path = Path(self.directory.name) / f"canonical_{i}.nc" + member.to_netcdf(path) + paths.append(path) + for sources in (members, paths): + for source, expected in zip(sources, members): + result = ux.open_dataset( + self.grid, source, chunks={"n_face": 2, "n_node": 1} + ) + self.addCleanup(result.close) + self.assertIs(result.uxgrid, self.grid) + self.assertEqual(result.temperature.dims, ("time", "n_face")) + self.assertEqual(result.velocity.dims, ("n_node",)) + self.assertEqual(result.temperature.chunks[-1], (2, 2)) + self.assertEqual(result.velocity.chunks, ((1, 1, 1, 1),)) + np.testing.assert_array_equal(result.temperature, expected.temperature) + np.testing.assert_array_equal(result.velocity, expected.velocity) + self.assert_grid_unchanged() + combined = ux.open_mfdataset( + self.grid, + paths, + combine="nested", + concat_dim="member", + chunks={"n_face": 2, "n_node": 1}, + ) + self.addCleanup(combined.close) + self.assertEqual(combined.temperature.dims, ("member", "time", "n_face")) + self.assertEqual(combined.velocity.dims, ("member", "n_node")) + np.testing.assert_array_equal( + combined.temperature, + np.stack([member.temperature.values for member in members]), + ) + self.assert_grid_unchanged() + # Canonical members must not disrupt later source-dimension members. + native = ux.open_dataset(self.grid, self.members[0]) + self.assertEqual(native.temperature.dims, ("time", "n_face")) + self.assert_grid_unchanged() + + def test_topology_grid_without_source_mapping(self): + """A topology-created Grid can reuse already-standardized data.""" + grid = ux.Grid.from_topology( + node_lon=self.grid.node_lon.values, + node_lat=self.grid.node_lat.values, + face_node_connectivity=self.grid.face_node_connectivity.values, + ) + member = xr.Dataset({"temperature": ("n_face", np.arange(4))}) + result = ux.open_dataset(grid, member, chunks={"n_face": 2}) + self.assertIs(result.uxgrid, grid) + self.assertEqual(result.temperature.dims, ("n_face",)) + self.assertEqual(result.temperature.chunks, ((2, 2),)) + np.testing.assert_array_equal(result.temperature, member.temperature) + self.assertEqual(grid._source_dims_dict, {}) + + def test_explicit_source_mapping_takes_precedence(self): + """An explicit source mapping wins even for a canonical-looking name.""" + self.grid._source_dims_dict = {"n_face": "n_node"} + result = ux.open_dataset( + self.grid, xr.Dataset({"velocity": ("n_face", np.arange(4))}) + ) + self.assertEqual(result.velocity.dims, ("n_node",)) + self.assertEqual(self.grid._source_dims_dict, {"n_face": "n_node"}) + + def test_structured_grid_chunks_use_source_dimensions(self): + """Structured grids retain their tuple-valued source mapping on reuse.""" + grid = ux.Grid.from_structured( + lon=xr.DataArray([0.0, 20.0, 40.0, 60.0], dims="x"), + lat=xr.DataArray([-30.0, 0.0, 30.0], dims="y"), + ) + source_dims = grid._source_dims_dict.copy() + member = xr.Dataset({"temperature": (("y", "x"), np.arange(12).reshape(3, 4))}) + path = Path(self.directory.name) / "structured.nc" + member.to_netcdf(path) + for source in (member, path): + chunks = {"x": 2, "y": 1} + result = ux.open_dataset(grid, source, chunks=chunks) + self.addCleanup(result.close) + self.assertIs(result.uxgrid, grid) + self.assertEqual(result.temperature.dims, ("n_face",)) + self.assertEqual(result.temperature.chunks, ((2, 2, 2, 2, 2, 2),)) + np.testing.assert_array_equal(result.temperature, np.arange(12)) + self.assertEqual(grid._source_dims_dict, source_dims) + self.assertEqual(chunks, {"x": 2, "y": 1}) + + def test_structured_grid_accepts_already_flattened_members(self): + """Do not flatten canonical structured member data a second time.""" + grid = ux.Grid.from_structured( + lon=xr.DataArray([0.0, 20.0, 40.0, 60.0], dims="x"), + lat=xr.DataArray([-30.0, 0.0, 30.0], dims="y"), + ) + source_dims = grid._source_dims_dict.copy() + members = [ + xr.Dataset({"temperature": ("n_face", np.arange(12) + i)}) for i in range(2) + ] + paths = [] + for i, member in enumerate(members): + path = Path(self.directory.name) / f"flattened_{i}.nc" + member.to_netcdf(path) + paths.append(path) + for sources in (members, paths): + for source, expected in zip(sources, members): + result = ux.open_dataset(grid, source, chunks={"n_face": 2}) + self.addCleanup(result.close) + self.assertIs(result.uxgrid, grid) + self.assertEqual(result.temperature.dims, ("n_face",)) + self.assertEqual(result.temperature.chunks, ((2, 2, 2, 2, 2, 2),)) + np.testing.assert_array_equal(result.temperature, expected.temperature) + self.assertEqual(grid._source_dims_dict, source_dims) + combined = ux.open_mfdataset( + grid, + paths, + combine="nested", + concat_dim="member", + chunks={"n_face": 2}, + ) + self.addCleanup(combined.close) + self.assertIs(combined.uxgrid, grid) + self.assertEqual(combined.temperature.dims, ("member", "n_face")) + self.assertEqual(combined.temperature.chunks, ((1, 1), (2, 2, 2, 2, 2, 2))) + np.testing.assert_array_equal( + combined.temperature, + np.stack([member.temperature.values for member in members]), + ) + self.assertEqual(grid._source_dims_dict, source_dims) + + def test_existing_grid_file_and_dataset_inputs(self): + """Previously supported grid inputs still construct equivalent grids.""" + for grid_source in (self.grid_path, self.grid_ds): + result = ux.open_dataset(grid_source, self.paths[0], chunks={"n_face": 2}) + self.addCleanup(result.close) + self.addCleanup(result.uxgrid._ds.close) + self.assertIsNot(result.uxgrid, self.grid) + self.assertEqual(result.temperature.dims, ("time", "n_face")) + self.assertEqual(result.temperature.chunks[-1], (2, 2)) + np.testing.assert_array_equal( + result.temperature, self.members[0].temperature + ) + xr.testing.assert_identical(result.uxgrid._ds.compute(), self.grid_snapshot) diff --git a/uxarray/core/api.py b/uxarray/core/api.py index 5c32b8b50..7ee4f0e33 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -367,7 +367,7 @@ def list_grid_names( def open_dataset( - grid_filename_or_obj: str | os.PathLike[Any] | dict | xr.Dataset, + grid_filename_or_obj: str | os.PathLike[Any] | dict | xr.Dataset | Grid, filename_or_obj: str | os.PathLike[Any] | xr.Dataset | None = None, chunks=None, chunk_grid: bool = True, @@ -379,13 +379,15 @@ def open_dataset( Parameters ---------- - grid_filename_or_obj : str | os.PathLike[Any] | dict | xr.Dataset + grid_filename_or_obj : str | os.PathLike[Any] | dict | xr.Dataset | Grid Grid information for the ``UxDataset``. Strings and Path objects are interpreted as a path to a grid file. Xarray Datasets assume that each member variable is in the UGRID conventions and will be used to create a ``ux.Grid``. Similarly, a dictionary containing UGRID variables can be used to create a ``ux.Grid``. A path to a directory containing grid files (e.g. a FESOM2 ASCII grid) is also accepted, but only when ``filename_or_obj`` is provided; directory input is not supported in the single-argument - form below. + form below. An existing ``ux.Grid`` is attached directly, allowing its + topology and cached properties to be shared across datasets. A separate + ``filename_or_obj`` is required when passing a ``Grid``. filename_or_obj : str | os.PathLike[Any] | xr.Dataset, optional String or Path object as a path to a netCDF file or an OpenDAP URL that stores the actual data set, or an already-open ``xarray.Dataset``. It @@ -396,7 +398,9 @@ def open_dataset( the grid and also treated as data; all variables it contains, including grid variables, will be included as data. chunks : int, dict, 'auto' or None, default: None - If provided, used to load the grid into dask arrays. + If provided, used to load the data into dask arrays. With an existing + ``Grid``, this also applies to an already-open ``xarray.Dataset`` + without modifying the input dataset. - ``chunks="auto"`` will use dask ``auto`` chunking taking into account the engine preferred chunks. @@ -406,12 +410,22 @@ def open_dataset( - ``chunks={}`` loads the data with dask using the engine's preferred chunk size, generally identical to the format's chunk size. If not available, a single chunk for all arrays. + + For structured data that still has its original spatial dimensions, + use those names (such as ``lon`` and ``lat``) to chunk before + flattening. You can also rechunk the returned ``n_face`` dimension. chunk_grid : bool, default: True - If valid chunks are passed in, determines whether to also apply the same chunks to the attached ``Grid`` + If valid chunks are passed in, determines whether to also apply the same + chunks when constructing a ``Grid``. An existing ``Grid`` is reused + unchanged, including its chunking; call ``Grid.chunk`` explicitly + before opening data to change its chunking. use_dual : bool, optional - Specify whether to use the primal (use_dual=False) or dual (use_dual=True) mesh if the file type is MPAS + Specify whether to use the primal (use_dual=False) or dual (use_dual=True) mesh if the file type is MPAS. + ``use_dual=True`` is not supported for an existing ``Grid``; select + the desired mesh when constructing the grid. grid_kwargs : dict, optional Additional arguments passed on to ``ux.open_grid`` when opening up a Grid File. + Must be empty or ``None`` when passing an existing ``Grid``. **kwargs Additional arguments passed on to ``xarray.open_dataset``. Refer to the :func:`xarray docs ` @@ -440,6 +454,14 @@ def open_dataset( Open a dataset stored in a single combined grid-and-data file >>> ux_ds = ux.open_dataset("combined_file.nc") + + Reuse one grid for multiple ensemble members + + >>> uxgrid = ux.open_grid("grid_file.nc") + >>> member_1 = ux.open_dataset(uxgrid, "member_1.nc") + >>> member_2 = ux.open_dataset(uxgrid, "member_2.nc") + >>> member_1.uxgrid is member_2.uxgrid + True """ if grid_kwargs is None: grid_kwargs = {} @@ -480,13 +502,20 @@ def open_dataset( # Load the data as a Xarray Dataset if isinstance(filename_or_obj, xr.Dataset): ds = filename_or_obj + if isinstance(grid_filename_or_obj, Grid) and corrected_chunks is not None: + data_chunks = corrected_chunks + if isinstance(data_chunks, dict): + data_chunks = { + dim: size for dim, size in data_chunks.items() if dim in ds.dims + } + ds = ds.chunk(data_chunks) else: ds = _open_dataset_with_fallback( filename_or_obj, chunks=corrected_chunks, **kwargs ) # Map original dimensions to the UGRID conventions - ds = _map_dims_to_ugrid(ds, uxgrid._source_dims_dict, uxgrid) + ds = _map_dataset_to_ugrid(ds, uxgrid) # Create a UXarray Dataset by linking the Xarray Dataset with a UXarray Grid source_datasets = ( @@ -497,7 +526,7 @@ def open_dataset( def open_mfdataset( - grid_filename_or_obj: str | os.PathLike[Any] | dict | xr.Dataset, + grid_filename_or_obj: str | os.PathLike[Any] | dict | xr.Dataset | Grid, paths: str | os.PathLike, chunks=None, chunk_grid: bool = True, @@ -510,15 +539,16 @@ def open_mfdataset( Parameters ---------- - grid_filename_or_obj : str | os.PathLike[Any] | dict | xr.dataset + grid_filename_or_obj : str | os.PathLike[Any] | dict | xr.Dataset | Grid Strings and Path objects are interpreted as a path to a grid file. Xarray Datasets assume that each member variable is in the UGRID conventions and will be used to create a ``ux.Grid``. Similarly, a dictionary - containing UGRID variables can be used to create a ``ux.Grid`` + containing UGRID variables can be used to create a ``ux.Grid``. An existing + ``ux.Grid`` is attached directly without reconstructing or modifying it. paths : string, required Either a string glob in the form ``"path/to/my/files/*.nc"`` or an explicit list of files to open. It is the same ``paths`` in ``xarray.open_mfdataset``. chunks : int, dict, 'auto' or None, default: None - If provided, used to load the grid into dask arrays. + If provided, used to load the data into dask arrays. - ``chunks="auto"`` will use dask ``auto`` chunking taking into account the engine preferred chunks. @@ -528,12 +558,22 @@ def open_mfdataset( - ``chunks={}`` loads the data with dask using the engine's preferred chunk size, generally identical to the format's chunk size. If not available, a single chunk for all arrays. + + For structured data that still has its original spatial dimensions, + use those names (such as ``lon`` and ``lat``) to chunk before + flattening. You can also rechunk the returned ``n_face`` dimension. chunk_grid : bool, default: True - If valid chunks are passed in, determines whether to also apply the same chunks to the attached ``Grid`` + If valid chunks are passed in, determines whether to also apply the same + chunks when constructing a ``Grid``. An existing ``Grid`` is reused + unchanged, including its chunking; call ``Grid.chunk`` explicitly + before opening data to change its chunking. use_dual : bool, optional - Specify whether to use the primal (use_dual=False) or dual (use_dual=True) mesh if the file type is mpas + Specify whether to use the primal (use_dual=False) or dual (use_dual=True) mesh if the file type is MPAS. + ``use_dual=True`` is not supported for an existing ``Grid``; select + the desired mesh when constructing the grid. grid_kwargs : dict, optional Additional arguments passed on to ``ux.open_grid`` when opening up a Grid File. + Must be empty or ``None`` when passing an existing ``Grid``. **kwargs Additional arguments passed on to ``xarray.open_mfdataset``. Refer to the :func:`xarray docs ` @@ -561,6 +601,11 @@ def open_mfdataset( 2. Open from a string glob >>> ux_ds = ux.open_mfdataset("grid_filename.g", "grid_filename_vortex_*.nc") + + 3. Reuse an existing grid + + >>> uxgrid = ux.open_grid("grid_filename.g") + >>> ux_ds = ux.open_mfdataset(uxgrid, "grid_filename_vortex_*.nc") """ if grid_kwargs is None: grid_kwargs = {} @@ -574,16 +619,52 @@ def open_mfdataset( ds = xr.open_mfdataset(paths, chunks=corrected_chunks, **kwargs) # Map original dimensions to the UGRID conventions - ds = _map_dims_to_ugrid(ds, uxgrid._source_dims_dict, uxgrid) + ds = _map_dataset_to_ugrid(ds, uxgrid) # Create a UXarray Dataset by linking the Xarray Dataset with a UXarray Grid return UxDataset(ds, uxgrid=uxgrid, source_datasets=str(paths)) +def _map_dataset_to_ugrid(ds, grid): + """Map member dimensions without modifying the Grid's source mapping.""" + source_dims = grid._source_dims_dict.copy() + if grid.source_grid_spec == "Structured": + lon_dim, lat_dim = source_dims["n_face"] + if "n_face" in ds.dims and lon_dim not in ds.dims and lat_dim not in ds.dims: + # This member is already flattened onto the structured Grid. + return ds + # Already-standardized dimensions must not be inferred by size: nodes and + # faces can have equal counts. Preserve any explicit source mapping. + for dim in ("n_face", "n_node", "n_edge"): + if dim in ds.dims: + source_dims.setdefault(dim, dim) + return _map_dims_to_ugrid(ds, source_dims, grid) + + def _get_grid( grid_filename_or_obj, chunks, chunk_grid, use_dual, grid_kwargs, **kwargs ): """Utility function to validate the input parameters and return a Grid.""" + if isinstance(grid_filename_or_obj, Grid): + if use_dual: + raise ValueError( + "use_dual=True is not supported with an existing Grid. " + "Select the desired mesh when constructing the Grid." + ) + if grid_kwargs: + raise ValueError( + "grid_kwargs must be empty when passing an existing Grid. " + "Apply these options when constructing the Grid." + ) + if isinstance(chunks, dict): + # Keep the caller's options and the shared Grid unchanged. Match + # UGRID chunk keys to the data's source dimensions without grid I/O. + chunks = chunks.copy() + for source_dim, ugrid_dim in grid_filename_or_obj._source_dims_dict.items(): + if isinstance(ugrid_dim, str) and ugrid_dim in chunks: + chunks[source_dim] = chunks[ugrid_dim] + return grid_filename_or_obj, chunks + if "latlon" in kwargs: warn( "'latlon is no longer a supported parameter", From b5821fff2104f72a1d4d62b7f5a12af902f1e16d 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, 16 Oct 2026 07:17:43 +0800 Subject: [PATCH 2/7] Close source dataset owners in Grid reuse test fixtures --- test/core/test_grid_reuse.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/test/core/test_grid_reuse.py b/test/core/test_grid_reuse.py index 22d0a8b9a..079f3abda 100644 --- a/test/core/test_grid_reuse.py +++ b/test/core/test_grid_reuse.py @@ -1,5 +1,6 @@ """Regression tests for opening ensemble members with a shared Grid.""" +from functools import partial from pathlib import Path from tempfile import TemporaryDirectory from unittest import TestCase @@ -17,6 +18,16 @@ def setUp(self): """Create a synthetic tetrahedral grid and two independent members.""" self.directory = TemporaryDirectory() self.addCleanup(self.directory.cleanup) + # Grid conversion and UxDataset construction can lose the source + # Dataset's close callback. Keep the actual backend owners alive and + # close them before removing files, including on Windows. + for name in ("open_dataset", "open_mfdataset"): + opener = getattr(xr, name) + patcher = patch.object( + xr, name, side_effect=partial(self._open_source, opener) + ) + patcher.start() + self.addCleanup(patcher.stop) self.grid_path = Path(self.directory.name) / "grid.nc" # Equal node and face counts ensure source names, rather than sizes, # distinguish node-centered and face-centered variables after reuse. @@ -60,6 +71,12 @@ def setUp(self): member.to_netcdf(path) self.paths.append(path) + def _open_source(self, opener, *args, **kwargs): + """Retain the original xarray resource owner without changing file I/O.""" + source = opener(*args, **kwargs) + self.addCleanup(source.close) + return source + def assert_grid_unchanged(self): """Opening data must not alter the shared grid or its source mapping.""" self.assertEqual(self.grid._source_dims_dict, self.source_dims) From 9b4fc457d38533e69da7109013c8da69bd92aa2b 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: Wed, 21 Oct 2026 21:01:38 +0800 Subject: [PATCH 3/7] Fix canonical structured members and overlapping chunk aliases --- test/core/test_grid_reuse.py | 162 +++++++++++++++++++++++++++++++++++ uxarray/core/api.py | 10 ++- 2 files changed, 168 insertions(+), 4 deletions(-) diff --git a/test/core/test_grid_reuse.py b/test/core/test_grid_reuse.py index 079f3abda..61189da6b 100644 --- a/test/core/test_grid_reuse.py +++ b/test/core/test_grid_reuse.py @@ -357,6 +357,168 @@ def test_structured_grid_accepts_already_flattened_members(self): ) self.assertEqual(grid._source_dims_dict, source_dims) + def test_structured_grid_accepts_canonical_node_and_edge_members(self): + """Canonical node and edge data have no source axes left to flatten.""" + grid = ux.Grid.from_structured( + lon=xr.DataArray([0.0, 20.0, 40.0, 60.0], dims="x"), + lat=xr.DataArray([-30.0, 0.0, 30.0], dims="y"), + ) + # n_edge lazily creates connectivity; snapshot the intended cache first. + counts = {"n_node": grid.n_node, "n_edge": grid.n_edge} + source_dims = grid._source_dims_dict.copy() + snapshot = grid._ds.copy(deep=True) + for dimension, size in counts.items(): + members = [ + xr.Dataset({"value": (dimension, np.arange(size) + i)}) + for i in range(2) + ] + member_snapshots = [member.copy(deep=True) for member in members] + paths = [] + for i, member in enumerate(members): + path = Path(self.directory.name) / f"structured_{dimension}_{i}.nc" + member.to_netcdf(path) + paths.append(path) + cases = ( + [("dataset", member, member.value.values) for member in members] + + [ + ("file", path, member.value.values) + for path, member in zip(paths, members) + ] + + [ + ( + "multiple", + paths, + np.stack([member.value.values for member in members]), + ) + ] + ) + chunks = {dimension: 2} + expected_chunks = (2,) * (size // 2) + ((1,) if size % 2 else ()) + for case_index, (kind, source, expected) in enumerate(cases): + with self.subTest( + dimension=dimension, kind=kind, case_index=case_index + ): + if kind == "multiple": + result = ux.open_mfdataset( + grid, + source, + combine="nested", + concat_dim="member", + chunks=chunks, + ) + expected_dims = ("member", dimension) + else: + result = ux.open_dataset(grid, source, chunks=chunks) + expected_dims = (dimension,) + self.addCleanup(result.close) + self.assertIs(result.uxgrid, grid) + self.assertEqual(result.value.dims, expected_dims) + self.assertIsInstance(result.value.data, da.Array) + all_chunks = (expected_chunks,) + if kind == "multiple": + all_chunks = ((1, 1),) + all_chunks + self.assertEqual(result.value.chunks, all_chunks) + np.testing.assert_array_equal(result.value.compute(), expected) + self.assertEqual(chunks, {dimension: 2}) + self.assertEqual(grid._source_dims_dict, source_dims) + xr.testing.assert_identical(grid._ds, snapshot) + for original, saved in zip(members, member_snapshots): + xr.testing.assert_identical(original, saved) + # Reuse after canonical data must still flatten ordinary face members. + native = xr.Dataset({"temperature": (("y", "x"), np.arange(12).reshape(3, 4))}) + result = ux.open_dataset(grid, native) + self.assertIs(result.uxgrid, grid) + self.assertEqual(result.temperature.dims, ("n_face",)) + np.testing.assert_array_equal(result.temperature, np.arange(12)) + self.assertEqual(grid._source_dims_dict, source_dims) + xr.testing.assert_identical(grid._ds, snapshot) + + def test_overlapping_chunk_aliases_read_only_requested_values(self): + """Alias translation is independent of order and newly created keys.""" + members = [ + member.assign(velocity=("n_face", np.arange(4) + i)) + for i, member in enumerate(self.members) + ] + member_snapshots = [member.copy(deep=True) for member in members] + paths = [] + for i, member in enumerate(members): + path = Path(self.directory.name) / f"overlapping_alias_{i}.nc" + member.to_netcdf(path) + paths.append(path) + for source_dims in ( + {"n_face": "n_node", "cells": "n_face"}, + {"cells": "n_face", "n_face": "n_node"}, + ): + grid = ux.Grid.from_topology( + node_lon=self.grid.node_lon.values, + node_lat=self.grid.node_lat.values, + face_node_connectivity=self.grid.face_node_connectivity.values, + dims_dict=source_dims.copy(), + ) + snapshot = grid._ds.copy(deep=True) + for requested in ( + {"n_face": 2, "n_node": 1, "time": 1}, + {"n_node": 1, "time": 1}, + ): + original_chunks = requested.copy() + expected_face_chunks = (2, 2) if "n_face" in requested else (4,) + cases = ( + [("dataset", member, member) for member in members] + + [("file", path, member) for path, member in zip(paths, members)] + + [("multiple", paths, None)] + ) + for case_index, (kind, source, member) in enumerate(cases): + with self.subTest( + source_dims=source_dims, + requested=requested, + kind=kind, + case_index=case_index, + ): + if kind == "multiple": + result = ux.open_mfdataset( + grid, + source, + combine="nested", + concat_dim="member", + chunks=requested, + ) + temperature = np.stack( + [m.temperature.values for m in members] + ) + velocity = np.stack([m.velocity.values for m in members]) + temperature_dims = ("member", "time", "n_face") + velocity_dims = ("member", "n_node") + else: + result = ux.open_dataset(grid, source, chunks=requested) + temperature = member.temperature.values + velocity = member.velocity.values + temperature_dims = ("time", "n_face") + velocity_dims = ("n_node",) + self.addCleanup(result.close) + self.assertIs(result.uxgrid, grid) + self.assertEqual(result.temperature.dims, temperature_dims) + self.assertEqual(result.velocity.dims, velocity_dims) + self.assertIsInstance(result.temperature.data, da.Array) + self.assertIsInstance(result.velocity.data, da.Array) + temperature_chunks = ((1, 1, 1), expected_face_chunks) + velocity_chunks = ((1, 1, 1, 1),) + if kind == "multiple": + temperature_chunks = ((1, 1),) + temperature_chunks + velocity_chunks = ((1, 1),) + velocity_chunks + self.assertEqual(result.temperature.chunks, temperature_chunks) + self.assertEqual(result.velocity.chunks, velocity_chunks) + np.testing.assert_array_equal( + result.temperature.compute(), temperature + ) + np.testing.assert_array_equal( + result.velocity.compute(), velocity + ) + self.assertEqual(requested, original_chunks) + self.assertEqual(grid._source_dims_dict, source_dims) + xr.testing.assert_identical(grid._ds, snapshot) + for original, saved in zip(members, member_snapshots): + xr.testing.assert_identical(original, saved) + def test_existing_grid_file_and_dataset_inputs(self): """Previously supported grid inputs still construct equivalent grids.""" for grid_source in (self.grid_path, self.grid_ds): diff --git a/uxarray/core/api.py b/uxarray/core/api.py index 7ee4f0e33..fa5695cc5 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -630,8 +630,9 @@ def _map_dataset_to_ugrid(ds, grid): source_dims = grid._source_dims_dict.copy() if grid.source_grid_spec == "Structured": lon_dim, lat_dim = source_dims["n_face"] - if "n_face" in ds.dims and lon_dim not in ds.dims and lat_dim not in ds.dims: - # This member is already flattened onto the structured Grid. + if lon_dim not in ds.dims and lat_dim not in ds.dims: + # No source spatial axes remain to flatten, including members + # located only on canonical nodes or edges. return ds # Already-standardized dimensions must not be inferred by size: nodes and # faces can have equal counts. Preserve any explicit source mapping. @@ -659,10 +660,11 @@ def _get_grid( if isinstance(chunks, dict): # Keep the caller's options and the shared Grid unchanged. Match # UGRID chunk keys to the data's source dimensions without grid I/O. + requested_chunks = chunks chunks = chunks.copy() for source_dim, ugrid_dim in grid_filename_or_obj._source_dims_dict.items(): - if isinstance(ugrid_dim, str) and ugrid_dim in chunks: - chunks[source_dim] = chunks[ugrid_dim] + if isinstance(ugrid_dim, str) and ugrid_dim in requested_chunks: + chunks[source_dim] = requested_chunks[ugrid_dim] return grid_filename_or_obj, chunks if "latlon" in kwargs: From 7acbe336fd673205a7002f0e778472f84d530ece 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 23:46:33 +0800 Subject: [PATCH 4/7] Address Grid reuse review: helper placement, contracts and benchmark scope --- benchmarks/{grid_reuse.py => open_dataset.py} | 4 ++-- test/core/test_grid_reuse.py | 1 + uxarray/core/api.py | 21 ++----------------- uxarray/core/utils.py | 17 +++++++++++++++ 4 files changed, 22 insertions(+), 21 deletions(-) rename benchmarks/{grid_reuse.py => open_dataset.py} (94%) diff --git a/benchmarks/grid_reuse.py b/benchmarks/open_dataset.py similarity index 94% rename from benchmarks/grid_reuse.py rename to benchmarks/open_dataset.py index 3351b5593..873cbee21 100644 --- a/benchmarks/grid_reuse.py +++ b/benchmarks/open_dataset.py @@ -1,4 +1,4 @@ -"""Opening synthetic ensemble data with a rebuilt or reusable grid.""" +"""Opening synthetic data with file, Dataset, or existing Grid inputs.""" from pathlib import Path from tempfile import TemporaryDirectory @@ -9,7 +9,7 @@ import uxarray as ux -class OpenDatasetGridReuse: +class OpenDataset: param_names = ["grid_input", "data_input"] params = [["file", "dataset", "grid"], ["file", "dataset"]] diff --git a/test/core/test_grid_reuse.py b/test/core/test_grid_reuse.py index 61189da6b..0cc9af4eb 100644 --- a/test/core/test_grid_reuse.py +++ b/test/core/test_grid_reuse.py @@ -526,6 +526,7 @@ def test_existing_grid_file_and_dataset_inputs(self): self.addCleanup(result.close) self.addCleanup(result.uxgrid._ds.close) self.assertIsNot(result.uxgrid, self.grid) + self.assertEqual(result.uxgrid, self.grid) self.assertEqual(result.temperature.dims, ("time", "n_face")) self.assertEqual(result.temperature.chunks[-1], (2, 2)) np.testing.assert_array_equal( diff --git a/uxarray/core/api.py b/uxarray/core/api.py index fa5695cc5..6e35bcea3 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -12,7 +12,7 @@ from uxarray.core.dataarray import UxDataArray from uxarray.core.dataset import UxDataset from uxarray.core.utils import ( - _map_dims_to_ugrid, + _map_dataset_to_ugrid, _open_dataset_with_fallback, match_chunks_to_ugrid, ) @@ -625,27 +625,10 @@ def open_mfdataset( return UxDataset(ds, uxgrid=uxgrid, source_datasets=str(paths)) -def _map_dataset_to_ugrid(ds, grid): - """Map member dimensions without modifying the Grid's source mapping.""" - source_dims = grid._source_dims_dict.copy() - if grid.source_grid_spec == "Structured": - lon_dim, lat_dim = source_dims["n_face"] - if lon_dim not in ds.dims and lat_dim not in ds.dims: - # No source spatial axes remain to flatten, including members - # located only on canonical nodes or edges. - return ds - # Already-standardized dimensions must not be inferred by size: nodes and - # faces can have equal counts. Preserve any explicit source mapping. - for dim in ("n_face", "n_node", "n_edge"): - if dim in ds.dims: - source_dims.setdefault(dim, dim) - return _map_dims_to_ugrid(ds, source_dims, grid) - - def _get_grid( grid_filename_or_obj, chunks, chunk_grid, use_dual, grid_kwargs, **kwargs ): - """Utility function to validate the input parameters and return a Grid.""" + """Validate input parameters and return the ``(grid, chunks)`` tuple.""" if isinstance(grid_filename_or_obj, Grid): if use_dual: raise ValueError( diff --git a/uxarray/core/utils.py b/uxarray/core/utils.py index 84c0a2dc6..5e2120424 100644 --- a/uxarray/core/utils.py +++ b/uxarray/core/utils.py @@ -51,6 +51,23 @@ def _open_dataset_with_fallback(filename_or_obj, chunks=None, **kwargs): raise fallback_error from default_engine_error +def _map_dataset_to_ugrid(ds, grid): + """Map member dimensions without modifying the Grid's source mapping.""" + source_dims = grid._source_dims_dict.copy() + if grid.source_grid_spec == "Structured": + lon_dim, lat_dim = source_dims["n_face"] + if lon_dim not in ds.dims and lat_dim not in ds.dims: + # No source spatial axes remain to flatten, including members + # located only on canonical nodes or edges. + return ds + # Already-standardized dimensions must not be inferred by size: nodes and + # faces can have equal counts. Preserve any explicit source mapping. + for dim in ("n_face", "n_node", "n_edge"): + if dim in ds.dims: + source_dims.setdefault(dim, dim) + return _map_dims_to_ugrid(ds, source_dims, grid) + + def _map_dims_to_ugrid( ds, _source_dims_dict, From 62f4d411ac133a36cc5461ec9caa984f2e5a068f Mon Sep 17 00:00:00 2001 From: Lubai <1174273035@qq.com> Date: Wed, 28 Oct 2026 02:42:39 +0800 Subject: [PATCH 5/7] Consolidate unstructured Grid reuse tests and clarify docs --- test/core/test_api.py | 301 ++++++++++++++++++++ test/core/test_grid_reuse.py | 535 ----------------------------------- uxarray/core/api.py | 30 +- 3 files changed, 314 insertions(+), 552 deletions(-) delete mode 100644 test/core/test_grid_reuse.py diff --git a/test/core/test_api.py b/test/core/test_api.py index df63595c0..376400641 100644 --- a/test/core/test_api.py +++ b/test/core/test_api.py @@ -1,3 +1,7 @@ +from contextlib import ExitStack +from functools import partial + +import dask.array as da import numpy.testing as nt import uxarray as ux import numpy as np @@ -89,6 +93,303 @@ def test_open_dataset_single_argument_rejects_invalid_combined_file(datasetpath) ux.open_dataset(data_path) +@pytest.fixture +def reuse_owners(tmp_path, monkeypatch): + """Close original backend owners even if Grid/UxDataset loses callbacks.""" + with ExitStack() as owners: + + def tracked_open(opener, *args, **kwargs): + """Track the actual backend owner, not just its UxDataset wrapper.""" + source = opener(*args, **kwargs) + owners.callback(source.close) + return source + + for name in ("open_dataset", "open_mfdataset"): + monkeypatch.setattr(xr, name, partial(tracked_open, getattr(xr, name))) + yield owners + + +@pytest.fixture +def reuse_grid(reuse_owners): + """Use the suite's small unstructured quad-hexagon grid for ordinary reuse.""" + grid = ux.open_grid(TEST_MESHFILES / "ugrid" / "quad-hexagon" / "grid.nc") + grid.get_ball_tree(coordinates="nodes") + return grid + + +@pytest.fixture +def reuse_open(tmp_path, reuse_owners): + """Exercise Dataset/file/multifile inputs with identical preservation checks.""" + counter = 0 + + def open_members(grid, members, kind, **kwargs): + """Open members and verify preservation before returning expected data.""" + nonlocal counter + snapshots = [member.copy(deep=True) for member in members] + arrays = [{name: member[name].data for name in member} for member in members] + grid_snapshot, mapping = grid._ds.copy(deep=True), grid._source_dims_dict.copy() + grid_dataset = grid._ds + # Preserve storage, not materialized .data arrays from lazy backends. + grid_arrays = {name: grid._ds[name].variable._data for name in grid._ds} + tree = grid._ball_tree + chunks = kwargs.get("chunks") + saved_chunks = chunks.copy() if isinstance(chunks, dict) else chunks + grid_kwargs = kwargs.get("grid_kwargs") + saved_grid_kwargs = grid_kwargs.copy() if grid_kwargs is not None else None + paths = [] + if kind != "dataset": + for member in members: + path = tmp_path / f"member_{counter}.nc" + counter += 1 + member.to_netcdf(path) + paths.append(path) + # Chunk ownership for True is a separate open review question. + if chunks is not None: + kwargs.setdefault("chunk_grid", False) + with patch("uxarray.core.api.open_grid") as construct: + try: + if kind == "multiple": + results = [ + ux.open_mfdataset( + grid, paths, combine="nested", concat_dim="member", **kwargs + ) + ] + expected = [xr.concat(members, dim="member")] + sources = [str(paths)] + elif kind == "dataset": + with patch("xarray.open_dataset") as backend: + results = [ + ux.open_dataset(grid, member, **kwargs) + for member in members + ] + backend.assert_not_called() + expected, sources = members, [None] * len(members) + else: + results = [ux.open_dataset(grid, path, **kwargs) for path in paths] + expected, sources = members, list(map(str, paths)) + for result, source in zip(results, sources): + reuse_owners.callback(result.close) + assert result.uxgrid is grid + assert result.source_datasets == source + return list(zip(results, expected)) + finally: + construct.assert_not_called() + assert grid._ds is grid_dataset + assert grid._source_dims_dict == mapping + assert grid._ball_tree is tree + xr.testing.assert_identical(grid._ds, grid_snapshot) + for name, data in grid_arrays.items(): + assert grid._ds[name].variable._data is data + assert chunks == saved_chunks + assert grid_kwargs == saved_grid_kwargs + for member, snapshot, data in zip(members, snapshots, arrays): + xr.testing.assert_identical(member, snapshot) + for name in member: + assert member[name].chunks == snapshot[name].chunks + assert member[name].data is data[name] + + return open_members + + +def _reuse_members(face_dim="n_face"): + """Generate two deterministic members, independently of grid fixture data.""" + return [ + xr.Dataset( + {"temperature": (("time", face_dim), np.arange(12).reshape(3, 4) + i)} + ) + for i in range(2) + ] + + +def _equal_size_grid(source_dims): + """Equal node/face counts make size inference and alias collisions observable.""" + return ux.Grid.from_topology( + node_lon=np.array([0.0, 120.0, -120.0, 0.0]), + node_lat=np.array([0.0, 0.0, 0.0, 90.0]), + face_node_connectivity=np.array([[0, 1, 2], [0, 3, 1], [1, 3, 2], [2, 3, 0]]), + dims_dict=source_dims.copy(), + ) + + +@pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) +@pytest.mark.parametrize("chunks", [None, -1, "auto", {}, {"n_face": 2, "time": 1}]) +def test_open_dataset_reuse_members(reuse_grid, reuse_open, kind, chunks): + """Reuse topology/cached tree without grid I/O, preserving member values.""" + for result, expected in reuse_open( + reuse_grid, _reuse_members(), kind, chunks=chunks, grid_kwargs={} + ): + assert result.temperature.dims == expected.temperature.dims + nt.assert_array_equal(result.temperature, expected.temperature) + assert isinstance(result.temperature.data, da.Array) == ( + kind == "multiple" or chunks is not None + ) + if chunks == {"n_face": 2, "time": 1}: + assert result.temperature.chunks[-2:] == ((1, 1, 1), (2, 2)) + + +@pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) +@pytest.mark.parametrize( + "chunks", + [ + {"n_face": 2, "n_node": 1, "time": 1}, + {"cells": 2, "vertices": 1, "time": 1}, + {"cells": 4, "vertices": 4, "n_face": 2, "n_node": 1, "time": 1}, + ], +) +def test_open_dataset_reuse_source_chunks(reuse_open, kind, chunks): + """Translate source/canonical chunks without mutating caller options/data.""" + grid = _equal_size_grid({"vertices": "n_node", "cells": "n_face"}) + members = [ + ds.assign(velocity=("vertices", np.arange(4) + i)) + for i, ds in enumerate(_reuse_members("cells")) + ] + for result, expected in reuse_open(grid, members, kind, chunks=chunks): + assert ( + result.temperature.dims == expected.temperature.rename(cells="n_face").dims + ) + assert result.velocity.dims == expected.velocity.rename(vertices="n_node").dims + assert result.temperature.chunks[-2:] == ((1, 1, 1), (2, 2)) + assert result.velocity.chunks[-1] == (1, 1, 1, 1) + nt.assert_array_equal(result.temperature, expected.temperature) + nt.assert_array_equal(result.velocity, expected.velocity) + + +@pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("face_chunks", [False, True]) +def test_open_dataset_reuse_overlapping_aliases(reuse_open, kind, reverse, face_chunks): + """Read requested values, not newly translated keys, in either mapping order.""" + items = [("n_face", "n_node"), ("cells", "n_face")] + grid = _equal_size_grid(dict(items[::-1] if reverse else items)) + chunks = {"n_node": 1, "time": 1} + if face_chunks: + chunks["n_face"] = 2 + members = [ + ds.assign(velocity=("n_face", np.arange(4) + i)) + for i, ds in enumerate(_reuse_members("cells")) + ] + for result, expected in reuse_open(grid, members, kind, chunks=chunks): + prefix = ("member",) if kind == "multiple" else () + assert result.temperature.dims == prefix + ("time", "n_face") + assert result.velocity.dims == prefix + ("n_node",) + assert result.temperature.chunks[-2:] == ( + (1, 1, 1), + (2, 2) if face_chunks else (4,), + ) + assert result.velocity.chunks[-1] == (1, 1, 1, 1) + nt.assert_array_equal(result.temperature, expected.temperature) + nt.assert_array_equal(result.velocity, expected.velocity) + + +@pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) +@pytest.mark.parametrize("source_dims", [{}, {"vertices": "n_node", "cells": "n_face"}]) +def test_open_dataset_reuse_canonical_equal_counts(reuse_open, kind, source_dims): + """Canonical dimensions remain distinct, including topology without a map.""" + grid = _equal_size_grid(source_dims) + members = [ + ds.assign(velocity=("n_node", np.arange(4) + i)) + for i, ds in enumerate(_reuse_members()) + ] + for result, expected in reuse_open( + grid, members, kind, chunks={"n_face": 2, "n_node": 1} + ): + assert result.temperature.dims == expected.temperature.dims + assert result.velocity.dims == expected.velocity.dims + assert result.temperature.chunks[-1] == (2, 2) + assert result.velocity.chunks[-1] == (1, 1, 1, 1) + nt.assert_array_equal(result.temperature, expected.temperature) + nt.assert_array_equal(result.velocity, expected.velocity) + if source_dims: + result, _ = reuse_open(grid, _reuse_members("cells")[:1], "dataset")[0] + assert result.temperature.dims == ("time", "n_face") + + +def test_open_dataset_reuse_face_node_sequence(reuse_open): + """Face-only members must not discard mappings needed by later node data.""" + grid = _equal_size_grid({"vertices": "n_node", "cells": "n_face"}) + members = [ + _reuse_members("cells")[0], + xr.Dataset({"velocity": ("vertices", np.arange(4))}), + _reuse_members("cells")[1], + ] + for result, expected in reuse_open(grid, members, "dataset"): + xr.testing.assert_equal( + result.to_xarray(), + expected.rename( + { + dim: grid._source_dims_dict[dim] + for dim in expected.dims + if dim in grid._source_dims_dict + } + ), + ) + + +@pytest.mark.parametrize("prechunked", [False, True]) +def test_open_dataset_reuse_input_dataset_chunks(reuse_grid, reuse_open, prechunked): + """Rechunk the returned Dataset without replacing the caller's arrays.""" + members = _reuse_members() + if prechunked: + members = [member.chunk({"n_face": 4, "time": 3}) for member in members] + for result, _ in reuse_open( + reuse_grid, members, "dataset", chunks={"n_face": 2, "time": 1} + ): + assert result.temperature.chunks == ((1, 1, 1), (2, 2)) + + +@pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) +def test_open_dataset_reuse_prechunked_grid(reuse_grid, reuse_open, kind): + """With chunk_grid=False, preserve an already-chunked Grid and its caches.""" + reuse_grid.chunk(n_node=2, n_face=2) + for result, _ in reuse_open( + reuse_grid, _reuse_members(), kind, chunks={"n_face": 1}, chunk_grid=False + ): + assert result.temperature.chunks[-1] == (1, 1, 1, 1) + assert reuse_grid.face_node_connectivity.chunks[0] == (2, 2) + + +@pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) +@pytest.mark.parametrize( + "options, message", + [ + ({"use_dual": True}, "use_dual"), + ({"grid_kwargs": {"drop_variables": "face_nodes"}}, "grid_kwargs"), + ], +) +def test_open_dataset_reuse_rejects_construction_options( + reuse_grid, reuse_open, kind, options, message +): + """Reject options that only apply when constructing a Grid.""" + with pytest.raises(ValueError, match=message): + reuse_open(reuse_grid, _reuse_members(), kind, **options) + + +def test_open_dataset_reuse_requires_data(reuse_grid): + """An existing Grid alone is not a member Dataset.""" + with pytest.raises(TypeError, match="filename_or_obj"): + ux.open_dataset(reuse_grid) + + +@pytest.mark.parametrize("kind", ["file", "dataset"]) +def test_open_dataset_reuse_existing_construction( + reuse_grid, reuse_owners, tmp_path, kind +): + """Existing grid-file and grid-Dataset inputs still construct equivalent grids.""" + path = TEST_MESHFILES / "ugrid" / "quad-hexagon" / "grid.nc" + grid_source = path if kind == "file" else xr.open_dataset(path) + member = _reuse_members()[0] + data_path = tmp_path / "construction.nc" + member.to_netcdf(data_path) + result = ux.open_dataset(grid_source, data_path, chunks={"n_face": 2}) + reuse_owners.callback(result.close) + assert result.uxgrid is not reuse_grid + assert result.uxgrid == reuse_grid + assert result.temperature.dims == ("time", "n_face") + assert result.temperature.chunks[-1] == (2, 2) + nt.assert_array_equal(result.temperature, member.temperature) + xr.testing.assert_identical(result.uxgrid._ds.compute(), reuse_grid._ds) + + def test_open_mf_dataset(gridpath, datasetpath, mesh_constants): """Loads multiple datasets with their grid topology file using uxarray's open_dataset call.""" diff --git a/test/core/test_grid_reuse.py b/test/core/test_grid_reuse.py deleted file mode 100644 index 0cc9af4eb..000000000 --- a/test/core/test_grid_reuse.py +++ /dev/null @@ -1,535 +0,0 @@ -"""Regression tests for opening ensemble members with a shared Grid.""" - -from functools import partial -from pathlib import Path -from tempfile import TemporaryDirectory -from unittest import TestCase -from unittest.mock import patch - -import dask.array as da -import numpy as np -import xarray as xr - -import uxarray as ux - - -class TestGridReuse(TestCase): - def setUp(self): - """Create a synthetic tetrahedral grid and two independent members.""" - self.directory = TemporaryDirectory() - self.addCleanup(self.directory.cleanup) - # Grid conversion and UxDataset construction can lose the source - # Dataset's close callback. Keep the actual backend owners alive and - # close them before removing files, including on Windows. - for name in ("open_dataset", "open_mfdataset"): - opener = getattr(xr, name) - patcher = patch.object( - xr, name, side_effect=partial(self._open_source, opener) - ) - patcher.start() - self.addCleanup(patcher.stop) - self.grid_path = Path(self.directory.name) / "grid.nc" - # Equal node and face counts ensure source names, rather than sizes, - # distinguish node-centered and face-centered variables after reuse. - self.grid_ds = xr.Dataset( - { - "mesh": xr.DataArray( - 0, - attrs={ - "cf_role": "mesh_topology", - "topology_dimension": 2, - "node_coordinates": "lon lat", - "face_node_connectivity": "face_nodes", - "node_dimension": "vertices", - "face_dimension": "cells", - }, - ), - "lon": ("vertices", [0.0, 120.0, -120.0, 0.0]), - "lat": ("vertices", [0.0, 0.0, 0.0, 90.0]), - "face_nodes": ( - ("cells", "corners"), - [[0, 1, 2], [0, 3, 1], [1, 3, 2], [2, 3, 0]], - {"start_index": 0}, - ), - } - ) - self.grid_ds.to_netcdf(self.grid_path) - self.grid = ux.open_grid(self.grid_path) - self.addCleanup(self.grid._ds.close) - self.source_dims = self.grid._source_dims_dict.copy() - self.grid_snapshot = self.grid._ds.copy(deep=True) - self.members = [ - xr.Dataset( - {"temperature": (("time", "cells"), np.arange(12).reshape(3, 4) + i)}, - coords={"time": [0, 1, 2]}, - ) - for i in range(2) - ] - self.paths = [] - for i, member in enumerate(self.members): - path = Path(self.directory.name) / f"member_{i}.nc" - member.to_netcdf(path) - self.paths.append(path) - - def _open_source(self, opener, *args, **kwargs): - """Retain the original xarray resource owner without changing file I/O.""" - source = opener(*args, **kwargs) - self.addCleanup(source.close) - return source - - def assert_grid_unchanged(self): - """Opening data must not alter the shared grid or its source mapping.""" - self.assertEqual(self.grid._source_dims_dict, self.source_dims) - xr.testing.assert_identical(self.grid._ds, self.grid_snapshot) - - def test_open_dataset_reuses_grid_for_file_and_dataset_members(self): - """A single Grid is attached by identity to independent member data.""" - for sources in (self.paths, self.members): - for source, expected in zip(sources, self.members): - with self.subTest(source=type(source).__name__): - result = ux.open_dataset(self.grid, source) - self.addCleanup(result.close) - self.assertIs(result.uxgrid, self.grid) - self.assertEqual(result.temperature.dims, ("time", "n_face")) - np.testing.assert_array_equal( - result.temperature, expected.temperature - ) - self.assertEqual( - result.source_datasets, - None if isinstance(source, xr.Dataset) else str(source), - ) - self.assert_grid_unchanged() - - def test_members_with_different_grid_dimensions_preserve_mapping(self): - """Opening face-only data must not lose the next member's node mapping.""" - node_data = xr.Dataset({"velocity": ("vertices", np.arange(4))}) - for member, name, dimension in ( - (self.members[0], "temperature", "n_face"), - (node_data, "velocity", "n_node"), - (self.members[1], "temperature", "n_face"), - ): - result = ux.open_dataset(self.grid, member) - self.assertEqual(result[name].dims[-1], dimension) - self.assert_grid_unchanged() - - def test_chunks_preserve_grid_and_caller_options(self): - """Chunk data with source or UGRID names without rechunking the Grid.""" - for source in (self.paths[0], self.members[0]): - for chunk_grid in (True, False): - for chunks in ( - {"n_face": 2, "time": 1}, - {"cells": 2, "time": 1}, - {"cells": 4, "n_face": 2, "time": 1}, - ): - with self.subTest( - source=type(source).__name__, - chunk_grid=chunk_grid, - chunks=chunks, - ): - original_chunks = chunks.copy() - grid_kwargs = {} - result = ux.open_dataset( - self.grid, - source, - chunks=chunks, - chunk_grid=chunk_grid, - grid_kwargs=grid_kwargs, - ) - self.addCleanup(result.close) - self.assertIs(result.uxgrid, self.grid) - self.assertIsInstance(result.temperature.data, da.Array) - self.assertEqual(result.temperature.chunks, ((1, 1, 1), (2, 2))) - np.testing.assert_array_equal( - result.temperature.compute(), self.members[0].temperature - ) - self.assertEqual(chunks, original_chunks) - self.assertEqual(grid_kwargs, {}) - self.assert_grid_unchanged() - - def test_scalar_and_auto_chunks(self): - """All non-dict chunk forms reach the member data unchanged.""" - for source in (self.paths[0], self.members[0]): - for chunks in (-1, "auto", {}, None): - with self.subTest(source=type(source).__name__, chunks=chunks): - result = ux.open_dataset(self.grid, source, chunks=chunks) - self.addCleanup(result.close) - self.assertEqual( - isinstance(result.temperature.data, da.Array), - chunks is not None, - ) - self.assert_grid_unchanged() - - def test_prechunked_grid_is_not_rechunked(self): - """Each member retains the same already-chunked grid arrays.""" - self.grid.chunk(n_node=2, n_face=2) - connectivity = self.grid.face_node_connectivity.data - result = ux.open_dataset(self.grid, self.paths[0], chunks={"n_face": 1}) - self.addCleanup(result.close) - self.assertIs(result.uxgrid, self.grid) - self.assertIs(self.grid.face_node_connectivity.data, connectivity) - self.assertEqual(result.temperature.chunks[-1], (1, 1, 1, 1)) - self.assertEqual(self.grid.face_node_connectivity.chunks[0], (2, 2)) - - def test_existing_grid_does_not_open_or_reconstruct_grid(self): - """An in-memory Grid/Dataset pair must not invoke file backends.""" - with ( - patch("uxarray.core.api.open_grid") as open_grid, - patch("xarray.open_dataset") as open_dataset, - ): - result = ux.open_dataset(self.grid, self.members[0], chunks={"n_face": 2}) - self.assertIs(result.uxgrid, self.grid) - open_grid.assert_not_called() - open_dataset.assert_not_called() - - def test_open_mfdataset_reuses_grid(self): - """The multi-file API shares the same reuse and chunking contract.""" - chunks = {"n_face": 2, "time": 1} - with patch("uxarray.core.api.open_grid") as open_grid: - result = ux.open_mfdataset( - self.grid, - self.paths, - chunks=chunks, - combine="nested", - concat_dim="member", - ) - self.addCleanup(result.close) - open_grid.assert_not_called() - self.assertIs(result.uxgrid, self.grid) - self.assertEqual(result.temperature.dims, ("member", "time", "n_face")) - self.assertEqual(result.temperature.chunks, ((1, 1), (1, 1, 1), (2, 2))) - np.testing.assert_array_equal( - result.temperature.compute(), - np.stack([member.temperature.values for member in self.members]), - ) - self.assertEqual(chunks, {"n_face": 2, "time": 1}) - self.assert_grid_unchanged() - - def test_grid_construction_options_are_rejected(self): - """Options requiring grid construction should not be silently ignored.""" - for opener, data in ( - (ux.open_dataset, self.members[0]), - (ux.open_mfdataset, self.paths), - ): - for options, message in ( - ({"use_dual": True}, "use_dual"), - ({"grid_kwargs": {"drop_variables": "face_nodes"}}, "grid_kwargs"), - ): - with self.subTest(opener=opener.__name__, options=options): - with self.assertRaisesRegex(ValueError, message): - opener(self.grid, data, **options) - self.assert_grid_unchanged() - - def test_grid_requires_separate_data(self): - """A Grid alone does not provide member data.""" - with self.assertRaisesRegex(TypeError, "filename_or_obj"): - ux.open_dataset(self.grid) - - def test_canonical_dimensions_are_not_inferred_from_equal_sizes(self): - """Canonical face and node dimensions remain distinct across members.""" - members = [ - member.rename({"cells": "n_face"}).assign( - velocity=("n_node", np.arange(4) + i) - ) - for i, member in enumerate(self.members) - ] - paths = [] - for i, member in enumerate(members): - path = Path(self.directory.name) / f"canonical_{i}.nc" - member.to_netcdf(path) - paths.append(path) - for sources in (members, paths): - for source, expected in zip(sources, members): - result = ux.open_dataset( - self.grid, source, chunks={"n_face": 2, "n_node": 1} - ) - self.addCleanup(result.close) - self.assertIs(result.uxgrid, self.grid) - self.assertEqual(result.temperature.dims, ("time", "n_face")) - self.assertEqual(result.velocity.dims, ("n_node",)) - self.assertEqual(result.temperature.chunks[-1], (2, 2)) - self.assertEqual(result.velocity.chunks, ((1, 1, 1, 1),)) - np.testing.assert_array_equal(result.temperature, expected.temperature) - np.testing.assert_array_equal(result.velocity, expected.velocity) - self.assert_grid_unchanged() - combined = ux.open_mfdataset( - self.grid, - paths, - combine="nested", - concat_dim="member", - chunks={"n_face": 2, "n_node": 1}, - ) - self.addCleanup(combined.close) - self.assertEqual(combined.temperature.dims, ("member", "time", "n_face")) - self.assertEqual(combined.velocity.dims, ("member", "n_node")) - np.testing.assert_array_equal( - combined.temperature, - np.stack([member.temperature.values for member in members]), - ) - self.assert_grid_unchanged() - # Canonical members must not disrupt later source-dimension members. - native = ux.open_dataset(self.grid, self.members[0]) - self.assertEqual(native.temperature.dims, ("time", "n_face")) - self.assert_grid_unchanged() - - def test_topology_grid_without_source_mapping(self): - """A topology-created Grid can reuse already-standardized data.""" - grid = ux.Grid.from_topology( - node_lon=self.grid.node_lon.values, - node_lat=self.grid.node_lat.values, - face_node_connectivity=self.grid.face_node_connectivity.values, - ) - member = xr.Dataset({"temperature": ("n_face", np.arange(4))}) - result = ux.open_dataset(grid, member, chunks={"n_face": 2}) - self.assertIs(result.uxgrid, grid) - self.assertEqual(result.temperature.dims, ("n_face",)) - self.assertEqual(result.temperature.chunks, ((2, 2),)) - np.testing.assert_array_equal(result.temperature, member.temperature) - self.assertEqual(grid._source_dims_dict, {}) - - def test_explicit_source_mapping_takes_precedence(self): - """An explicit source mapping wins even for a canonical-looking name.""" - self.grid._source_dims_dict = {"n_face": "n_node"} - result = ux.open_dataset( - self.grid, xr.Dataset({"velocity": ("n_face", np.arange(4))}) - ) - self.assertEqual(result.velocity.dims, ("n_node",)) - self.assertEqual(self.grid._source_dims_dict, {"n_face": "n_node"}) - - def test_structured_grid_chunks_use_source_dimensions(self): - """Structured grids retain their tuple-valued source mapping on reuse.""" - grid = ux.Grid.from_structured( - lon=xr.DataArray([0.0, 20.0, 40.0, 60.0], dims="x"), - lat=xr.DataArray([-30.0, 0.0, 30.0], dims="y"), - ) - source_dims = grid._source_dims_dict.copy() - member = xr.Dataset({"temperature": (("y", "x"), np.arange(12).reshape(3, 4))}) - path = Path(self.directory.name) / "structured.nc" - member.to_netcdf(path) - for source in (member, path): - chunks = {"x": 2, "y": 1} - result = ux.open_dataset(grid, source, chunks=chunks) - self.addCleanup(result.close) - self.assertIs(result.uxgrid, grid) - self.assertEqual(result.temperature.dims, ("n_face",)) - self.assertEqual(result.temperature.chunks, ((2, 2, 2, 2, 2, 2),)) - np.testing.assert_array_equal(result.temperature, np.arange(12)) - self.assertEqual(grid._source_dims_dict, source_dims) - self.assertEqual(chunks, {"x": 2, "y": 1}) - - def test_structured_grid_accepts_already_flattened_members(self): - """Do not flatten canonical structured member data a second time.""" - grid = ux.Grid.from_structured( - lon=xr.DataArray([0.0, 20.0, 40.0, 60.0], dims="x"), - lat=xr.DataArray([-30.0, 0.0, 30.0], dims="y"), - ) - source_dims = grid._source_dims_dict.copy() - members = [ - xr.Dataset({"temperature": ("n_face", np.arange(12) + i)}) for i in range(2) - ] - paths = [] - for i, member in enumerate(members): - path = Path(self.directory.name) / f"flattened_{i}.nc" - member.to_netcdf(path) - paths.append(path) - for sources in (members, paths): - for source, expected in zip(sources, members): - result = ux.open_dataset(grid, source, chunks={"n_face": 2}) - self.addCleanup(result.close) - self.assertIs(result.uxgrid, grid) - self.assertEqual(result.temperature.dims, ("n_face",)) - self.assertEqual(result.temperature.chunks, ((2, 2, 2, 2, 2, 2),)) - np.testing.assert_array_equal(result.temperature, expected.temperature) - self.assertEqual(grid._source_dims_dict, source_dims) - combined = ux.open_mfdataset( - grid, - paths, - combine="nested", - concat_dim="member", - chunks={"n_face": 2}, - ) - self.addCleanup(combined.close) - self.assertIs(combined.uxgrid, grid) - self.assertEqual(combined.temperature.dims, ("member", "n_face")) - self.assertEqual(combined.temperature.chunks, ((1, 1), (2, 2, 2, 2, 2, 2))) - np.testing.assert_array_equal( - combined.temperature, - np.stack([member.temperature.values for member in members]), - ) - self.assertEqual(grid._source_dims_dict, source_dims) - - def test_structured_grid_accepts_canonical_node_and_edge_members(self): - """Canonical node and edge data have no source axes left to flatten.""" - grid = ux.Grid.from_structured( - lon=xr.DataArray([0.0, 20.0, 40.0, 60.0], dims="x"), - lat=xr.DataArray([-30.0, 0.0, 30.0], dims="y"), - ) - # n_edge lazily creates connectivity; snapshot the intended cache first. - counts = {"n_node": grid.n_node, "n_edge": grid.n_edge} - source_dims = grid._source_dims_dict.copy() - snapshot = grid._ds.copy(deep=True) - for dimension, size in counts.items(): - members = [ - xr.Dataset({"value": (dimension, np.arange(size) + i)}) - for i in range(2) - ] - member_snapshots = [member.copy(deep=True) for member in members] - paths = [] - for i, member in enumerate(members): - path = Path(self.directory.name) / f"structured_{dimension}_{i}.nc" - member.to_netcdf(path) - paths.append(path) - cases = ( - [("dataset", member, member.value.values) for member in members] - + [ - ("file", path, member.value.values) - for path, member in zip(paths, members) - ] - + [ - ( - "multiple", - paths, - np.stack([member.value.values for member in members]), - ) - ] - ) - chunks = {dimension: 2} - expected_chunks = (2,) * (size // 2) + ((1,) if size % 2 else ()) - for case_index, (kind, source, expected) in enumerate(cases): - with self.subTest( - dimension=dimension, kind=kind, case_index=case_index - ): - if kind == "multiple": - result = ux.open_mfdataset( - grid, - source, - combine="nested", - concat_dim="member", - chunks=chunks, - ) - expected_dims = ("member", dimension) - else: - result = ux.open_dataset(grid, source, chunks=chunks) - expected_dims = (dimension,) - self.addCleanup(result.close) - self.assertIs(result.uxgrid, grid) - self.assertEqual(result.value.dims, expected_dims) - self.assertIsInstance(result.value.data, da.Array) - all_chunks = (expected_chunks,) - if kind == "multiple": - all_chunks = ((1, 1),) + all_chunks - self.assertEqual(result.value.chunks, all_chunks) - np.testing.assert_array_equal(result.value.compute(), expected) - self.assertEqual(chunks, {dimension: 2}) - self.assertEqual(grid._source_dims_dict, source_dims) - xr.testing.assert_identical(grid._ds, snapshot) - for original, saved in zip(members, member_snapshots): - xr.testing.assert_identical(original, saved) - # Reuse after canonical data must still flatten ordinary face members. - native = xr.Dataset({"temperature": (("y", "x"), np.arange(12).reshape(3, 4))}) - result = ux.open_dataset(grid, native) - self.assertIs(result.uxgrid, grid) - self.assertEqual(result.temperature.dims, ("n_face",)) - np.testing.assert_array_equal(result.temperature, np.arange(12)) - self.assertEqual(grid._source_dims_dict, source_dims) - xr.testing.assert_identical(grid._ds, snapshot) - - def test_overlapping_chunk_aliases_read_only_requested_values(self): - """Alias translation is independent of order and newly created keys.""" - members = [ - member.assign(velocity=("n_face", np.arange(4) + i)) - for i, member in enumerate(self.members) - ] - member_snapshots = [member.copy(deep=True) for member in members] - paths = [] - for i, member in enumerate(members): - path = Path(self.directory.name) / f"overlapping_alias_{i}.nc" - member.to_netcdf(path) - paths.append(path) - for source_dims in ( - {"n_face": "n_node", "cells": "n_face"}, - {"cells": "n_face", "n_face": "n_node"}, - ): - grid = ux.Grid.from_topology( - node_lon=self.grid.node_lon.values, - node_lat=self.grid.node_lat.values, - face_node_connectivity=self.grid.face_node_connectivity.values, - dims_dict=source_dims.copy(), - ) - snapshot = grid._ds.copy(deep=True) - for requested in ( - {"n_face": 2, "n_node": 1, "time": 1}, - {"n_node": 1, "time": 1}, - ): - original_chunks = requested.copy() - expected_face_chunks = (2, 2) if "n_face" in requested else (4,) - cases = ( - [("dataset", member, member) for member in members] - + [("file", path, member) for path, member in zip(paths, members)] - + [("multiple", paths, None)] - ) - for case_index, (kind, source, member) in enumerate(cases): - with self.subTest( - source_dims=source_dims, - requested=requested, - kind=kind, - case_index=case_index, - ): - if kind == "multiple": - result = ux.open_mfdataset( - grid, - source, - combine="nested", - concat_dim="member", - chunks=requested, - ) - temperature = np.stack( - [m.temperature.values for m in members] - ) - velocity = np.stack([m.velocity.values for m in members]) - temperature_dims = ("member", "time", "n_face") - velocity_dims = ("member", "n_node") - else: - result = ux.open_dataset(grid, source, chunks=requested) - temperature = member.temperature.values - velocity = member.velocity.values - temperature_dims = ("time", "n_face") - velocity_dims = ("n_node",) - self.addCleanup(result.close) - self.assertIs(result.uxgrid, grid) - self.assertEqual(result.temperature.dims, temperature_dims) - self.assertEqual(result.velocity.dims, velocity_dims) - self.assertIsInstance(result.temperature.data, da.Array) - self.assertIsInstance(result.velocity.data, da.Array) - temperature_chunks = ((1, 1, 1), expected_face_chunks) - velocity_chunks = ((1, 1, 1, 1),) - if kind == "multiple": - temperature_chunks = ((1, 1),) + temperature_chunks - velocity_chunks = ((1, 1),) + velocity_chunks - self.assertEqual(result.temperature.chunks, temperature_chunks) - self.assertEqual(result.velocity.chunks, velocity_chunks) - np.testing.assert_array_equal( - result.temperature.compute(), temperature - ) - np.testing.assert_array_equal( - result.velocity.compute(), velocity - ) - self.assertEqual(requested, original_chunks) - self.assertEqual(grid._source_dims_dict, source_dims) - xr.testing.assert_identical(grid._ds, snapshot) - for original, saved in zip(members, member_snapshots): - xr.testing.assert_identical(original, saved) - - def test_existing_grid_file_and_dataset_inputs(self): - """Previously supported grid inputs still construct equivalent grids.""" - for grid_source in (self.grid_path, self.grid_ds): - result = ux.open_dataset(grid_source, self.paths[0], chunks={"n_face": 2}) - self.addCleanup(result.close) - self.addCleanup(result.uxgrid._ds.close) - self.assertIsNot(result.uxgrid, self.grid) - self.assertEqual(result.uxgrid, self.grid) - self.assertEqual(result.temperature.dims, ("time", "n_face")) - self.assertEqual(result.temperature.chunks[-1], (2, 2)) - np.testing.assert_array_equal( - result.temperature, self.members[0].temperature - ) - xr.testing.assert_identical(result.uxgrid._ds.compute(), self.grid_snapshot) diff --git a/uxarray/core/api.py b/uxarray/core/api.py index 6e35bcea3..3a259c954 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -375,19 +375,19 @@ def open_dataset( grid_kwargs: dict[str, Any] | None = None, **kwargs: dict[str, Any], ) -> UxDataset: - """Wraps ``xarray.open_dataset()`` for loading in a dataset paired with a grid file. + """Wraps ``xarray.open_dataset()`` for loading data paired with an unstructured grid. Parameters ---------- grid_filename_or_obj : str | os.PathLike[Any] | dict | xr.Dataset | Grid Grid information for the ``UxDataset``. Strings and Path objects are interpreted as a path to a grid - file. Xarray Datasets assume that each member variable is in the UGRID conventions and will be used to + file. An ``xarray.Dataset`` containing grid information in a supported unstructured format is used to create a ``ux.Grid``. Similarly, a dictionary containing UGRID variables can be used to create a ``ux.Grid``. A path to a directory containing grid files (e.g. a FESOM2 ASCII grid) is also accepted, but only when ``filename_or_obj`` is provided; directory input is not supported in the single-argument - form below. An existing ``ux.Grid`` is attached directly, allowing its - topology and cached properties to be shared across datasets. A separate - ``filename_or_obj`` is required when passing a ``Grid``. + form below. An existing ``ux.Grid`` can be used directly, allowing its + topology and cached properties to be shared across datasets as long as + they conform to it. filename_or_obj : str | os.PathLike[Any] | xr.Dataset, optional String or Path object as a path to a netCDF file or an OpenDAP URL that stores the actual data set, or an already-open ``xarray.Dataset``. It @@ -398,9 +398,11 @@ def open_dataset( the grid and also treated as data; all variables it contains, including grid variables, will be included as data. chunks : int, dict, 'auto' or None, default: None - If provided, used to load the data into dask arrays. With an existing - ``Grid``, this also applies to an already-open ``xarray.Dataset`` - without modifying the input dataset. + If provided, used to load file-based data into dask arrays. For an + ``xarray.Dataset`` data input, this parameter is applied when + ``grid_filename_or_obj`` is an existing ``Grid``. In that case, the + chunks apply to the returned ``UxDataset``; the input dataset retains + its original dimensions, values, and chunking. - ``chunks="auto"`` will use dask ``auto`` chunking taking into account the engine preferred chunks. @@ -411,9 +413,6 @@ def open_dataset( size, generally identical to the format's chunk size. If not available, a single chunk for all arrays. - For structured data that still has its original spatial dimensions, - use those names (such as ``lon`` and ``lat``) to chunk before - flattening. You can also rechunk the returned ``n_face`` dimension. chunk_grid : bool, default: True If valid chunks are passed in, determines whether to also apply the same chunks when constructing a ``Grid``. An existing ``Grid`` is reused @@ -540,9 +539,9 @@ def open_mfdataset( Parameters ---------- grid_filename_or_obj : str | os.PathLike[Any] | dict | xr.Dataset | Grid - Strings and Path objects are interpreted as a path to a grid file. Xarray Datasets assume that - each member variable is in the UGRID conventions and will be used to create a ``ux.Grid``. Similarly, a dictionary - containing UGRID variables can be used to create a ``ux.Grid``. An existing + Strings and Path objects are interpreted as a path to a grid file. An ``xarray.Dataset`` containing + grid information in a supported unstructured format is used to create a ``ux.Grid``. Similarly, a + dictionary containing UGRID variables can be used to create a ``ux.Grid``. An existing ``ux.Grid`` is attached directly without reconstructing or modifying it. paths : string, required Either a string glob in the form ``"path/to/my/files/*.nc"`` or an explicit @@ -559,9 +558,6 @@ def open_mfdataset( size, generally identical to the format's chunk size. If not available, a single chunk for all arrays. - For structured data that still has its original spatial dimensions, - use those names (such as ``lon`` and ``lat``) to chunk before - flattening. You can also rechunk the returned ``n_face`` dimension. chunk_grid : bool, default: True If valid chunks are passed in, determines whether to also apply the same chunks when constructing a ``Grid``. An existing ``Grid`` is reused From d2c2ad9fef5400967cf3d93f9cd64d0ac3aaf1d1 Mon Sep 17 00:00:00 2001 From: Lubai <1174273035@qq.com> Date: Wed, 28 Oct 2026 03:53:42 +0800 Subject: [PATCH 6/7] docs: leave Grid conformity details to follow-up --- uxarray/core/api.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/uxarray/core/api.py b/uxarray/core/api.py index 3a259c954..ed3c2ab2f 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -386,8 +386,7 @@ def open_dataset( ``ux.Grid``. A path to a directory containing grid files (e.g. a FESOM2 ASCII grid) is also accepted, but only when ``filename_or_obj`` is provided; directory input is not supported in the single-argument form below. An existing ``ux.Grid`` can be used directly, allowing its - topology and cached properties to be shared across datasets as long as - they conform to it. + topology and cached properties to be shared across datasets. filename_or_obj : str | os.PathLike[Any] | xr.Dataset, optional String or Path object as a path to a netCDF file or an OpenDAP URL that stores the actual data set, or an already-open ``xarray.Dataset``. It From 723f6d8a273b42aa8fe26ec755ed0cc7048d14cc Mon Sep 17 00:00:00 2001 From: Lubai <1174273035@qq.com> Date: Tue, 3 Nov 2026 23:38:27 +0800 Subject: [PATCH 7/7] Reject implicit rechunking when reusing an existing Grid --- test/core/test_api.py | 40 ++++++++++++++++++++++++++-------------- uxarray/core/api.py | 18 ++++++++++++++---- 2 files changed, 40 insertions(+), 18 deletions(-) diff --git a/test/core/test_api.py b/test/core/test_api.py index b1905b97c..871ebe68c 100644 --- a/test/core/test_api.py +++ b/test/core/test_api.py @@ -143,9 +143,6 @@ def open_members(grid, members, kind, **kwargs): counter += 1 member.to_netcdf(path) paths.append(path) - # Chunk ownership for True is a separate open review question. - if chunks is not None: - kwargs.setdefault("chunk_grid", False) with patch("uxarray.core.api.open_grid") as construct: try: if kind == "multiple": @@ -212,11 +209,15 @@ def _equal_size_grid(source_dims): @pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) -@pytest.mark.parametrize("chunks", [None, -1, "auto", {}, {"n_face": 2, "time": 1}]) -def test_open_dataset_reuse_members(reuse_grid, reuse_open, kind, chunks): +@pytest.mark.parametrize("chunks, chunk_grid", [ + (None, True), (None, False), (-1, False), ("auto", False), + ({}, False), ({"n_face": 2, "time": 1}, False), +]) +def test_open_dataset_reuse_members(reuse_grid, reuse_open, kind, chunks, chunk_grid): """Reuse topology/cached tree without grid I/O, preserving member values.""" for result, expected in reuse_open( - reuse_grid, _reuse_members(), kind, chunks=chunks, grid_kwargs={} + reuse_grid, _reuse_members(), kind, chunks=chunks, + chunk_grid=chunk_grid, grid_kwargs={} ): assert result.temperature.dims == expected.temperature.dims nt.assert_array_equal(result.temperature, expected.temperature) @@ -243,7 +244,7 @@ def test_open_dataset_reuse_source_chunks(reuse_open, kind, chunks): ds.assign(velocity=("vertices", np.arange(4) + i)) for i, ds in enumerate(_reuse_members("cells")) ] - for result, expected in reuse_open(grid, members, kind, chunks=chunks): + for result, expected in reuse_open(grid, members, kind, chunks=chunks, chunk_grid=False): assert ( result.temperature.dims == expected.temperature.rename(cells="n_face").dims ) @@ -268,7 +269,7 @@ def test_open_dataset_reuse_overlapping_aliases(reuse_open, kind, reverse, face_ ds.assign(velocity=("n_face", np.arange(4) + i)) for i, ds in enumerate(_reuse_members("cells")) ] - for result, expected in reuse_open(grid, members, kind, chunks=chunks): + for result, expected in reuse_open(grid, members, kind, chunks=chunks, chunk_grid=False): prefix = ("member",) if kind == "multiple" else () assert result.temperature.dims == prefix + ("time", "n_face") assert result.velocity.dims == prefix + ("n_node",) @@ -291,7 +292,7 @@ def test_open_dataset_reuse_canonical_equal_counts(reuse_open, kind, source_dims for i, ds in enumerate(_reuse_members()) ] for result, expected in reuse_open( - grid, members, kind, chunks={"n_face": 2, "n_node": 1} + grid, members, kind, chunks={"n_face": 2, "n_node": 1}, chunk_grid=False ): assert result.temperature.dims == expected.temperature.dims assert result.velocity.dims == expected.velocity.dims @@ -332,22 +333,33 @@ def test_open_dataset_reuse_input_dataset_chunks(reuse_grid, reuse_open, prechun if prechunked: members = [member.chunk({"n_face": 4, "time": 3}) for member in members] for result, _ in reuse_open( - reuse_grid, members, "dataset", chunks={"n_face": 2, "time": 1} + reuse_grid, members, "dataset", chunks={"n_face": 2, "time": 1}, chunk_grid=False ): assert result.temperature.chunks == ((1, 1, 1), (2, 2)) @pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) -def test_open_dataset_reuse_prechunked_grid(reuse_grid, reuse_open, kind): - """With chunk_grid=False, preserve an already-chunked Grid and its caches.""" +@pytest.mark.parametrize("chunks, chunk_grid", [(None, True), (None, False), ({"n_face": 1}, False)]) +def test_open_dataset_reuse_prechunked_grid(reuse_grid, reuse_open, kind, chunks, chunk_grid): + """Preserve an already-chunked Grid and its caches for supported options.""" reuse_grid.chunk(n_node=2, n_face=2) for result, _ in reuse_open( - reuse_grid, _reuse_members(), kind, chunks={"n_face": 1}, chunk_grid=False + reuse_grid, _reuse_members(), kind, chunks=chunks, chunk_grid=chunk_grid ): - assert result.temperature.chunks[-1] == (1, 1, 1, 1) + if chunks is not None: + assert result.temperature.chunks[-1] == (1, 1, 1, 1) assert reuse_grid.face_node_connectivity.chunks[0] == (2, 2) +@pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) +@pytest.mark.parametrize("chunks", [-1, "auto", {}, {"n_face": 2, "time": 1}]) +@pytest.mark.parametrize("options", [{}, {"chunk_grid": True}]) +def test_open_dataset_reuse_rejects_grid_chunks(reuse_grid, reuse_open, kind, chunks, options): + """Reject implicit or explicit grid rechunking without changing shared state.""" + with pytest.raises(ValueError, match=r"Grid, non-None chunks, and chunk_grid=True.*Grid\.chunk\(\)"): + reuse_open(reuse_grid, _reuse_members(), kind, chunks=chunks, **options) + + @pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) @pytest.mark.parametrize( "options, message", diff --git a/uxarray/core/api.py b/uxarray/core/api.py index ec8f386ad..9039aa5ae 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -415,8 +415,10 @@ def open_dataset( chunk_grid : bool, default: True If valid chunks are passed in, determines whether to also apply the same chunks when constructing a ``Grid``. An existing ``Grid`` is reused - unchanged, including its chunking; call ``Grid.chunk`` explicitly - before opening data to change its chunking. + unchanged, including its chunking. With an existing ``Grid``, passing + non-None ``chunks`` and ``chunk_grid=True`` raises ``ValueError``. + Set ``chunk_grid=False`` to chunk only the data, or call ``Grid.chunk()`` + explicitly before opening data to change the grid's chunking. use_dual : bool, optional Specify whether to use the primal (use_dual=False) or dual (use_dual=True) mesh if the file type is MPAS. ``use_dual=True`` is not supported for an existing ``Grid``; select @@ -560,8 +562,10 @@ def open_mfdataset( chunk_grid : bool, default: True If valid chunks are passed in, determines whether to also apply the same chunks when constructing a ``Grid``. An existing ``Grid`` is reused - unchanged, including its chunking; call ``Grid.chunk`` explicitly - before opening data to change its chunking. + unchanged, including its chunking. With an existing ``Grid``, passing + non-None ``chunks`` and ``chunk_grid=True`` raises ``ValueError``. + Set ``chunk_grid=False`` to chunk only the data, or call ``Grid.chunk()`` + explicitly before opening data to change the grid's chunking. use_dual : bool, optional Specify whether to use the primal (use_dual=False) or dual (use_dual=True) mesh if the file type is MPAS. ``use_dual=True`` is not supported for an existing ``Grid``; select @@ -643,6 +647,12 @@ def _get_grid( "grid_kwargs must be empty when passing an existing Grid. " "Apply these options when constructing the Grid." ) + if chunks is not None and chunk_grid: + raise ValueError( + "Providing Grid, non-None chunks, and chunk_grid=True simultaneously " + "is not supported. Call Grid.chunk() explicitly to change the grid's " + "chunking, or set chunk_grid=False to chunk only the data." + ) if isinstance(chunks, dict): # Keep the caller's options and the shared Grid unchanged. Match # UGRID chunk keys to the data's source dimensions without grid I/O.