Merge pull request #1282 from UXARRAY/rajeeja/fix_to_netcdf_ugrid

Deprecate Grid.encode_as, Add esmf writer, Add tests, fix to_netcdf bug after healpix mesh creation

Author: rajeeja

test/test_esmf.py

@@ -2,11 +2,14 @@
 import os
 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"
 esmf_ne30_data_path = current_path / 'meshfiles' / "esmf" / "ne30" / "ne30pg3.data.nc"
+gridfile_ne30 = current_path / 'meshfiles' / "ugrid" / "outCSne30" / "outCSne30.ug"
 
 def test_read_esmf():
     """Tests the reading of an ESMF grid file and its encoding into the UGRID
@@ -35,3 +38,72 @@ def test_read_esmf_dataset():
 
     for dim in dims:
         assert dim in uxds.dims
+
+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

@@ -33,9 +33,30 @@ def test_mixed_exodus():
     """Read/write an exodus file with two types of faces (triangle and quadrilaterals) and writes a ugrid file."""
     uxgrid = ux.open_grid(exo2_filename)
 
-    uxgrid.encode_as("UGRID")
-    uxgrid.encode_as("Exodus")
-    # Add assertions or checks as needed
+    ugrid_obj = uxgrid.to_xarray("UGRID")
+    exo_obj = uxgrid.to_xarray("Exodus")
+
+    ugrid_obj.to_netcdf("test_ugrid.nc")
+    exo_obj.to_netcdf("test_exo.exo")
+
+    ugrid_load_saved = ux.open_grid("test_ugrid.nc")
+    exodus_load_saved = ux.open_grid("test_exo.exo")
+
+    # Face node connectivity comparison
+    assert np.array_equal(ugrid_load_saved.face_node_connectivity.values, uxgrid.face_node_connectivity.values)
+    assert np.array_equal(uxgrid.face_node_connectivity.values, exodus_load_saved.face_node_connectivity.values)
+
+    # Node coordinates comparison
+    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
+    ugrid_load_saved._ds.close()
+    exodus_load_saved._ds.close()
+    del ugrid_load_saved, exodus_load_saved
+    os.remove("test_ugrid.nc")
+    os.remove("test_exo.exo")
 
 def test_standardized_dtype_and_fill():
     """Test to see if Mesh2_Face_Nodes uses the expected integer datatype and expected fill value as set in constants.py."""

test/test_grid.py

@@ -96,15 +96,115 @@ def test_grid_with_holes():
     assert grid_without_holes.global_sphere_coverage
 
 
-def test_grid_encode_as():
-    """Reads a ugrid file and encodes it as `xarray.Dataset` in various types."""
-    grid_CSne30.encode_as("UGRID")
-    grid_RLL1deg.encode_as("UGRID")
-    grid_RLL10deg_CSne4.encode_as("UGRID")
+def test_grid_ugrid_exodus_roundtrip():
+    """Test round-trip serialization of grid objects through UGRID and Exodus xarray formats.
 
-    grid_CSne30.encode_as("Exodus")
-    grid_RLL1deg.encode_as("Exodus")
-    grid_RLL10deg_CSne4.encode_as("Exodus")
+    Validates that 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.
+
+    The test verifies:
+    - Successful conversion to UGRID and Exodus xarray formats
+    - 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
+    """
+
+    # Convert grids to xarray.Dataset objects in different formats
+    ugrid_datasets = {
+        'CSne30': grid_CSne30.to_xarray("UGRID"),
+        'RLL1deg': grid_RLL1deg.to_xarray("UGRID"),
+        'RLL10deg_CSne4': grid_RLL10deg_CSne4.to_xarray("UGRID")
+    }
+
+    exodus_datasets = {
+        'CSne30': grid_CSne30.to_xarray("Exodus"),
+        'RLL1deg': grid_RLL1deg.to_xarray("Exodus"),
+        'RLL10deg_CSne4': grid_RLL10deg_CSne4.to_xarray("Exodus")
+    }
+
+    # Define test cases with corresponding grid objects
+    test_grids = {
+        'CSne30': grid_CSne30,
+        'RLL1deg': grid_RLL1deg,
+        'RLL10deg_CSne4': grid_RLL10deg_CSne4
+    }
+
+    # Perform round-trip validation for each grid type
+    test_files = []
+
+    for grid_name in test_grids.keys():
+        ugrid_dataset = ugrid_datasets[grid_name]
+        exodus_dataset = exodus_datasets[grid_name]
+        original_grid = test_grids[grid_name]
+
+        # 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)
+        exodus_dataset.to_netcdf(exodus_filepath)
+
+        # 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=f"UGRID face connectivity mismatch for {grid_name}"
+        )
+        np.testing.assert_array_equal(
+            original_grid.face_node_connectivity.values,
+            reloaded_exodus.face_node_connectivity.values,
+            err_msg=f"Exodus face connectivity mismatch for {grid_name}"
+        )
+
+        # 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=f"UGRID longitude mismatch for {grid_name}",
+            rtol=ERROR_TOLERANCE
+        )
+        np.testing.assert_allclose(
+            original_grid.node_lon.values,
+            reloaded_exodus.node_lon.values,
+            err_msg=f"Exodus longitude mismatch for {grid_name}",
+            rtol=ERROR_TOLERANCE
+        )
+        np.testing.assert_allclose(
+            original_grid.node_lat.values,
+            reloaded_ugrid.node_lat.values,
+            err_msg=f"UGRID latitude mismatch for {grid_name}",
+            rtol=ERROR_TOLERANCE
+        )
+        np.testing.assert_allclose(
+            original_grid.node_lat.values,
+            reloaded_exodus.node_lat.values,
+            err_msg=f"Exodus latitude mismatch for {grid_name}",
+            rtol=ERROR_TOLERANCE
+        )
+
+    # This might be need for windows "ermissionError: [WinError 32] -- file accessed by another process"
+    reloaded_exodus._ds.close()
+    reloaded_ugrid._ds.close()
+    del reloaded_exodus
+    del reloaded_ugrid
+
+    # Clean up temporary test files
+    for filepath in test_files:
+        if os.path.exists(filepath):
+            os.remove(filepath)
 
 
 def test_grid_init_verts():
@@ -147,21 +247,21 @@ def test_grid_init_verts():
 
     assert vgrid.n_face == 6
     assert vgrid.n_node == 8
-    vgrid.encode_as("UGRID")
+    vgrid.to_xarray("UGRID")
 
     faces_verts_one = np.array([
         np.array([[150, 10], [160, 20], [150, 30], [135, 30], [125, 20], [135, 10]])
     ])
     vgrid = ux.open_grid(faces_verts_one, latlon=True)
     assert vgrid.n_face == 1
     assert vgrid.n_node == 6
-    vgrid.encode_as("UGRID")
+    vgrid.to_xarray("UGRID")
 
     faces_verts_single_face = np.array([[150, 10], [160, 20], [150, 30], [135, 30], [125, 20], [135, 10]])
     vgrid = ux.open_grid(faces_verts_single_face, latlon=True)
     assert vgrid.n_face == 1
     assert vgrid.n_node == 6
-    vgrid.encode_as("UGRID")
+    vgrid.to_xarray("UGRID")
 
 
 def test_grid_init_verts_different_input_datatype():
@@ -174,7 +274,7 @@ def test_grid_init_verts_different_input_datatype():
     vgrid = ux.open_grid(faces_verts_ndarray, latlon=True)
     assert vgrid.n_face == 3
     assert vgrid.n_node == 14
-    vgrid.encode_as("UGRID")
+    vgrid.to_xarray("UGRID")
 
     faces_verts_list = [[[150, 10], [160, 20], [150, 30], [135, 30], [125, 20], [135, 10]],
                         [[125, 20], [135, 30], [125, 60], [110, 60], [100, 30], [105, 20]],
@@ -183,7 +283,7 @@ def test_grid_init_verts_different_input_datatype():
     assert vgrid.n_face == 3
     assert vgrid.n_node == 14
     assert vgrid.validate()
-    vgrid.encode_as("UGRID")
+    vgrid.to_xarray("UGRID")
 
     faces_verts_tuples = [
         ((150, 10), (160, 20), (150, 30), (135, 30), (125, 20), (135, 10)),
@@ -194,7 +294,7 @@ def test_grid_init_verts_different_input_datatype():
     assert vgrid.n_face == 3
     assert vgrid.n_node == 14
     assert vgrid.validate()
-    vgrid.encode_as("UGRID")
+    vgrid.to_xarray("UGRID")
 
 
 def test_grid_init_verts_fill_values():

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__)))
@@ -99,3 +100,89 @@ def test_invalid_cells():
     xrda = xr.DataArray(data=np.ones(11), dims=['cell']).to_dataset(name='cell')
     with pytest.raises(ValueError):
         uxda = ux.UxDataset.from_healpix(xrda)
+
+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
+    original_grid = ux.Grid.from_healpix(zoom=3)
+
+    # Access node coordinates to ensure they're generated before encoding
+    _ = 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