diff --git a/test/test_api.py b/test/test_api.py index e6b05698c..879551073 100644 --- a/test/test_api.py +++ b/test/test_api.py @@ -4,6 +4,8 @@ import numpy.testing as nt import uxarray as ux +import xarray as xr +import numpy as np try: import constants @@ -14,7 +16,6 @@ class TestAPI(TestCase): - geoflow_data_path = current_path / "meshfiles" / "ugrid" / "geoflow-small" gridfile_geoflow = current_path / "meshfiles" / "ugrid" / "geoflow-small" / "grid.nc" geoflow_data_v1 = geoflow_data_path / "v1.nc" @@ -128,3 +129,14 @@ def test_copy_dataarray(self): # Check that the deep copy is a deep copy assert (v1_uxdata_array_copy_deep.uxgrid == v1_uxdata_array.uxgrid) assert (v1_uxdata_array_copy_deep.uxgrid is not v1_uxdata_array.uxgrid) + + def test_open_dataset_grid_kwargs(self): + """Drops ``Mesh2_face_nodes`` from the inputted grid file using + ``grid_kwargs``""" + + with self.assertRaises(KeyError): + # attempt to open a dataset after dropping face nodes should raise a KeyError + uxds = ux.open_dataset( + self.gridfile_ne30, + self.dsfile_var2_ne30, + grid_kwargs={"drop_variables": "Mesh2_face_nodes"}) diff --git a/uxarray/core/api.py b/uxarray/core/api.py index 6a73f632f..e4f096562 100644 --- a/uxarray/core/api.py +++ b/uxarray/core/api.py @@ -106,6 +106,7 @@ def open_dataset(grid_filename_or_obj: str, islatlon: Optional[bool] = False, isconcave: Optional[bool] = False, use_dual: Optional[bool] = False, + grid_kwargs: Optional[Dict[str, Any]] = {}, **kwargs: Dict[str, Any]) -> UxDataset: """Wraps ``xarray.open_dataset()``, given a grid topology definition with a single dataset file or object with corresponding data. @@ -141,6 +142,12 @@ def open_dataset(grid_filename_or_obj: str, 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 + grid_kwargs : Dict[str, Any], optional + Additional arguments passed on to ``xarray.open_dataset`` when opening up a Grid File. Refer to the + [xarray + docs](https://xarray.pydata.org/en/stable/generated/xarray.open_dataset.html) + for accepted keyword arguments. + **kwargs : Dict[str, Any] Additional arguments passed on to ``xarray.open_dataset``. Refer to the [xarray @@ -160,20 +167,21 @@ def open_dataset(grid_filename_or_obj: str, >>> import uxarray as ux >>> ux_ds = ux.open_dataset("grid_filename.g", "grid_filename_vortex.nc") + :param grid_kwargs: """ - ## Grid definition + # Grid definition uxgrid = open_grid(grid_filename_or_obj, gridspec=gridspec, vertices=vertices, islatlon=islatlon, isconcave=isconcave, - use_dual=use_dual) + use_dual=use_dual, + **grid_kwargs) - ## UxDataset + # UxDataset ds = xr.open_dataset(filename_or_obj, decode_times=False, **kwargs) # type: ignore - uxds = UxDataset(ds, uxgrid=uxgrid, source_datasets=str(filename_or_obj)) return uxds @@ -186,6 +194,7 @@ def open_mfdataset(grid_filename_or_obj: str, islatlon: Optional[bool] = False, isconcave: Optional[bool] = False, use_dual: Optional[bool] = False, + grid_kwargs: Optional[Dict[str, Any]] = {}, **kwargs: Dict[str, Any]) -> UxDataset: """Wraps ``xarray.open_mfdataset()``, given a single grid topology file with multiple dataset paths with corresponding data. @@ -218,7 +227,13 @@ def open_mfdataset(grid_filename_or_obj: str, Path or URL to the source grid file. For diagnostic/reporting purposes only. 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 + + grid_kwargs : Dict[str, Any], optional + Additional arguments passed on to ``xarray.open_dataset`` when opening up a Grid File. Refer to the + [xarray + docs](https://xarray.pydata.org/en/stable/generated/xarray.open_dataset.html) + for accepted keyword arguments. **kwargs : Dict[str, Any] Additional arguments passed on to ``xarray.open_mfdataset``. Refer to the @@ -248,17 +263,17 @@ def open_mfdataset(grid_filename_or_obj: str, >>> ux_ds = ux.open_mfdataset("grid_filename.g", "grid_filename_vortex_*.nc") """ - ## Grid definition + # Grid definition uxgrid = open_grid(grid_filename_or_obj, gridspec=gridspec, vertices=vertices, islatlon=islatlon, isconcave=isconcave, - use_dual=use_dual) + use_dual=use_dual, + **grid_kwargs) - ## UxDataset + # UxDataset ds = xr.open_mfdataset(paths, decode_times=False, **kwargs) # type: ignore - uxds = UxDataset(ds, uxgrid=uxgrid, source_datasets=str(paths)) return uxds