o Enhance dual mesh construction with comprehensive testing and stability improvements

Author: rajeeja

test/test_grid.py

@@ -15,7 +15,7 @@
 
 from uxarray.grid.coordinates import _populate_node_latlon, _lonlat_rad_to_xyz, _xyz_to_lonlat_rad_scalar
 
-from uxarray.constants import INT_FILL_VALUE, ERROR_TOLERANCE
+from uxarray.constants import INT_FILL_VALUE, ERROR_TOLERANCE, INT_DTYPE
 
 from uxarray.grid.arcs import extreme_gca_latitude
 
@@ -748,6 +748,139 @@ def test_dual_duplicate():
         dataset.get_dual()
 
 
+def test_dual_mesh_parallel():
+    """Test dual mesh construction with parallel processing enabled."""
+    import numba
+    from uxarray.grid.dual import construct_faces
+
+    # Test with a simple grid
+    grid = ux.open_grid(gridfile_mpas, use_dual=False)
+
+    # Get the inputs for construct_faces
+    dual_node_x = grid.face_x.values
+    dual_node_y = grid.face_y.values
+    dual_node_z = grid.face_z.values
+    node_x = grid.node_x.values
+    node_y = grid.node_y.values
+    node_z = grid.node_z.values
+    node_face_connectivity = grid.node_face_connectivity.values
+
+    # Get an array with the number of edges for each face
+    n_edges_mask = node_face_connectivity != INT_FILL_VALUE
+    n_edges = np.sum(n_edges_mask, axis=1)
+    max_edges = len(node_face_connectivity[0])
+
+    valid_node_indices = np.where(n_edges >= 3)[0]
+    construct_node_face_connectivity = np.full(
+        (len(valid_node_indices), max_edges), INT_FILL_VALUE, dtype=INT_DTYPE
+    )
+
+    # Test that construct_faces works (this tests the numba compilation)
+    try:
+        result = construct_faces(
+            valid_node_indices,
+            n_edges,
+            dual_node_x,
+            dual_node_y,
+            dual_node_z,
+            node_face_connectivity,
+            node_x,
+            node_y,
+            node_z,
+            construct_node_face_connectivity,
+            max_edges,
+        )
+
+        # Verify the result has the expected shape and properties
+        assert result is not None
+        assert result.dtype == INT_DTYPE
+        assert result.ndim == 2
+
+        # Compare with the standard dual construction
+        dual = grid.get_dual()
+        expected_shape = dual.face_node_connectivity.values.shape
+
+        # The shapes should match since we're using the same algorithm
+        assert result.shape == expected_shape
+
+    except numba.errors.NumbaError as e:
+        pytest.skip(f"Numba compilation failed: {e}")
+    except Exception as e:
+        pytest.fail(f"Dual mesh parallel construction failed: {e}")
+
+
+def test_dual_mesh_parallel_validation():
+    """Test that parallel dual construction produces identical results and validates threading."""
+    import time
+    from uxarray.grid.dual import construct_faces
+
+    grid = ux.open_grid(gridfile_mpas, use_dual=False)
+
+    # Prepare inputs for direct function call
+    dual_node_x = grid.face_x.values
+    dual_node_y = grid.face_y.values
+    dual_node_z = grid.face_z.values
+    node_x = grid.node_x.values
+    node_y = grid.node_y.values
+    node_z = grid.node_z.values
+    node_face_connectivity = grid.node_face_connectivity.values
+
+    n_edges_mask = node_face_connectivity != INT_FILL_VALUE
+    n_edges = np.sum(n_edges_mask, axis=1)
+    max_edges = node_face_connectivity.shape[1]
+
+    valid_node_indices = np.where(n_edges >= 3)[0]
+
+    # Test multiple runs for consistency
+    results = []
+    times = []
+
+    for i in range(3):
+        construct_node_face_connectivity = np.full(
+            (len(valid_node_indices), max_edges), INT_FILL_VALUE, dtype=INT_DTYPE
+        )
+
+        start_time = time.time()
+        result = construct_faces(
+            valid_node_indices,
+            n_edges,
+            dual_node_x,
+            dual_node_y,
+            dual_node_z,
+            node_face_connectivity,
+            node_x,
+            node_y,
+            node_z,
+            construct_node_face_connectivity,
+            max_edges,
+        )
+        elapsed = time.time() - start_time
+
+        results.append(result)
+        times.append(elapsed)
+
+    # Verify all results are identical (parallel consistency)
+    for i in range(1, len(results)):
+        nt.assert_array_equal(results[0], results[i],
+                             err_msg=f"Parallel dual construction run {i} differs from run 0")
+
+    # Verify results are valid
+    for result in results:
+        assert result is not None
+        assert result.dtype == INT_DTYPE
+        assert result.ndim == 2
+        assert result.shape[0] == len(valid_node_indices)
+        assert result.shape[1] == max_edges
+
+    # Verify performance is reasonable (should complete in reasonable time)
+    avg_time = np.mean(times)
+    assert avg_time < 30.0, f"Dual construction too slow: {avg_time:.2f}s"
+
+    # Verify against standard implementation
+    dual_standard = grid.get_dual()
+    assert results[0].shape == dual_standard.face_node_connectivity.values.shape
+
+
 def test_normalize_existing_coordinates_non_norm_initial():
     gridfile_mpas = current_path / "meshfiles" / "mpas" / "QU" / "mesh.QU.1920km.151026.nc"
     from uxarray.grid.validation import _check_normalization

uxarray/grid/dual.py

@@ -34,8 +34,9 @@ def construct_dual(grid):
     # Get an array with the number of edges for each face
     n_edges_mask = node_face_connectivity != INT_FILL_VALUE
     n_edges = np.sum(n_edges_mask, axis=1)
-    max_edges = len(node_face_connectivity[0])
+    max_edges = node_face_connectivity.shape[1]
 
+    # Only nodes with 3+ edges can form valid dual faces
     valid_node_indices = np.where(n_edges >= 3)[0]
 
     construct_node_face_connectivity = np.full(
@@ -80,25 +81,25 @@ def construct_faces(
     Parameters
     ----------
     valid_node_indices: np.ndarray
-        number of valid nodes in the primal mesh
+        Array of node indices with at least 3 connections in the primal mesh
     n_edges: np.ndarray
-        array of the number of edges for each dual face
+        Array of the number of edges for each node in the primal mesh
     dual_node_x: np.ndarray
-        x node coordinates for the dual mesh
+        x coordinates for the dual mesh nodes (face centers of primal mesh)
     dual_node_y: np.ndarray
-        y node coordinates for the dual mesh
+        y coordinates for the dual mesh nodes (face centers of primal mesh)
     dual_node_z: np.ndarray
-        z node coordinates for the dual mesh
+        z coordinates for the dual mesh nodes (face centers of primal mesh)
     node_face_connectivity: np.ndarray
-        `node_face_connectivity` of the primal mesh
+        Node-to-face connectivity of the primal mesh
     node_x: np.ndarray
-        x node coordinates from the primal mesh
+        x coordinates of nodes from the primal mesh
     node_y: np.ndarray
-        y node coordinates from the primal mesh
+        y coordinates of nodes from the primal mesh
     node_z: np.ndarray
-        z node coordinates from the primal mesh
+        z coordinates of nodes from the primal mesh
     construct_node_face_connectivity: np.ndarray
-        Empty array to store connectivity
+        Pre-allocated array to store the dual mesh connectivity
     max_edges: int
         The max number of edges in a face
 
@@ -107,6 +108,13 @@ def construct_faces(
     --------
     construct_node_face_connectivity : ndarray
         Constructed node_face_connectivity for the dual mesh
+
+    Notes
+    -----
+    In dual mesh construction, the "valid node indices" are face indices from
+    the primal mesh's node_face_connectivity that are not fill values. These
+    represent the actual faces that each primal node connects to, which become
+    the nodes of the dual mesh faces.
     """
     n_valid = valid_node_indices.shape[0]
 
@@ -118,14 +126,13 @@ def construct_faces(
             [INT_FILL_VALUE for _ in range(n_edges[i])], dtype=INT_DTYPE
         )
 
-        # Get a list of the valid non fill value nodes
+        # Get the face indices this node connects to (these become dual face nodes)
         connected_faces = node_face_connectivity[i][0 : n_edges[i]]
-        index = 0
 
         # Connect the face centers around the node to make dual face
-        for node_idx in connected_faces:
-            temp_face[index] = node_idx
-            index += 1
+        for index, node_idx in enumerate(connected_faces):
+            if node_idx != INT_FILL_VALUE:
+                temp_face[index] = node_idx
 
         # Order the nodes using the angles so the faces have nodes in counter-clockwise sequence
         node_central = np.array([node_x[i], node_y[i], node_z[i]])
@@ -138,7 +145,7 @@ def construct_faces(
         )
 
         # Order the face nodes properly in a counter-clockwise fashion
-        if temp_face[0] is not INT_FILL_VALUE:
+        if temp_face[0] != INT_FILL_VALUE:
             _face = _order_nodes(
                 temp_face,
                 node_0,
@@ -192,10 +199,18 @@ def _order_nodes(
     final_face : np.ndarray
         The face in proper counter-clockwise order
     """
+    # Add numerical stability check for degenerate cases
+    if n_edges < 3:
+        return np.full(max_edges, INT_FILL_VALUE, dtype=INT_DTYPE)
+
     node_zero = node_0 - node_central
+    node_zero_mag = np.linalg.norm(node_zero)
+
+    # Check for numerical stability
+    if node_zero_mag < 1e-15:
+        return np.full(max_edges, INT_FILL_VALUE, dtype=INT_DTYPE)
 
     node_cross = np.cross(node_0, node_central)
-    node_zero_mag = np.linalg.norm(node_zero)
 
     d_angles = np.zeros(n_edges, dtype=np.float64)
     d_angles[0] = 0.0
@@ -214,11 +229,16 @@ def _order_nodes(
             node_diff = sub - node_central
             node_diff_mag = np.linalg.norm(node_diff)
 
+            # Skip if node difference is too small (numerical stability)
+            if node_diff_mag < 1e-15:
+                d_angles[j] = 0.0
+                continue
+
             d_side = np.dot(node_cross, node_diff)
             d_dot_norm = np.dot(node_zero, node_diff) / (node_zero_mag * node_diff_mag)
 
-            if d_dot_norm > 1.0:
-                d_dot_norm = 1.0
+            # Clamp to valid range for arccos to avoid numerical errors
+            d_dot_norm = max(-1.0, min(1.0, d_dot_norm))
 
             d_angles[j] = np.arccos(d_dot_norm)