Merge branch 'main' into philipc2/2023-09-1

UXARRAY · Sep 22, 2023 · d846c4a · d846c4a
2 parents 419f81f + f812cf5
commit d846c4a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/docs/user_api/index.rst b/docs/user_api/index.rst
@@ -71,6 +71,7 @@ IO
 .. autosummary::
    :toctree: _autosummary
 
+   UxDataArray.to_dataset
    UxDataArray.to_geodataframe
    UxDataArray.to_polycollection
 

diff --git a/test/test_dataarray.py b/test/test_dataarray.py
@@ -7,6 +7,8 @@
 
 from uxarray.grid.geometry import _build_polygon_shells, _build_corrected_polygon_shells
 
+from uxarray.core.dataset import UxDataset
+
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
 gridfile_ne30 = current_path / "meshfiles" / "ugrid" / "outCSne30" / "outCSne30.ug"
@@ -19,6 +21,17 @@
 dsfile_v1_geoflow = current_path / "meshfiles" / "ugrid" / "geoflow-small" / "v1.nc"
 
 
+class TestDataArray(TestCase):
+
+    def test_to_dataset(self):
+        """Tests the conversion of UxDataArrays to a UXDataset."""
+        uxds = ux.open_dataset(gridfile_ne30, dsfile_var2_ne30)
+        uxds_converted = uxds['psi'].to_dataset()
+
+        assert isinstance(uxds_converted, UxDataset)
+        assert uxds_converted.uxgrid == uxds.uxgrid
+
+
 class TestGeometryConversions(TestCase):
 
     def test_to_geodataframe(self):

diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py
@@ -3,9 +3,14 @@
 import xarray as xr
 import numpy as np
 
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 from uxarray.grid import Grid
+import uxarray.core.dataset
+
+if TYPE_CHECKING:
+    from uxarray.core.dataarray import UxDataArray
+    from uxarray.core.dataset import UxDataset
 
 
 class UxDataArray(xr.DataArray):
@@ -99,6 +104,11 @@ def uxgrid(self):
     def uxgrid(self, ugrid_obj):
         self._uxgrid = ugrid_obj
 
+    def to_dataset(self) -> UxDataset:
+        """Convert a UxDataArray to a UxDataset."""
+        xrds = super().to_dataset()
+        return uxarray.core.dataset.UxDataset(xrds, uxgrid=self.uxgrid)
+
     def to_geodataframe(self,
                         override=False,
                         cache=True,