diff --git a/test/io/test_esmf.py b/test/io/test_esmf.py index fbf910e1c..8885a3ec4 100644 --- a/test/io/test_esmf.py +++ b/test/io/test_esmf.py @@ -1,9 +1,8 @@ -import uxarray as ux -import os +import numpy as np import pytest +import uxarray as ux import xarray as xr -import numpy as np -from uxarray.constants import ERROR_TOLERANCE +from uxarray.constants import INT_FILL_VALUE def test_read_esmf(gridpath): @@ -35,71 +34,55 @@ 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. - - 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) +def test_encode_esmf_structure(gridpath): + """Encoding to ESMF produces the variables the format requires. - 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")) + uxgrid = ux.open_grid(gridpath("ugrid", "outCSne30", "outCSne30.ug")) + esmf_dataset = uxgrid.to_xarray("ESMF") - # 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 + assert 'numElementConn' in esmf_dataset + + # 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() + ) + - # 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) +@pytest.mark.parametrize("mask_and_scale", [True, False]) +def test_read_esmf_padding_independent_of_cf_decoding(mask_and_scale, tmp_path): + """Padding is recognized whether or not xarray decoded the fill value. + + CF decoding replaces the ``-1`` padding with NaN and promotes elementConn to + float; an undecoded read hands back the raw int32. Neither survives the cast + to INT_DTYPE as INT_FILL_VALUE, so both must be identified before it. + """ + 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=np.array([ + [0, 1, 2, 3], + [1, 4, 2, INT_FILL_VALUE], + [0, 3, 4, INT_FILL_VALUE], + ]), + fill_value=INT_FILL_VALUE, + ) + + path = tmp_path / "esmf_ragged.nc" + uxgrid.to_xarray("ESMF").to_netcdf(path) + + with xr.open_dataset(path, mask_and_scale=mask_and_scale) as ds: + reloaded = ux.open_grid(ds) + + np.testing.assert_array_equal( + reloaded.face_node_connectivity.values, + uxgrid.face_node_connectivity.values, + ) + np.testing.assert_array_equal(reloaded.n_nodes_per_face.values, [4, 3, 3]) diff --git a/test/io/test_exodus.py b/test/io/test_exodus.py index ea37f5010..9a06b7076 100644 --- a/test/io/test_exodus.py +++ b/test/io/test_exodus.py @@ -2,6 +2,7 @@ import numpy as np import pytest import uxarray as ux +import xarray as xr from uxarray.constants import INT_DTYPE, INT_FILL_VALUE @@ -21,8 +22,76 @@ 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_read_exodus_ignores_third_party_elem_num_map(gridpath, tmp_path): + """Element IDs from another writer must not reorder the mesh. + + ``_encode_exodus`` records the original face order in ``elem_num_map`` to undo + its own block grouping. Exodus defines that variable as a user-facing ID, which + a third-party file is free to renumber, so the reader only honors it on files + it wrote -- identified via ``qa_records``. + """ + source = gridpath("exodus", "mixed", "mixed.exo") + original = ux.open_grid(source) + + # Same mesh, same blocks, same connectivity -- only the element IDs change + with xr.open_dataset(source) as ds: + renumbered = ds.load() + n_elem = renumbered["elem_num_map"].sizes["num_elem"] + renumbered["elem_num_map"][:] = np.arange(n_elem, 0, -1) + + assert b"uxarray" not in np.asarray( + renumbered["qa_records"].values, dtype="S" + ).ravel(), "fixture is no longer a third-party file" + + path = tmp_path / "mixed_renumbered.exo" + renumbered.to_netcdf(path) + + np.testing.assert_array_equal( + ux.open_grid(path).face_node_connectivity.values, + original.face_node_connectivity.values, + ) 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..3316afca9 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, not node indices, so its reader rebuilds and +# renumbers nodes. Geometry and face sizes survive a round trip; the index each +# node lands on does not. +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,132 @@ 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", + ) + + @pytest.mark.parametrize("fmt", WRITABLE_FORMATS) + def test_ragged_grid_round_trip_preserves_face_sizes( + self, fmt, ragged_grid, tmp_path + ): + """A short face comes back short, in every writable format. + + Node indices may be renumbered, but a triangle must not reload as a quad. + SCRIP is the interesting case: with no corner fill value it repeats a + corner instead, and the reader has to collapse those repeats. Leaving them + keeps the node count honest while still widening the face. + """ + reloaded = _write_and_reload(ragged_grid, fmt, tmp_path) + + assert_array_equal( + reloaded.n_nodes_per_face.values, + ragged_grid.n_nodes_per_face.values, + err_msg=f"{fmt}: face sizes not preserved", + ) + assert reloaded.n_edge == ragged_grid.n_edge, ( + f"{fmt}: edge count changed -- a face gained a duplicate vertex" + ) + + # No face may name the same node twice + for face, row in enumerate(reloaded.face_node_connectivity.values): + nodes = [i for i in row if i != INT_FILL_VALUE] + assert len(nodes) == len(set(nodes)), ( + f"{fmt}: face {face} references a node more than once" + ) diff --git a/test/io/test_scrip.py b/test/io/test_scrip.py index 4b30ddb0c..09fe10e4f 100644 --- a/test/io/test_scrip.py +++ b/test/io/test_scrip.py @@ -7,7 +7,7 @@ import uxarray as ux from uxarray.constants import INT_DTYPE, INT_FILL_VALUE -from uxarray.io._scrip import _detect_multigrid +from uxarray.io._scrip import _collapse_repeated_corners, _detect_multigrid def test_read_ugrid(gridpath, mesh_constants): @@ -137,3 +137,43 @@ def test_open_multigrid_mask_active_value_per_grid_override(gridpath): assert grids["ocn"].n_face == expected_ocn assert grids["atm"].n_face == expected_atm + + +def test_collapse_repeated_corners(): + """SCRIP pads a short face by repeating a corner; the reader undoes that. + + Duplicates drop keeping the first occurrence, so winding order is preserved + and a closed ring collapses like a trailing repeat. + """ + face_nodes = np.array( + [ + [0, 1, 2, 3], # a real quad, untouched + [4, 5, 6, 6], # trailing repeat, as _encode_scrip writes it + [7, 8, 9, 7], # closed ring, repeating the first corner + [1, 1, 2, 3], # repeat in the middle of the row + ], + dtype=INT_DTYPE, + ) + + nt.assert_array_equal( + _collapse_repeated_corners(face_nodes), + np.array( + [ + [0, 1, 2, 3], + [4, 5, 6, INT_FILL_VALUE], + [7, 8, 9, INT_FILL_VALUE], + [1, 2, 3, INT_FILL_VALUE], + ] + ), + ) + + +def test_collapse_repeated_corners_keeps_degenerate_faces(): + """A cell that would collapse below three vertices is left alone. + + It is degenerate in the source file; packing it to one or two nodes would + invent a face no polygon routine can use, so the degeneracy stays visible. + """ + face_nodes = np.array([[0, 0, 0, 0], [1, 1, 2, 2]], dtype=INT_DTYPE) + + nt.assert_array_equal(_collapse_repeated_corners(face_nodes), face_nodes) diff --git a/uxarray/io/_esmf.py b/uxarray/io/_esmf.py index 538fbf171..6879ddaa3 100644 --- a/uxarray/io/_esmf.py +++ b/uxarray/io/_esmf.py @@ -93,12 +93,22 @@ 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"] + face_dim, node_dim = element_conn.dims + + # "numElementConn" gives the face size, so locate the padding positionally. + # Matching the sentinel means guessing: CF decoding turns it into NaN, and the + # cast below preserves neither NaN nor the raw value as INT_FILL_VALUE. + positions = xr.DataArray( + np.arange(element_conn.sizes[node_dim], dtype=INT_DTYPE), dims=node_dim ) + fill_mask = (positions >= n_nodes_per_face).transpose(face_dim, node_dim) + + # NaN is never a usable index, whatever "numElementConn" claims + fill_mask = fill_mask | element_conn.isnull() + + 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 +154,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..18acce5be 100644 --- a/uxarray/io/_exodus.py +++ b/uxarray/io/_exodus.py @@ -9,6 +9,23 @@ from uxarray.grid.connectivity import _replace_fill_values from uxarray.grid.coordinates import _lonlat_rad_to_xyz, _xyz_to_lonlat_deg +# Producer name written into (and looked for in) "qa_records" +_WRITER_NAME = "uxarray" + + +def _written_by_uxarray(ext_ds): + """Whether ``ext_ds`` was produced by :func:`_encode_exodus`. + + Exodus defines ``elem_num_map`` as each element's user-facing ID, not its + original position, so only files this writer produced can be reordered from + it. A third-party file keeps the face order it was written with. + """ + if "qa_records" not in ext_ds: + return False + + records = np.asarray(ext_ds["qa_records"].values, dtype="S").ravel() + return _WRITER_NAME.encode() in records + # Exodus Number is one-based. def _read_exodus(ext_ds): @@ -91,6 +108,18 @@ def _read_exodus(ext_ds): else: face_nodes = np.vstack(padded_blocks) + if "elem_num_map" in ext_ds and _written_by_uxarray(ext_ds): + # _encode_exodus groups elements into blocks by face size, permuting a + # mixed mesh, and records each original position here. Undo that, but only + # for a genuine permutation -- Exodus also allows arbitrary 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 @@ -174,7 +203,7 @@ def _encode_exodus(ds, outfile=None): # --- QA Records --- ux_exodus_version = "1.0" - qa_records = [["uxarray"], [ux_exodus_version], [date], [time]] + qa_records = [[_WRITER_NAME], [ux_exodus_version], [date], [time]] exo_ds["qa_records"] = xr.DataArray( data=np.array(qa_records, dtype="S"), dims=["num_qa_rec", "four"], @@ -200,8 +229,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 +244,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 +279,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 b24edac60..218dd3c95 100644 --- a/uxarray/io/_scrip.py +++ b/uxarray/io/_scrip.py @@ -37,6 +37,49 @@ def _values_in_degrees(data_array): return values +def _collapse_repeated_corners(face_nodes): + """Turn SCRIP's repeated corners back into fill values. + + SCRIP has no fill value for corners, so a face with fewer than + ``grid_corners`` vertices repeats one of them. Left in place each repeat reads + as a real vertex, keeping the face full width and adding a zero-length edge. + + Duplicates are dropped keeping the first occurrence, so winding order survives + and a closed ring (last corner repeating the first) collapses too. + + Parameters + ---------- + face_nodes : numpy.ndarray + ``(n_face, n_max_face_nodes)`` node indices, with repeated corners. + + Returns + ------- + numpy.ndarray + Connectivity with repeats replaced by ``INT_FILL_VALUE``, packed left. + """ + n_face, n_corners = face_nodes.shape + + keep = np.ones(face_nodes.shape, dtype=bool) + for corner in range(1, n_corners): + keep[:, corner] = (face_nodes[:, [corner]] != face_nodes[:, :corner]).all( + axis=1 + ) + + # A cell that would collapse below three vertices is degenerate in the source + # file; leave it alone rather than invent a one- or two-node face. + keep[keep.sum(axis=1) < 3] = True + + if keep.all(): + return face_nodes + + collapsed = np.full(face_nodes.shape, INT_FILL_VALUE, dtype=face_nodes.dtype) + rows = np.broadcast_to(np.arange(n_face)[:, None], face_nodes.shape) + destination = np.cumsum(keep, axis=1) - 1 + collapsed[rows[keep], destination[keep]] = face_nodes[keep] + + return collapsed + + def _to_ugrid(in_ds, out_ds): """If input dataset (``in_ds``) file is an unstructured SCRIP file, function will reassign SCRIP variables to UGRID conventions in output file @@ -81,6 +124,9 @@ def _to_ugrid(in_ds, out_ds): # Reshape face nodes array into original shape for use in 'face_node_connectivity' unq_inv = np.reshape(unq_inv, (len(in_ds.grid_size), len(in_ds.grid_corners))) + # Recover the real face sizes from the degenerate corners SCRIP pads with + unq_inv = _collapse_repeated_corners(unq_inv) + # Create node_lon & node_lat out_ds[ugrid.NODE_COORDINATES[0]] = xr.DataArray( unq_lon, dims=[ugrid.NODE_DIM], attrs=ugrid.NODE_LON_ATTRS @@ -104,10 +150,11 @@ def _to_ugrid(in_ds, out_ds): attrs=ugrid.FACE_LAT_ATTRS, ) - # standardize fill values and data type face nodes + # standardize fill values and data type face nodes. The padding here comes + # from the collapse above; SCRIP has no fill value of its own. face_nodes = _replace_fill_values( xr.DataArray(data=unq_inv), - original_fill=-1, + original_fill=INT_FILL_VALUE, new_fill=INT_FILL_VALUE, new_dtype=INT_DTYPE, ) @@ -202,30 +249,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) + conn = face_node_connectivity.values.astype(INT_DTYPE) + valid_nodes_mask = conn != INT_FILL_VALUE - # 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] + # 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) - # 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 --- + 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(