diff --git a/test/io/test_esmf.py b/test/io/test_esmf.py index fbf910e1c..2c5104dba 100644 --- a/test/io/test_esmf.py +++ b/test/io/test_esmf.py @@ -1,9 +1,5 @@ import uxarray as ux -import os -import pytest import xarray as xr -import numpy as np -from uxarray.constants import ERROR_TOLERANCE def test_read_esmf(gridpath): @@ -35,71 +31,23 @@ def test_read_esmf_dataset(gridpath, datasetpath): for dim in dims: assert dim in uxds.dims -def test_esmf_round_trip_consistency(gridpath): - """Test round-trip serialization of grid objects through ESMF xarray format. +def test_encode_esmf_structure(gridpath): + """Encoding to ESMF produces the variables the format requires. - 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 + Round-trip fidelity is covered for every writable format by + ``TestIOWriteRoundTrip`` in test_io_common.py; this only pins down the + ESMF-specific layout. """ - # Load original grid - original_grid = ux.open_grid(gridpath("ugrid", "outCSne30", "outCSne30.ug")) - - # Convert to ESMF xarray format - esmf_dataset = original_grid.to_xarray("ESMF") + uxgrid = ux.open_grid(gridpath("ugrid", "outCSne30", "outCSne30.ug")) + esmf_dataset = uxgrid.to_xarray("ESMF") - # Verify dataset structure assert isinstance(esmf_dataset, xr.Dataset) assert 'nodeCoords' in esmf_dataset assert 'elementConn' in esmf_dataset + assert 'numElementConn' 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) + # elementConn is 1-based with -1 marking unused slots + assert esmf_dataset['elementConn'].attrs['_FillValue'] == -1 + assert esmf_dataset['numElementConn'].values.sum() == ( + uxgrid.n_nodes_per_face.values.sum() + ) diff --git a/test/io/test_exodus.py b/test/io/test_exodus.py index ea37f5010..bae63c9f8 100644 --- a/test/io/test_exodus.py +++ b/test/io/test_exodus.py @@ -21,8 +21,47 @@ def test_init_verts(): def test_encode_exodus(gridpath): """Read a UGRID dataset and encode that as an Exodus format.""" uxgrid = ux.open_grid(gridpath("exodus", "outCSne8", "outCSne8.g")) - # Add encoding logic and assertions as needed - pass # Placeholder for actual implementation + exo_ds = uxgrid.to_xarray("Exodus") + + # A uniform quad mesh belongs in exactly one block, typed for a quad + blocks = [v for v in exo_ds.data_vars if v.startswith("connect")] + assert blocks == ["connect1"] + assert exo_ds["connect1"].attrs["elem_type"] == "SHELL4" + assert exo_ds["connect1"].shape == (uxgrid.n_face, 4) + +def test_encode_exodus_mixed_blocks(): + """Faces of different sizes go into separate, correctly typed blocks. + + Exodus element blocks are homogeneous, so a mixed mesh has to be split by + face size. Getting the fill value wrong collapses everything into one + max-width block and writes the padding out as a node index. + """ + face_node_connectivity = np.array([ + [0, 1, 2, 3], + [1, 4, 2, INT_FILL_VALUE], + [0, 3, 4, INT_FILL_VALUE], + ]) + uxgrid = ux.Grid.from_topology( + node_lon=np.array([0.0, 10.0, 10.0, 0.0, 20.0]), + node_lat=np.array([0.0, 0.0, 10.0, 10.0, 0.0]), + face_node_connectivity=face_node_connectivity, + fill_value=INT_FILL_VALUE, + ) + + exo_ds = uxgrid.to_xarray("Exodus") + + blocks = sorted(v for v in exo_ds.data_vars if v.startswith("connect")) + assert blocks == ["connect1", "connect2"] + + by_type = {exo_ds[b].attrs["elem_type"]: exo_ds[b] for b in blocks} + assert set(by_type) == {"TRI", "SHELL4"} + assert by_type["TRI"].shape == (2, 3) + assert by_type["SHELL4"].shape == (1, 4) + + # Blocks are written grouped by type, so the original face order has to be + # recorded or face-centered data silently misaligns on the way back in. + assert "elem_num_map" in exo_ds + assert sorted(exo_ds["elem_num_map"].values.tolist()) == [1, 2, 3] def test_mixed_exodus(gridpath): """Read/write an exodus file with two types of faces (triangle and quadrilaterals) and writes a ugrid file.""" diff --git a/test/io/test_io_common.py b/test/io/test_io_common.py index 8e72409a9..10220cff1 100644 --- a/test/io/test_io_common.py +++ b/test/io/test_io_common.py @@ -35,6 +35,11 @@ # Formats that support writing WRITABLE_FORMATS = ["ugrid", "exodus", "scrip", "esmf"] +# SCRIP stores corner coordinates rather than node indices, so its reader +# rebuilds nodes by deduplicating coordinates and renumbers them in the process. +# Its geometry survives a round trip; its connectivity is not restored verbatim. +EXACT_CONNECTIVITY_FORMATS = ["ugrid", "exodus", "esmf"] + # Format conversion test pairs - removed for now as format conversion # requires more sophisticated handling than simple to_netcdf @@ -72,6 +77,69 @@ def grid_from_format(request, test_data_dir): return grid +# File suffix to write each format under. Exodus is sniffed by extension. +FORMAT_SUFFIX = {"ugrid": ".nc", "exodus": ".exo", "scrip": ".nc", "esmf": ".nc"} + +# Formats that write node indices rather than coordinates, as +# {format: (variable name prefix, on-disk fill value or None)}. Exodus splits +# connectivity across one exactly-sized connect per element block, so it has +# no padding to skip; ESMF writes a single padded array. +ENCODED_INDEX_VARS = {"exodus": ("connect", None), "esmf": ("elementConn", -1)} + +RAGGED_FACE_NODES = np.array( + [ + [0, 1, 2, 3], # quad + [1, 4, 2, INT_FILL_VALUE], # triangle + [0, 3, 4, INT_FILL_VALUE], # triangle + ] +) +RAGGED_NODE_LON = np.array([0.0, 10.0, 10.0, 0.0, 20.0]) +RAGGED_NODE_LAT = np.array([0.0, 0.0, 10.0, 10.0, 0.0]) + + +@pytest.fixture +def ragged_grid(): + """A grid whose faces are not all the same size. + + Every uniform grid pads nothing, so the fill-value paths in the encoders are + only reachable with mixed face sizes. + """ + return ux.Grid.from_topology( + node_lon=RAGGED_NODE_LON, + node_lat=RAGGED_NODE_LAT, + face_node_connectivity=RAGGED_FACE_NODES, + fill_value=INT_FILL_VALUE, + ) + + +def _write_and_reload(grid, fmt, directory): + """Encode ``grid`` as ``fmt``, write it out, and read it back.""" + path = directory / f"round_trip_{fmt}{FORMAT_SUFFIX[fmt]}" + grid.to_xarray(fmt).to_netcdf(path) + return ux.open_grid(path) + + +def _face_geometry(grid): + """Describe each face by its corner coordinates instead of node indices. + + Lets formats that renumber nodes, or that pad a short face by repeating a + vertex, be compared against the grid they were written from. + """ + conn = grid.face_node_connectivity.values + lon = grid.node_lon.values + lat = grid.node_lat.values + + faces = [] + for row in conn: + corners = { + (round(float(lon[i]), 6), round(float(lat[i]), 6)) + for i in row + if i != INT_FILL_VALUE + } + faces.append(tuple(sorted(corners))) + return sorted(faces) + + class TestIOCommon: """Common IO tests across all formats. Helps catch format-specific regressions early and keep behavior consistent. @@ -139,3 +207,108 @@ def test_standardized_dtype_and_fill(self, grid_from_format): # Check that face_node_connectivity uses the standardized fill value assert grid.face_node_connectivity._FillValue == INT_FILL_VALUE + + +class TestIOWriteRoundTrip: + """Write each format back out and read it in again. + + The encoders historically broke on padded connectivity: a fill value that + gets offset, narrowed to a smaller dtype, or written out as a coordinate + comes back as a real vertex. Nothing raises when that happens -- the mesh + just quietly gains nodes -- so these tests assert on the reloaded topology + rather than on the write succeeding. + """ + + @pytest.mark.parametrize("fmt", EXACT_CONNECTIVITY_FORMATS) + def test_uniform_grid_round_trip(self, fmt, gridpath, tmp_path): + """A grid with uniform face sizes survives a write/read cycle intact.""" + original = ux.open_grid(gridpath("ugrid", "outCSne30", "outCSne30.ug")) + reloaded = _write_and_reload(original, fmt, tmp_path) + + assert_array_equal( + original.face_node_connectivity.values, + reloaded.face_node_connectivity.values, + err_msg=f"{fmt}: face connectivity changed across a round trip", + ) + assert_allclose( + original.node_lon.values, reloaded.node_lon.values, rtol=ERROR_TOLERANCE + ) + assert_allclose( + original.node_lat.values, reloaded.node_lat.values, rtol=ERROR_TOLERANCE + ) + + @pytest.mark.parametrize("fmt", WRITABLE_FORMATS) + def test_ragged_grid_round_trip_adds_no_nodes(self, fmt, ragged_grid, tmp_path): + """Padding must not survive a round trip as a usable vertex. + + Covers the whole family at once: an unguarded index offset, a narrowing + cast that truncates the fill value into a small valid index, and padding + written out as NaN coordinates that dedupe into a phantom node. + """ + reloaded = _write_and_reload(ragged_grid, fmt, tmp_path) + + assert reloaded.n_face == ragged_grid.n_face + assert reloaded.n_node == ragged_grid.n_node, ( + f"{fmt}: round trip changed the node count, " + "which means padding became a real vertex" + ) + assert not np.isnan(reloaded.node_lon.values).any(), f"{fmt}: NaN node_lon" + assert not np.isnan(reloaded.node_lat.values).any(), f"{fmt}: NaN node_lat" + + # Anything that is not the fill value has to be a usable index. A + # negative leftover is the dangerous case: it is a valid Python index + # that silently wraps to the end of the coordinate array. + conn = reloaded.face_node_connectivity.values + valid = conn[conn != INT_FILL_VALUE] + assert valid.min() >= 0, f"{fmt}: negative index left in connectivity" + assert valid.max() < reloaded.n_node, f"{fmt}: out-of-range node index" + + # Node renumbering is allowed; changing the shape of a face is not. + assert _face_geometry(reloaded) == _face_geometry(ragged_grid), ( + f"{fmt}: face geometry changed across a round trip" + ) + + @pytest.mark.parametrize("fmt", list(ENCODED_INDEX_VARS)) + def test_ragged_grid_encodes_usable_indices(self, fmt, ragged_grid): + """Every index written out must name a real node or be the fill value. + + A round trip can hide this. Exodus connectivity is int64, so the + writer's +1 offset and the reader's -1 cancel exactly at + INT_FILL_VALUE: the grid reloads intact even when the file holds an + index no other Exodus reader could use. Check the encoded output + directly rather than trusting the trip back. + """ + encoded = ragged_grid.to_xarray(fmt) + prefix, fill = ENCODED_INDEX_VARS[fmt] + + index_vars = [v for v in encoded.data_vars if v.startswith(prefix)] + assert index_vars, f"{fmt}: no connectivity variable written" + + for name in index_vars: + values = encoded[name].values + if fill is not None: + values = values[values != fill] + assert values.min() >= 1, f"{fmt}: {name} holds an index below 1" + assert values.max() <= ragged_grid.n_node, ( + f"{fmt}: {name} indexes a node that does not exist" + ) + + @pytest.mark.parametrize("fmt", EXACT_CONNECTIVITY_FORMATS) + def test_ragged_grid_round_trip_is_exact(self, fmt, ragged_grid, tmp_path): + """Index-based formats restore ragged connectivity verbatim. + + Face order matters as much as face content: a reordered mesh silently + misaligns face-centered data with the faces it describes. + """ + reloaded = _write_and_reload(ragged_grid, fmt, tmp_path) + + assert_array_equal( + ragged_grid.face_node_connectivity.values, + reloaded.face_node_connectivity.values, + err_msg=f"{fmt}: ragged connectivity not preserved", + ) + assert_array_equal( + reloaded.n_nodes_per_face.values, + np.array([4, 3, 3]), + err_msg=f"{fmt}: face sizes not preserved", + ) diff --git a/uxarray/io/_esmf.py b/uxarray/io/_esmf.py index 195f0c91e..0c416f11d 100644 --- a/uxarray/io/_esmf.py +++ b/uxarray/io/_esmf.py @@ -93,12 +93,19 @@ def _read_esmf(in_ds): # assume start index is 1 if one is not provided start_index = 1 - face_node_connectivity = in_ds["elementConn"].astype(INT_DTYPE) - face_node_connectivity = xr.where( - face_node_connectivity != INT_FILL_VALUE, - face_node_connectivity - start_index, - face_node_connectivity, + element_conn = in_ds["elementConn"] + + # CF decoding turns the ESMF fill value into NaN, while an undecoded read + # leaves the raw sentinel in place. Identify the padding before the integer + # cast, which preserves neither form (NaN casts to a platform-dependent + # value, not to INT_FILL_VALUE). + fill_value = element_conn.encoding.get( + "_FillValue", element_conn.attrs.get("_FillValue", -1) ) + fill_mask = element_conn.isnull() | (element_conn == fill_value) + + face_node_connectivity = element_conn.fillna(0).astype(INT_DTYPE) - start_index + face_node_connectivity = xr.where(fill_mask, INT_FILL_VALUE, face_node_connectivity) out_ds["face_node_connectivity"] = xr.DataArray( data=face_node_connectivity, @@ -144,8 +151,17 @@ def _encode_esmf(ds: xr.Dataset) -> xr.Dataset: # Face Node Connectivity (elementConn) if "face_node_connectivity" in ds: # ESMF elementConn is 1-based, with -1 for unused; UGRID is 0-based + face_node_conn = ds["face_node_connectivity"] + + # Only offset the valid indices. Applying the offset to INT_FILL_VALUE and + # letting it fall through to the int32 encoding below truncates it into a + # small, valid node index, silently turning padding into real vertices. + element_conn = xr.where( + face_node_conn == INT_FILL_VALUE, -1, face_node_conn + 1 + ) + out_ds["elementConn"] = xr.DataArray( - ds["face_node_connectivity"] + 1, + element_conn, dims=("elementCount", "maxNodePElement"), attrs={ "long_name": "Node Indices that define the element connectivity", diff --git a/uxarray/io/_exodus.py b/uxarray/io/_exodus.py index e567fa3bc..a5968ed8e 100644 --- a/uxarray/io/_exodus.py +++ b/uxarray/io/_exodus.py @@ -91,6 +91,18 @@ def _read_exodus(ext_ds): else: face_nodes = np.vstack(padded_blocks) + if "elem_num_map" in ext_ds: + # Blocks are stored grouped by element type; elem_num_map gives each + # element's original position. Only honor it when it is a genuine + # permutation, since Exodus also allows arbitrary user-assigned IDs. + elem_num_map = ext_ds["elem_num_map"].values.astype(INT_DTYPE) - 1 + if elem_num_map.shape == (face_nodes.shape[0],) and np.array_equal( + np.sort(elem_num_map), np.arange(face_nodes.shape[0]) + ): + unpermuted = np.empty_like(face_nodes) + unpermuted[elem_num_map] = face_nodes + face_nodes = unpermuted + # standardize fill values and data type face nodes face_nodes = _replace_fill_values( grid_var=xr.DataArray(face_nodes - 1), # Wrap numpy array in a DataArray @@ -200,8 +212,10 @@ def _encode_exodus(ds, outfile=None): conn_nofill = [] for row in ds["face_node_connectivity"].values: - # Find the index of the first fill value (-1) - fill_val_idx = np.where(row == -1)[0] + # Find the index of the first fill value. Padding is stored as + # INT_FILL_VALUE, not -1; matching on -1 never fires, so every face is + # treated as full width and the padding is written out as a node index. + fill_val_idx = np.where(row == INT_FILL_VALUE)[0] if fill_val_idx.size > 0: num_nodes = fill_val_idx[0] @@ -213,7 +227,13 @@ def _encode_exodus(ds, outfile=None): conn_nofill.append(row.astype(int).tolist()) num_blks = np.count_nonzero(num_el_all_blks) - conn_nofill.sort(key=len) + + # Exodus element blocks are homogeneous, so a mixed mesh has to be regrouped + # by face size. Sort stably and carry each face's original position along, so + # the ordering can be written out below and restored on read. + block_order = sorted(range(len(conn_nofill)), key=lambda i: len(conn_nofill[i])) + conn_nofill = [conn_nofill[i] for i in block_order] + nonzero_el_index_blks = np.nonzero(num_el_all_blks)[0] start = 0 @@ -242,6 +262,13 @@ def _encode_exodus(ds, outfile=None): # Correctly increment the start index for the next block start += num_elem_in_blk + # Record where each written element came from in the original face ordering. + # Without this a mixed mesh comes back permuted, silently misaligning any + # face-centered data with the faces it describes. + exo_ds["elem_num_map"] = xr.DataArray( + data=np.asarray(block_order, dtype=np.int64) + 1, dims=["num_elem"] + ) + # --- Element Block Properties --- prop1_vals = np.arange(1, num_blks + 1, 1, dtype=np.int32) exo_ds["eb_prop1"] = xr.DataArray( diff --git a/uxarray/io/_scrip.py b/uxarray/io/_scrip.py index 21edb1377..f3673363e 100644 --- a/uxarray/io/_scrip.py +++ b/uxarray/io/_scrip.py @@ -202,30 +202,21 @@ def _encode_scrip(face_node_connectivity, node_lon, node_lat, face_areas): n_face = face_node_connectivity.shape[0] n_max_nodes = face_node_connectivity.shape[1] - # --- Core logic enhanced with Implementation 2's robust method --- - # Flatten the connectivity array to easily work with all node indices - f_nodes_flat = face_node_connectivity.values.astype(int).ravel() - - # Create a mask to identify valid nodes vs. fill values - valid_nodes_mask = f_nodes_flat != INT_FILL_VALUE - - # Create arrays to hold final lat/lon data, filled with NaN - lat_nodes_flat = np.full(f_nodes_flat.shape, np.nan, dtype=np.float64) - lon_nodes_flat = np.full(f_nodes_flat.shape, np.nan, dtype=np.float64) - - # Get the flattened indices of the valid nodes (where the mask is True) - valid_indices = np.where(valid_nodes_mask)[0] - # Get the actual node indices from the connectivity array for those valid positions - valid_node_ids = f_nodes_flat[valid_indices] - - # Use the valid indices to populate the coordinate arrays correctly - lon_nodes_flat[valid_indices] = node_lon.values[valid_node_ids] - lat_nodes_flat[valid_indices] = node_lat.values[valid_node_ids] - - # Reshape the 1D arrays back to 2D - reshp_lat = lat_nodes_flat.reshape((n_face, n_max_nodes)) - reshp_lon = lon_nodes_flat.reshape((n_face, n_max_nodes)) - # --- End of enhanced logic --- + conn = face_node_connectivity.values.astype(INT_DTYPE) + valid_nodes_mask = conn != INT_FILL_VALUE + + # SCRIP has no fill value for corners. A face with fewer than grid_corners + # vertices is written as a degenerate polygon that repeats its last valid + # corner. Writing NaN into the padded slots instead makes the reader dedupe + # them into a phantom NaN node, which both inflates n_node and poisons the + # node coordinates for every downstream geometry calculation. + last_valid = np.maximum.accumulate( + np.where(valid_nodes_mask, np.arange(n_max_nodes), 0), axis=1 + ) + padded_conn = np.take_along_axis(conn, last_valid, axis=1) + + reshp_lon = node_lon.values[padded_conn] + reshp_lat = node_lat.values[padded_conn] # Add data to new scrip output file ds["grid_corner_lat"] = xr.DataArray(