o When handling different element types - Fix the Exodus reader. Add num_elem to Exodus writer, addresses #1242. o Modify tests - they are far from nice good, needs a complete overhaul of all io module

Author: rajeeja

test/test_esmf.py

@@ -3,7 +3,8 @@
 from pathlib import Path
 import pytest
 import xarray as xr
-
+import numpy as np
+from uxarray.constants import ERROR_TOLERANCE
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
 esmf_ne30_grid_path = current_path / 'meshfiles' / "esmf" / "ne30" / "ne30pg3.grid.nc"
@@ -38,10 +39,71 @@ def test_read_esmf_dataset():
     for dim in dims:
         assert dim in uxds.dims
 
-def test_write_esmf():
-    """Tests the writing of a UxDataset to an ESMF Grid file."""
-    uxds = ux.open_grid(gridfile_ne30)
-    out_ds = uxds.to_xarray("ESMF")
-    assert isinstance(out_ds, xr.Dataset)
-    assert 'nodeCoords' in out_ds
-    assert 'elementConn' in out_ds
+def test_esmf_round_trip_consistency():
+    """Test round-trip serialization of grid objects through ESMF xarray format.
+
+    Validates that grid objects can be successfully converted to ESMF xarray.Dataset
+    format, serialized to disk, and reloaded while maintaining numerical accuracy
+    and topological integrity.
+
+    The test verifies:
+    - Successful conversion to ESMF xarray format
+    - File I/O round-trip consistency
+    - Preservation of face-node connectivity (exact)
+    - Preservation of node coordinates (within numerical tolerance)
+
+    Raises:
+        AssertionError: If any round-trip validation fails
+    """
+    # Load original grid
+    original_grid = ux.open_grid(gridfile_ne30)
+
+    # Convert to ESMF xarray format
+    esmf_dataset = original_grid.to_xarray("ESMF")
+
+    # Verify dataset structure
+    assert isinstance(esmf_dataset, xr.Dataset)
+    assert 'nodeCoords' in esmf_dataset
+    assert 'elementConn' in esmf_dataset
+
+    # Define output file path
+    esmf_filepath = "test_esmf_ne30.nc"
+
+    # Remove existing test file to ensure clean state
+    if os.path.exists(esmf_filepath):
+        os.remove(esmf_filepath)
+
+    try:
+        # Serialize dataset to disk
+        esmf_dataset.to_netcdf(esmf_filepath)
+
+        # Reload grid from serialized file
+        reloaded_grid = ux.open_grid(esmf_filepath)
+
+        # Validate topological consistency (face-node connectivity)
+        # Integer connectivity arrays must be exactly preserved
+        np.testing.assert_array_equal(
+            original_grid.face_node_connectivity.values,
+            reloaded_grid.face_node_connectivity.values,
+            err_msg="ESMF face connectivity mismatch"
+        )
+
+        # Validate coordinate consistency with numerical tolerance
+        # Coordinate transformations and I/O precision may introduce minor differences
+        np.testing.assert_allclose(
+            original_grid.node_lon.values,
+            reloaded_grid.node_lon.values,
+            err_msg="ESMF longitude mismatch",
+            rtol=ERROR_TOLERANCE
+        )
+        np.testing.assert_allclose(
+            original_grid.node_lat.values,
+            reloaded_grid.node_lat.values,
+            err_msg="ESMF latitude mismatch",
+            rtol=ERROR_TOLERANCE
+        )
+
+    finally:
+        # Clean up temporary test file
+        if os.path.exists(esmf_filepath):
+            os.remove(esmf_filepath)

test/test_exodus.py

@@ -50,6 +50,7 @@ def test_mixed_exodus():
     assert np.array_equal(ugrid_load_saved.node_lon.values, uxgrid.node_lon.values)
     assert np.array_equal(uxgrid.node_lon.values, exodus_load_saved.node_lon.values)
     assert np.array_equal(ugrid_load_saved.node_lat.values, uxgrid.node_lat.values)
+
     # Cleanup
     os.remove("test_ugrid.nc")
     os.remove("test_exo.exo")

test/test_grid.py

@@ -144,6 +144,8 @@ def test_grid_ugrid_exodus_roundtrip():
         # Define output file paths
         ugrid_filepath = f"test_ugrid_{grid_name}.nc"
         exodus_filepath = f"test_exodus_{grid_name}.exo"
+        test_files.append(ugrid_filepath)
+        test_files.append(exodus_filepath)
 
         # Serialize datasets to disk
         ugrid_dataset.to_netcdf(ugrid_filepath)

test/test_healpix.py

@@ -6,6 +6,7 @@
 import xarray as xr
 import pandas as pd
 from pathlib import Path
+from uxarray.constants import ERROR_TOLERANCE
 
 
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
@@ -100,35 +101,88 @@ def test_invalid_cells():
     with pytest.raises(ValueError):
         uxda = ux.UxDataset.from_healpix(xrda)
 
-def test_healpix_to_netcdf(tmp_path):
-    """Test that HEALPix grid can be encoded as UGRID and saved to netCDF.
-       Using pytest tmp_path fixture to create a temporary file.
+def test_healpix_round_trip_consistency(tmp_path):
+    """Test round-trip serialization of HEALPix grid through UGRID and Exodus formats.
+
+    Validates that HEALPix grid objects can be successfully converted to xarray.Dataset
+    objects in both UGRID and Exodus formats, serialized to disk, and reloaded
+    while maintaining numerical accuracy and topological integrity.
+
+    Args:
+        tmp_path: pytest fixture providing temporary directory
+
+    Raises:
+        AssertionError: If any round-trip validation fails
     """
     # Create HEALPix grid
-    h = ux.Grid.from_healpix(zoom=3)
+    original_grid = ux.Grid.from_healpix(zoom=3)
 
     # Access node coordinates to ensure they're generated before encoding
-    _ = h.node_lon
-    _ = h.node_lat
-
-    # Convert to different formats
-    uxa_ugrid = h.to_xarray("UGRID")
-    uxa_exodus = h.to_xarray("Exodus")
-
-    tmp_filename_ugrid = tmp_path / "healpix_test_ugrid.nc"
-    tmp_filename_exodus = tmp_path / "healpix_test_exodus.exo"
-
-    # Save to netCDF
-    uxa_ugrid.to_netcdf(tmp_filename_ugrid)
-    uxa_exodus.to_netcdf(tmp_filename_exodus)
-
-    # Assertions
-    assert tmp_filename_ugrid.exists()
-    assert tmp_filename_ugrid.stat().st_size > 0
-    assert tmp_filename_exodus.exists()
-    assert tmp_filename_exodus.stat().st_size > 0
-
-    loaded_grid_ugrid = ux.open_grid(tmp_filename_ugrid)
-    loaded_grid_exodus = ux.open_grid(tmp_filename_exodus)
-    assert loaded_grid_ugrid.n_face == h.n_face
-    assert loaded_grid_exodus.n_face == h.n_face
+    _ = original_grid.node_lon
+    _ = original_grid.node_lat
+
+    # Convert to xarray.Dataset objects in different formats
+    ugrid_dataset = original_grid.to_xarray("UGRID")
+    exodus_dataset = original_grid.to_xarray("Exodus")
+
+    # Define output file paths using tmp_path fixture
+    ugrid_filepath = tmp_path / "healpix_test_ugrid.nc"
+    exodus_filepath = tmp_path / "healpix_test_exodus.exo"
+
+    # Serialize datasets to disk
+    ugrid_dataset.to_netcdf(ugrid_filepath)
+    exodus_dataset.to_netcdf(exodus_filepath)
+
+    # Verify files were created successfully
+    assert ugrid_filepath.exists()
+    assert ugrid_filepath.stat().st_size > 0
+    assert exodus_filepath.exists()
+    assert exodus_filepath.stat().st_size > 0
+
+    # Reload grids from serialized files
+    reloaded_ugrid = ux.open_grid(ugrid_filepath)
+    reloaded_exodus = ux.open_grid(exodus_filepath)
+
+    # Validate topological consistency (face-node connectivity)
+    # Integer connectivity arrays must be exactly preserved
+    np.testing.assert_array_equal(
+        original_grid.face_node_connectivity.values,
+        reloaded_ugrid.face_node_connectivity.values,
+        err_msg="UGRID face connectivity mismatch for HEALPix"
+    )
+    np.testing.assert_array_equal(
+        original_grid.face_node_connectivity.values,
+        reloaded_exodus.face_node_connectivity.values,
+        err_msg="Exodus face connectivity mismatch for HEALPix"
+    )
+
+    # Validate coordinate consistency with numerical tolerance
+    # Coordinate transformations and I/O precision may introduce minor differences
+    np.testing.assert_allclose(
+        original_grid.node_lon.values,
+        reloaded_ugrid.node_lon.values,
+        err_msg="UGRID longitude mismatch for HEALPix",
+        rtol=ERROR_TOLERANCE
+    )
+    np.testing.assert_allclose(
+        original_grid.node_lon.values,
+        reloaded_exodus.node_lon.values,
+        err_msg="Exodus longitude mismatch for HEALPix",
+        rtol=ERROR_TOLERANCE
+    )
+    np.testing.assert_allclose(
+        original_grid.node_lat.values,
+        reloaded_ugrid.node_lat.values,
+        err_msg="UGRID latitude mismatch for HEALPix",
+        rtol=ERROR_TOLERANCE
+    )
+    np.testing.assert_allclose(
+        original_grid.node_lat.values,
+        reloaded_exodus.node_lat.values,
+        err_msg="Exodus latitude mismatch for HEALPix",
+        rtol=ERROR_TOLERANCE
+    )
+
+    # Validate grid dimensions are preserved
+    assert reloaded_ugrid.n_face == original_grid.n_face
+    assert reloaded_exodus.n_face == original_grid.n_face

test/test_scrip.py

@@ -20,6 +20,7 @@
 gridfile_RLL1deg = current_path / "meshfiles" / "ugrid" / "outRLL1deg" / "outRLL1deg.ug"
 gridfile_RLL10deg_ne4 = current_path / "meshfiles" / "ugrid" / "ov_RLL10deg_CSne4" / "ov_RLL10deg_CSne4.ug"
 gridfile_exo_ne8 = current_path / "meshfiles" / "exodus" / "outCSne8" / "outCSne8.g"
+gridfile_scrip = current_path / "meshfiles" / "scrip" / "outCSne8" /"outCSne8.nc"
 
 def test_read_ugrid():
     """Reads a ugrid file."""
@@ -47,8 +48,20 @@ def test_read_ugrid():
 
 def test_to_xarray_ugrid():
     """Read an Exodus dataset and convert it to UGRID format using to_xarray."""
-    ux_grid = ux.open_grid(gridfile_exo_ne8)
-    ux_grid.to_xarray("UGRID")
+    ux_grid = ux.open_grid(gridfile_scrip)
+    xr_obj = ux_grid.to_xarray("UGRID")
+    xr_obj.to_netcdf("scrip_ugrid_csne8.nc")
+    reloaded_grid = ux.open_grid("scrip_ugrid_csne8.nc")
+    # Check that the grid topology is perfectly preserved
+    nt.assert_array_equal(ux_grid.face_node_connectivity.values,
+                          reloaded_grid.face_node_connectivity.values)
+
+    # Check that node coordinates are numerically close
+    nt.assert_allclose(ux_grid.node_lon.values, reloaded_grid.node_lon.values)
+    nt.assert_allclose(ux_grid.node_lat.values, reloaded_grid.node_lat.values)
+
+    # Cleanup
+    os.remove("scrip_ugrid_csne8.nc")
 
 def test_standardized_dtype_and_fill():
     """Test to see if Mesh2_Face_Nodes uses the expected integer datatype

test/test_ugrid.py

@@ -46,7 +46,19 @@ def test_read_ugrid():
 def test_to_xarray_ugrid():
     """Read an Exodus dataset and convert it to UGRID format using to_xarray."""
     ux_grid = ux.open_grid(gridfile_exo_ne8)
-    ux_grid.to_xarray("UGRID")
+    xr_obj = ux_grid.to_xarray("UGRID")
+    xr_obj.to_netcdf("ugrid_exo_csne8.nc")
+    reloaded_grid = ux.open_grid("ugrid_exo_csne8.nc")
+    # Check that the grid topology is perfectly preserved
+    nt.assert_array_equal(ux_grid.face_node_connectivity.values,
+                          reloaded_grid.face_node_connectivity.values)
+
+    # Check that node coordinates are numerically close
+    nt.assert_allclose(ux_grid.node_lon.values, reloaded_grid.node_lon.values)
+    nt.assert_allclose(ux_grid.node_lat.values, reloaded_grid.node_lat.values)
+
+    # Cleanup
+    os.remove("ugrid_exo_csne8.nc")
 
 def test_standardized_dtype_and_fill():
     """Test to see if Mesh2_Face_Nodes uses the expected integer datatype and expected fill value."""

uxarray/io/_exodus.py

@@ -64,18 +64,36 @@ def _read_exodus(ext_ds):
                     data=ext_ds.coordz, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_Z_ATTRS
                 )
         elif "connect" in key:
-            # check if num face nodes is less than max.
-            if value.data.shape[1] <= max_face_nodes:
-                face_nodes = value
-            else:
-                raise RuntimeError("found face_nodes_dim greater than n_max_face_nodes")
+            # This variable will be populated in the next step
+            pass
 
     # outside the k,v for loop
     # set the face nodes data compiled in "connect" section
+    connect_list = []
+    for key, value in ext_ds.variables.items():
+        if "connect" in key:
+            connect_list.append(value.data)
+
+    padded_blocks = []
+    for block in connect_list:
+        num_nodes = block.shape[1]
+        pad_width = max_face_nodes - num_nodes
+
+        # Pad with 0, as Exodus uses 0 for non-existent nodes
+        padded_block = np.pad(
+            block, ((0, 0), (0, pad_width)), "constant", constant_values=0
+        )
+        padded_blocks.append(padded_block)
+
+    # Prevent error on empty grids
+    if not padded_blocks:
+        face_nodes = np.empty((0, max_face_nodes), dtype=INT_DTYPE)
+    else:
+        face_nodes = np.vstack(padded_blocks)
 
     # standardize fill values and data type face nodes
     face_nodes = _replace_fill_values(
-        grid_var=face_nodes[:] - 1,
+        grid_var=xr.DataArray(face_nodes - 1),  # Wrap numpy array in a DataArray
         original_fill=-1,
         new_fill=INT_FILL_VALUE,
         new_dtype=INT_DTYPE,
@@ -151,6 +169,9 @@ def _encode_exodus(ds, outfile=None):
 
     exo_ds["time_whole"] = xr.DataArray(data=[], dims=["time_step"])
 
+    # --- Add num_elem dimension ---
+    exo_ds.attrs["num_elem"] = ds.sizes["n_face"]
+
     # --- QA Records ---
     ux_exodus_version = "1.0"
     qa_records = [["uxarray"], [ux_exodus_version], [date], [time]]
@@ -161,7 +182,6 @@ def _encode_exodus(ds, outfile=None):
 
     # --- Node Coordinates ---
     if "node_x" not in ds:
-        print("HERE", ds["node_lon"].values)
         node_lon_rad = np.deg2rad(ds["node_lon"].values)
         node_lat_rad = np.deg2rad(ds["node_lat"].values)
         x, y, z = _lonlat_rad_to_xyz(node_lon_rad, node_lat_rad)