Correctly check dimensions with HEALPix (#1278)

* update n_node property and add extra check to dimension renaming

* update n_node dimension

* add test case

Author: philipc2

test/test_healpix.py

@@ -4,6 +4,7 @@
 import pytest
 import os
 import xarray as xr
+import pandas as pd
 from pathlib import Path
 
 
@@ -38,6 +39,17 @@ def test_boundaries(resolution_level):
     # check for the correct number of boundary nodes
     assert (uxgrid.n_node == uxgrid.n_face + 2)
 
+@pytest.mark.parametrize("pixels_only", [True, False])
+def test_time_dimension_roundtrip(pixels_only):
+    ds = xr.open_dataset(ds_path)
+    dummy_time = pd.to_datetime(["2025-01-01T00:00:00"])
+    ds_time = ds.expand_dims(time=dummy_time)
+    uxds = ux.UxDataset.from_healpix(ds_time, pixels_only=pixels_only)
+
+    # Ensure time dimension is preserved and that the conversion worked
+    assert "time" in uxds.dims
+    assert uxds.sizes["time"] == 1
+
 def test_dataset():
     uxds = ux.UxDataset.from_healpix(ds_path)
 

uxarray/core/utils.py

@@ -44,21 +44,17 @@ def _map_dims_to_ugrid(
             # drop dimensions not present in the original dataset
             _source_dims_dict.pop(key)
 
-        # only check edge dimension if it is present (to avoid overhead of computing connectivity)
-        if "n_edge" in grid._ds.dims:
-            n_edge = grid._ds.sizes["n_edge"]
-        else:
-            n_edge = None
-
-        for dim in set(ds.dims) ^ _source_dims_dict.keys():
-            # obtain dimensions that were not parsed source_dims_dict and attempt to match to a grid element
-            if ds.sizes[dim] == grid.n_face:
-                _source_dims_dict[dim] = "n_face"
-            elif ds.sizes[dim] == grid.n_node:
-                _source_dims_dict[dim] = "n_node"
-            elif n_edge is not None:
-                if ds.sizes[dim] == n_edge:
-                    _source_dims_dict[dim] = "n_edge"
+        # build a reverse map
+        size_to_name = {
+            grid._ds.sizes[name]: name
+            for name in ("n_face", "n_node", "n_edge")
+            if name in grid._ds.dims
+        }
+
+        for dim in set(ds.dims) - _source_dims_dict.keys():
+            name = size_to_name.get(ds.sizes[dim])
+            if name:
+                _source_dims_dict[dim] = name
 
         # rename dimensions to follow the UGRID conventions
         ds = ds.swap_dims(_source_dims_dict)

uxarray/grid/grid.py

@@ -817,6 +817,8 @@ def attrs(self) -> dict:
     @property
     def n_node(self) -> int:
         """Total number of nodes."""
+        if "face_node_connectivity" not in self._ds:
+            _ = self.face_node_connectivity
         return self._ds.sizes["n_node"]
 
     @property