Merge pull request #390 from UXARRAY/philipc2/grid-kwargs

Separate `kwargs` for Opening a Grid File in Open Dataset

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"})

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