diff --git a/benchmarks/open_dataset.py b/benchmarks/open_dataset.py new file mode 100644 index 000000000..873cbee21 --- /dev/null +++ b/benchmarks/open_dataset.py @@ -0,0 +1,49 @@ +"""Opening synthetic data with file, Dataset, or existing Grid inputs.""" + +from pathlib import Path +from tempfile import TemporaryDirectory + +import numpy as np +import xarray as xr + +import uxarray as ux + + +class OpenDataset: + 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_api.py b/test/core/test_api.py index ed29c2bce..871ebe68c 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,315 @@ 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) + 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, 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, + chunk_grid=chunk_grid, 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, chunk_grid=False): + 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, chunk_grid=False): + 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}, chunk_grid=False + ): + 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}, chunk_grid=False + ): + assert result.temperature.chunks == ((1, 1, 1), (2, 2)) + + +@pytest.mark.parametrize("kind", ["dataset", "file", "multiple"]) +@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=chunks, chunk_grid=chunk_grid + ): + 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", + [ + ({"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/uxarray/core/api.py b/uxarray/core/api.py index 015c01b71..9039aa5ae 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, ) @@ -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, @@ -375,17 +375,18 @@ 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_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. + form below. An existing ``ux.Grid`` can be used directly, allowing its + 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 @@ -396,7 +397,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 grid into dask arrays. + 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. @@ -406,12 +411,21 @@ 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. + 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. 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 + 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 - 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`` + 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. 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 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,21 @@ 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. + 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. 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 + 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 +600,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 = {} @@ -578,7 +622,7 @@ 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)) @@ -591,7 +635,34 @@ def _get_grid( use_dual, grid_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( + "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 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. + 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 requested_chunks: + chunks[source_dim] = requested_chunks[ugrid_dim] + return grid_filename_or_obj, chunks + grid_kwargs["data_chunks"] = chunks grid_kwargs["chunk_grid"] = chunk_grid grid_kwargs["return_chunks"] = True diff --git a/uxarray/core/utils.py b/uxarray/core/utils.py index 441b10ad8..cb22baafc 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,