diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index 5e3f53d02..fd871d6fa 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -111,6 +111,9 @@ def teardown(self, resolution, *args, **kwargs): def time_n_nodes_per_face(self, resolution): _ = self.uxgrid.n_nodes_per_face.compute() + def time_n_nodes_per_face(self, resolution): + _ = self.uxgrid.n_nodes_per_face + def time_face_node(self, resolution): _ = self.uxgrid.face_node_connectivity.compute() diff --git a/test/grid/geometry/test_centroids.py b/test/grid/geometry/test_centroids.py index 1ee8c7b38..887dcc0f0 100644 --- a/test/grid/geometry/test_centroids.py +++ b/test/grid/geometry/test_centroids.py @@ -63,13 +63,16 @@ def test_edge_centroids_from_triangle(): grid = ux.open_grid(test_triangle, latlon=False) _populate_edge_centroids(grid) - centroid_x = np.mean(grid.node_x[grid.edge_node_connectivity[0][0:]]) - centroid_y = np.mean(grid.node_y[grid.edge_node_connectivity[0][0:]]) - centroid_z = np.mean(grid.node_z[grid.edge_node_connectivity[0][0:]]) + edge_nodes = grid.edge_node_connectivity.values - assert centroid_x == grid.edge_x[0] - assert centroid_y == grid.edge_y[0] - assert centroid_z == grid.edge_z[0] + centroid_x = grid.node_x.values[edge_nodes].mean(axis=1) + centroid_y = grid.node_y.values[edge_nodes].mean(axis=1) + centroid_z = grid.node_z.values[edge_nodes].mean(axis=1) + centroid_x, centroid_y, centroid_z = _normalize_xyz(centroid_x, centroid_y, centroid_z) + + nt.assert_array_almost_equal(grid.edge_x.values, centroid_x) + nt.assert_array_almost_equal(grid.edge_y.values, centroid_y) + nt.assert_array_almost_equal(grid.edge_z.values, centroid_z) def test_edge_centroids_from_mpas(gridpath): """Test computed centroid values compared to values from a MPAS dataset.""" diff --git a/test/grid/grid/test_connectivity.py b/test/grid/grid/test_connectivity.py index 7ae6087a5..afaccd5b2 100644 --- a/test/grid/grid/test_connectivity.py +++ b/test/grid/grid/test_connectivity.py @@ -3,10 +3,12 @@ import pytest import uxarray as ux -from uxarray.constants import INT_FILL_VALUE, ERROR_TOLERANCE +from uxarray.constants import INT_DTYPE, INT_FILL_VALUE, ERROR_TOLERANCE from uxarray.grid.connectivity import (_populate_face_edge_connectivity, _build_edge_face_connectivity, _build_edge_node_connectivity, _build_face_face_connectivity, _populate_face_face_connectivity) +from uxarray.grid.utils import (_adaptive_sort_bucket, _insertion_sort_bucket, + MIN_ADAPTIVE_SORT_SIZE) def test_connectivity_build_n_nodes_per_face(gridpath): @@ -63,6 +65,83 @@ def test_connectivity_build_face_edges_connectivity(gridpath): assert np.all(valid_edges >= 0) assert np.all(valid_edges < uxgrid.n_edge) +@pytest.mark.parametrize("grid_parts", [("ugrid", "outCSne30", "outCSne30.ug"), + ("ugrid", "quad-hexagon", "grid.nc"), + ("ugrid", "geoflow-small", "grid.nc")]) +def test_connectivity_edge_node_canonical_order(gridpath, grid_parts): + """Test that constructed edges are numbered in lexicographic node order.""" + uxgrid = ux.open_grid(gridpath(*grid_parts)) + edge_nodes = uxgrid.edge_node_connectivity.values + + # Each edge is stored as an ascending node pair + assert np.all(edge_nodes[:, 0] < edge_nodes[:, 1]) + + # Edges are numbered lexicographically by that pair, with no duplicates + lexicographic_order = np.lexsort((edge_nodes[:, 1], edge_nodes[:, 0])) + nt.assert_array_equal(lexicographic_order, np.arange(uxgrid.n_edge)) + assert len(np.unique(edge_nodes, axis=0)) == uxgrid.n_edge + +@pytest.mark.parametrize("sort", [_insertion_sort_bucket, _adaptive_sort_bucket], + ids=["insertion", "adaptive"]) +def test_connectivity_bucket_sort(sort): + """Test that each bucket sort orders its own slice and nothing else. + + The bucket sizes straddle ``MIN_ADAPTIVE_SORT_SIZE``: the small ones cannot accumulate + enough shifts to exhaust the budget, so the metered sort stays on its insertion path, + while the 500 element bucket is shuffled far past the budget and falls back to the heap + sort. Keys repeat, since an interior edge reaches its bucket once per adjacent face. + """ + rng = np.random.default_rng(0) + + sizes = [5, MIN_ADAPTIVE_SORT_SIZE, MIN_ADAPTIVE_SORT_SIZE + 1, 500] + bounds = np.cumsum([0] + sizes) + n_half_edge = int(bounds[-1]) + buckets = list(zip(bounds[:-1], bounds[1:])) + + keys = rng.integers(0, 40, n_half_edge).astype(INT_DTYPE) + order = rng.permutation(n_half_edge).astype(INT_DTYPE) + + # the key each half edge must still be paired with once the permutation has moved it + key_for = np.empty(n_half_edge, dtype=INT_DTYPE) + key_for[order] = keys + + expected_keys = np.concatenate([np.sort(keys[start:end]) for start, end in buckets]) + + got_keys, got_order = keys.copy(), order.copy() + for start, end in buckets: + shuffle = rng.permutation(end - start) + got_keys[start:end] = got_keys[start:end][shuffle] + got_order[start:end] = got_order[start:end][shuffle] + + sort(got_keys, got_order, start, end - start) + + nt.assert_array_equal(got_keys, expected_keys) + + # sorted keys alone would pass even if the permutation had been scrambled independently + nt.assert_array_equal(key_for[got_order], got_keys) + nt.assert_array_equal(np.sort(got_order), np.arange(n_half_edge)) + + +def test_connectivity_face_edge_positional_alignment(gridpath): + """Test that face_edge_connectivity[i, j] is the edge between face nodes j and j+1.""" + uxgrid = ux.open_grid(gridpath("ugrid", "outCSne30", "outCSne30.ug")) + + face_nodes = uxgrid.face_node_connectivity.values + face_edges = uxgrid.face_edge_connectivity.values + edge_nodes = uxgrid.edge_node_connectivity.values + + for face_idx, n_edges in enumerate(uxgrid.n_nodes_per_face.values): + for cur in range(n_edges): + start_node = face_nodes[face_idx, cur] + end_node = face_nodes[face_idx, (cur + 1) % n_edges] + + expected = sorted((start_node, end_node)) + actual = sorted(edge_nodes[face_edges[face_idx, cur]]) + assert actual == expected + + # Remaining slots stay padded + assert np.all(face_edges[face_idx, n_edges:] == INT_FILL_VALUE) + def test_connectivity_build_face_edges_connectivity_fillvalues(): """Test face-edge connectivity with fill values.""" # Create a simple grid with mixed face types diff --git a/uxarray/grid/connectivity.py b/uxarray/grid/connectivity.py index ac9658979..a35d1311f 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -1,9 +1,15 @@ import numpy as np import xarray as xr -from numba import njit +from numba import njit, prange from uxarray.constants import INT_DTYPE, INT_FILL_VALUE from uxarray.conventions import ugrid +from uxarray.grid.utils import ( + _build_pair_index, + _count_unique_in_bucket, + _search_bucket, + _sort_bucket, +) def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): @@ -125,8 +131,8 @@ def _populate_n_nodes_per_face(grid): it within the internal (``Grid._ds``) and through the attribute (``Grid.n_nodes_per_face``).""" - n_nodes_per_face = _build_n_nodes_per_face( - grid.face_node_connectivity.values, grid.n_face, grid.n_max_face_nodes + n_nodes_per_face = ( + (grid.face_node_connectivity != INT_FILL_VALUE).sum(axis=1).astype(INT_DTYPE) ) if n_nodes_per_face.ndim == 0: @@ -141,98 +147,227 @@ def _populate_n_nodes_per_face(grid): ) -@njit(cache=True) -def _build_n_nodes_per_face(face_nodes, n_face, n_max_face_nodes): - """Constructs ``n_nodes_per_face``, which contains the number of non-fill- - value nodes for each face in ``face_node_connectivity``""" - - n_face, n_max_face_nodes = face_nodes.shape - n_nodes_per_face = np.empty(n_face, dtype=INT_DTYPE) - for i in range(n_face): - c = 0 - for j in range(n_max_face_nodes): - if face_nodes[i, j] != INT_FILL_VALUE: - c += 1 - n_nodes_per_face[i] = c - return n_nodes_per_face - - def _populate_edge_node_connectivity(grid): """Constructs the UGRID connectivity variable (``edge_node_connectivity``) and stores it within the internal (``Grid._ds``) and through the attribute (``Grid.edge_node_connectivity``).""" - edge_nodes, inverse_indices, fill_value_mask = _build_edge_node_connectivity( - grid.face_node_connectivity.values, grid.n_face, grid.n_max_face_nodes - ) + # Check edge coordinates already exist, if they do this might cause issues + + if "n_edge" in grid.dims: + stale = sorted(n for n in grid._ds if ugrid.EDGE_DIM in grid._ds[n].dims) + raise ValueError( + f"Constructing 'edge_node_connectivity' on a grid that already has " + f"edge-centered variables ({', '.join(stale)}). Constructed edges are " + f"numbered in lexicographic node-pair order, which need not match the " + f"numbering those variables were stored with; they may no longer refer " + f"to the same edges." + ) - edge_node_attrs = ugrid.EDGE_NODE_CONNECTIVITY_ATTRS - edge_node_attrs["inverse_indices"] = inverse_indices + # This is in lieu of an xarray equivalent to `da.compute(a, b)`. We traverse the + # grid once to gather both variables, possibly as chunks if dask is enabled + computed = xr.Dataset( + { + "face_nodes": grid.face_node_connectivity.variable, + "n_nodes_per_face": grid.n_nodes_per_face.variable, + } + ).compute() + + edge_nodes, face_edges = _build_edge_node_connectivity( + computed.face_nodes.data, computed.n_nodes_per_face.data, grid.n_node + ) - # add edge_node_connectivity to internal dataset grid._ds["edge_node_connectivity"] = xr.DataArray( - edge_nodes, dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, attrs=edge_node_attrs + edge_nodes, + dims=ugrid.EDGE_NODE_CONNECTIVITY_DIMS, + attrs=ugrid.EDGE_NODE_CONNECTIVITY_ATTRS, ) + grid._ds["face_edge_connectivity"] = xr.DataArray( + face_edges, + dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, + attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, + ) -def _build_edge_node_connectivity(face_nodes, n_face, n_max_face_nodes): - """Constructs the UGRID connectivity variable (``edge_node_connectivity``) - and stores it within the internal (``Grid._ds``) and through the attribute - (``Grid.edge_node_connectivity``). - Additionally, the attributes (``inverse_indices``) and - (``fill_value_mask``) are stored for constructing other - connectivity variables. +@njit(cache=True, inline="always") +def _canonical_half_edge(face_node_connectivity, face_idx, local_idx, n_edges): + """The ``(low, high)`` node pair of the half edge leaving face slot ``local_idx``, + wrapping back to slot 0 after ``n_edges``.""" + start_node = face_node_connectivity[face_idx, local_idx] + end_node = face_node_connectivity[face_idx, (local_idx + 1) % n_edges] + + if start_node > end_node: + return end_node, start_node + return start_node, end_node + + +@njit(cache=True) +def _count_half_edges_per_node(face_node_connectivity, n_nodes_per_face, n_node): + """Bucket offsets keyed on each half edge's lower node: bucket ``a`` will occupy + ``[bucket_offset[a], bucket_offset[a + 1])``.""" + bucket_offset = np.zeros(n_node + 1, dtype=INT_DTYPE) + + for face_idx in range(face_node_connectivity.shape[0]): + n_edges = n_nodes_per_face[face_idx] + for local_idx in range(n_edges): + node_a, _ = _canonical_half_edge( + face_node_connectivity, face_idx, local_idx, n_edges + ) + bucket_offset[node_a + 1] += 1 + + for n in range(n_node): + bucket_offset[n + 1] += bucket_offset[n] + + return bucket_offset + + +@njit(cache=True) +def _scatter_half_edges( + face_node_connectivity, n_nodes_per_face, bucket_offset, n_half_edge +): + """Fills every bucket with its half edges, leaving ``bucket_offset`` as it found it. + + Each half edge is identified by ``half_edge_slot``, its flattened position + ``face_idx * n_max_face_nodes + local_idx`` in the face node connectivity, and keyed on + ``end_node``, the higher of its two nodes.""" + n_max_face_nodes = face_node_connectivity.shape[1] + + half_edge_slot = np.empty(n_half_edge, dtype=INT_DTYPE) + end_node = np.empty(n_half_edge, dtype=INT_DTYPE) + + for face_idx in range(face_node_connectivity.shape[0]): + n_edges = n_nodes_per_face[face_idx] + for local_idx in range(n_edges): + node_a, node_b = _canonical_half_edge( + face_node_connectivity, face_idx, local_idx, n_edges + ) + + slot = bucket_offset[node_a] + half_edge_slot[slot] = face_idx * n_max_face_nodes + local_idx + end_node[slot] = node_b + bucket_offset[node_a] = slot + 1 + + # The scatter left each entry at its bucket's end, i.e. one slot right of where the + # convention above wants it. One backward pass puts it back. + for n in range(bucket_offset.shape[0] - 1, 0, -1): + bucket_offset[n] = bucket_offset[n - 1] + bucket_offset[0] = 0 + + return half_edge_slot, end_node + + +@njit(cache=True) +def _emit_bucket_edges( + end_node, + half_edge_slot, + bucket_start, + bucket_end, + node_a, + first_edge_idx, + edge_node_connectivity, + face_edge_flat, +): + """Numbers a sorted bucket's unique edges from ``first_edge_idx`` and points each of its + half edges at the edge it belongs to.""" + edge_idx = first_edge_idx - 1 + previous_end_node = INT_FILL_VALUE + + for i in range(bucket_start, bucket_end): + if end_node[i] != previous_end_node: + # Duplicate half edges are adjacent, so a new key starts a new edge + edge_idx += 1 + edge_node_connectivity[edge_idx, 0] = node_a + edge_node_connectivity[edge_idx, 1] = end_node[i] + previous_end_node = end_node[i] + + face_edge_flat[half_edge_slot[i]] = edge_idx + + +@njit(cache=True, parallel=True) +def _build_edge_node_connectivity(face_node_connectivity, n_nodes_per_face, n_node): + """Constructs the ``edge_node_connectivity`` variable, which represents the indices of the two nodes that make up + each edge. Additionally, the ``face_edge_connectivity`` is derived during construction, which represents the + indices of the edges that make up each face. + + Each edge is stored as an ascending ``(node_a, node_b)`` pair, and the edges are numbered in lexicographic + order of that pair. + + Every half edge is bucketed on its lower node, each bucket is sorted on its higher node, and the duplicates + that this makes adjacent are then collapsed into one edge apiece. Parameters ---------- - repopulate : bool, optional - Flag used to indicate if we want to overwrite the existed `edge_node_connectivity` and generate a new - inverse_indices, default is False - """ - - padded_face_nodes = close_face_nodes(face_nodes, n_face, n_max_face_nodes) + face_node_connectivity : np.ndarray + Face Node Connectivity + n_nodes_per_face : np.ndarray + Number of nodes/edges per face + n_node : int + Total number of nodes, used as the number of buckets for the counting sort - # array of empty edge nodes where each entry is a pair of indices - edge_nodes = np.empty((n_face * n_max_face_nodes, 2), dtype=INT_DTYPE) + Returns + ------- + edge_node_connectivity : np.ndarray + Edge Node Connectivity with shape (n_edge, 2) + face_edge_connectivity : np.ndarray + Face Edge Connectivity with shape (n_face, n_max_face_edges) - # first index includes starting node up to non-padded value - edge_nodes[:, 0] = padded_face_nodes[:, :-1].ravel() + """ + # ``np.full`` rather than ``np.full_like``, which would inherit a Fortran-ordered + # prototype's layout and make the flat view below unobtainable + face_edge_connectivity = np.full( + face_node_connectivity.shape, INT_FILL_VALUE, dtype=INT_DTYPE + ) - # second index includes second node up to padded value - edge_nodes[:, 1] = padded_face_nodes[:, 1:].ravel() + n_half_edge = np.sum(n_nodes_per_face) - # sorted edge nodes - edge_nodes.sort(axis=1) + if n_half_edge == 0: + return np.empty((0, 2), dtype=INT_DTYPE), face_edge_connectivity - # unique edge nodes - edge_nodes_unique, inverse_indices = np.unique( - edge_nodes, return_inverse=True, axis=0 + bucket_offset = _count_half_edges_per_node( + face_node_connectivity, n_nodes_per_face, n_node ) - # find all edge nodes that contain a fill value - fill_value_mask = np.logical_or( - edge_nodes_unique[:, 0] == INT_FILL_VALUE, - edge_nodes_unique[:, 1] == INT_FILL_VALUE, + half_edge_slot, end_node = _scatter_half_edges( + face_node_connectivity, n_nodes_per_face, bucket_offset, n_half_edge ) - # all edge nodes that do not contain a fill value - non_fill_value_mask = np.logical_not(fill_value_mask) - edge_nodes_unique = edge_nodes_unique[non_fill_value_mask] + # Sort each bucket and count its unique edges while the bucket is in cache. Buckets are + # disjoint, so this runs one bucket per thread. + unique_per_bucket = np.empty(n_node, dtype=INT_DTYPE) + for n in prange(n_node): + bucket_start = bucket_offset[n] + bucket_end = bucket_offset[n + 1] - # Update inverse_indices accordingly - indices_to_update = np.where(fill_value_mask)[0] - - remove_mask = np.isin(inverse_indices, indices_to_update) - inverse_indices[remove_mask] = INT_FILL_VALUE + _sort_bucket(end_node, half_edge_slot, bucket_start, bucket_end - bucket_start) + unique_per_bucket[n] = _count_unique_in_bucket( + end_node, bucket_start, bucket_end + ) - # Compute the indices where inverse_indices exceeds the values in indices_to_update - indexes = np.searchsorted(indices_to_update, inverse_indices, side="right") - # subtract the corresponding indexes from `inverse_indices` - for i in range(len(inverse_indices)): - if inverse_indices[i] != INT_FILL_VALUE: - inverse_indices[i] -= indexes[i] + # Hand each bucket the edge index its first unique edge takes, so the emit below can run + # one bucket per thread as well + edge_offset = np.empty(n_node + 1, dtype=INT_DTYPE) + n_edge = 0 + for n in range(n_node): + edge_offset[n] = n_edge + n_edge += unique_per_bucket[n] + edge_offset[n_node] = n_edge + + edge_node_connectivity = np.empty((n_edge, 2), dtype=INT_DTYPE) + face_edge_flat = face_edge_connectivity.reshape(-1) + + for n in prange(n_node): + _emit_bucket_edges( + end_node, + half_edge_slot, + bucket_offset[n], + bucket_offset[n + 1], + n, + edge_offset[n], + edge_node_connectivity, + face_edge_flat, + ) - return edge_nodes_unique, inverse_indices, fill_value_mask + return edge_node_connectivity, face_edge_connectivity def _populate_edge_face_connectivity(grid): @@ -252,8 +387,8 @@ def _populate_edge_face_connectivity(grid): @njit(cache=True) def _build_edge_face_connectivity(face_edges, n_nodes_per_face, n_edge): - """Helper for (``edge_face_connectivity``) construction.""" - edge_faces = np.ones(shape=(n_edge, 2), dtype=face_edges.dtype) * INT_FILL_VALUE + """Helper for (``edge_faces``) construction.""" + edge_faces = np.full((n_edge, 2), INT_FILL_VALUE, dtype=INT_DTYPE) for face_idx, (cur_face_edges, n_edges) in enumerate( zip(face_edges, n_nodes_per_face) @@ -274,29 +409,87 @@ def _populate_face_edge_connectivity(grid): and stores it within the internal (``Grid._ds``) and through the attribute (``Grid.face_edge_connectivity``).""" - if ( - "edge_node_connectivity" not in grid._ds - or "inverse_indices" not in grid._ds["edge_node_connectivity"].attrs - ): + if "edge_node_connectivity" not in grid._ds: + # Constructing the edges derives this variable in the same pass _populate_edge_node_connectivity(grid) + return + + # In lieu of an xarray equivalent to `da.compute(a, b)`, we can batch these variables as + # an xarray Dataset and re-extract after graph traversal. + computed = xr.Dataset( + { + "face_nodes": grid.face_node_connectivity.variable, + "n_nodes_per_face": grid.n_nodes_per_face.variable, + "edge_nodes": grid.edge_node_connectivity.variable, + } + ).compute() face_edges = _build_face_edge_connectivity( - grid.edge_node_connectivity.attrs["inverse_indices"], - grid.n_face, - grid.n_max_face_nodes, + computed.face_nodes.data, + computed.n_nodes_per_face.data, + computed.edge_nodes.data, + grid.n_node, ) grid._ds["face_edge_connectivity"] = xr.DataArray( - data=face_edges, + face_edges, dims=ugrid.FACE_EDGE_CONNECTIVITY_DIMS, attrs=ugrid.FACE_EDGE_CONNECTIVITY_ATTRS, ) -def _build_face_edge_connectivity(inverse_indices, n_face, n_max_face_nodes): - """Helper for (``face_edge_connectivity``) construction.""" - inverse_indices = inverse_indices.reshape(n_face, n_max_face_nodes) - return inverse_indices +@njit(cache=True, parallel=True) +def _build_face_edge_connectivity( + face_node_connectivity, n_nodes_per_face, edge_node_connectivity, n_node +): + """Constructs the ``face_edge_connectivity`` variable, which represents the indices of the edges that make up + each face, by looking each face's edges up in an existing ``edge_node_connectivity``. The edges keep the + numbering they arrived with. + + Edges are bucketed on their lower node so that each of a face's edges can be found by a binary search of one + bucket. Edges already in the canonical order that :func:`_build_edge_node_connectivity` emits are bucketed + without being sorted again. + + Parameters + ---------- + face_node_connectivity : np.ndarray + Face Node Connectivity + n_nodes_per_face : np.ndarray + Number of nodes/edges per face + edge_node_connectivity : np.ndarray + Edge Node Connectivity with shape (n_edge, 2), in any order or orientation + n_node : int + Total number of nodes, used as the number of buckets + + Returns + ------- + face_edge_connectivity : np.ndarray + Face Edge Connectivity with shape (n_face, n_max_face_edges). Edges of a face that are absent from + ``edge_node_connectivity`` are left as ``INT_FILL_VALUE``, as are the padding slots of a face with + fewer than ``n_max_face_edges`` edges. + + """ + face_edge_connectivity = np.full( + face_node_connectivity.shape, INT_FILL_VALUE, dtype=INT_DTYPE + ) + + bucket_offset, end_node, edge_id = _build_pair_index(edge_node_connectivity, n_node) + + for face_idx in prange(face_node_connectivity.shape[0]): + n_edges = n_nodes_per_face[face_idx] + for local_idx in range(n_edges): + node_a, node_b = _canonical_half_edge( + face_node_connectivity, face_idx, local_idx, n_edges + ) + face_edge_connectivity[face_idx, local_idx] = _search_bucket( + end_node, + edge_id, + bucket_offset[node_a], + bucket_offset[node_a + 1], + node_b, + ) + + return face_edge_connectivity def _populate_node_face_connectivity(grid): @@ -304,7 +497,7 @@ def _populate_node_face_connectivity(grid): and stores it within the internal (``Grid._ds``) and through the attribute (``Grid.node_face_connectivity``).""" - node_faces, n_max_faces_per_node = _build_node_faces_connectivity( + node_faces, n_max_faces_per_node = _build_node_face_connectivity( grid.face_node_connectivity.values, grid.n_node ) @@ -315,7 +508,7 @@ def _populate_node_face_connectivity(grid): ) -def _build_node_faces_connectivity(face_nodes, n_node): +def _build_node_face_connectivity(face_nodes, n_node): """Builds the `Grid.node_faces_connectivity`: integer DataArray of size (n_node, n_max_faces_per_node) (optional) A DataArray of indices indicating faces that are neighboring each node. @@ -419,7 +612,9 @@ def _populate_face_face_connectivity(grid): """Constructs the UGRID connectivity variable (``face_face_connectivity``) and stores it within the internal (``Grid._ds``) and through the attribute (``Grid.face_face_connectivity``).""" - face_face = _build_face_face_connectivity(grid) + face_face = _build_face_face_connectivity( + grid.edge_face_connectivity.values, grid.n_face, grid.n_max_face_nodes + ) grid._ds["face_face_connectivity"] = xr.DataArray( data=face_face, @@ -428,28 +623,21 @@ def _populate_face_face_connectivity(grid): ) -def _build_face_face_connectivity(grid): - """Returns face-face connectivity.""" - - # Dictionary to store each faces adjacent faces - face_neighbors = {i: [] for i in range(grid.n_face)} +@njit(cache=True) +def _build_face_face_connectivity(edge_face_connectivity, n_face, n_max_face_nodes): + face_face_connectivity = np.full( + (n_face, n_max_face_nodes), INT_FILL_VALUE, INT_DTYPE + ) + face_index_position = np.zeros(n_face, dtype=INT_DTYPE) - # Loop through each edge_face and add to the dictionary every face that shares an edge - for edge_face in grid.edge_face_connectivity.values: - face1, face2 = edge_face - if face1 != INT_FILL_VALUE and face2 != INT_FILL_VALUE: - # Append to each face's dictionary index the opposite face index - face_neighbors[face1].append(face2) - face_neighbors[face2].append(face1) + for edge_faces in edge_face_connectivity: + face_a, face_b = edge_faces + if face_a != INT_FILL_VALUE and face_b != INT_FILL_VALUE: + face_face_connectivity[face_a, face_index_position[face_a]] = face_b + face_index_position[face_a] += 1 - # Convert to an array and pad it with fill values - face_face_conn = list(face_neighbors.values()) - face_face_connectivity = [ - np.pad( - arr, (0, grid.n_max_face_edges - len(arr)), constant_values=INT_FILL_VALUE - ) - for arr in face_face_conn - ] + face_face_connectivity[face_b, face_index_position[face_b]] = face_a + face_index_position[face_b] += 1 return face_face_connectivity diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index 6fa02b069..2dd7af5e3 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -1270,7 +1270,12 @@ def edge_node_connectivity(self) -> xr.DataArray: Connectivity variable representing the indices of nodes (mesh vertices) that define each edge. Each row (i.e., each edge) contains exactly two node indices that define the start and end points of the edge. - The nodes are stored in an arbitrary order. + Constructed edges are stored as ascending node pairs and numbered in lexicographic order of that pair; edges + read from a file keep the order and orientation they were stored in. + + The result is cached after the first access; subsequent calls return the stored value without recomputing it. + Computing edge_node_connectivity always derives face_edge_connectivity as part of the same pass and + overwrites any existing face_edge_connectivity value, regardless of whether one was already present. Returns ------- @@ -1317,6 +1322,11 @@ def face_edge_connectivity(self) -> xr.DataArray: rows containing fewer than :py:attr:`~uxarray.Grid.n_max_face_edges` indices are padded with the fill value defined in :py:attr:`~uxarray.constants.INT_FILL_VALUE`. + The result is cached after the first access; subsequent calls return the stored value without recomputing it. + If edge_node_connectivity has not yet been computed, it is derived together with face_edge_connectivity in + the same pass. If edge_node_connectivity is already present, face_edge_connectivity is instead derived + independently from the existing connectivity data. + Returns ------- face_edge_connectivity : :py:class:`xarray.DataArray` diff --git a/uxarray/grid/utils.py b/uxarray/grid/utils.py index 9287bc325..6510c85d0 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -2,7 +2,7 @@ import xarray as xr from numba import njit -from uxarray.constants import INT_FILL_VALUE +from uxarray.constants import INT_DTYPE, INT_FILL_VALUE @njit(cache=True) @@ -443,3 +443,255 @@ def setter(self, value): self._ds[key] = value return setter + + +# Bucket sorting and searching for the counting sorts in ``uxarray.grid.connectivity``, which +# bucket half edges or edges on one of their two nodes and then order each bucket by the other. +# Nothing below knows about meshes: a bucket is a contiguous ``[bucket_start, bucket_end)`` slice +# of a key array, carrying an equally long ``payload`` array that every reordering moves in step +# so the two stay aligned. +# +# ``_sort_bucket`` orders one bucket and is the entry point for the sorts; the kernels beneath it +# are chosen by bucket size and are exposed only for testing. ``_build_pair_index`` runs the whole +# count/scatter/sort sequence for a caller starting from an ``(n, 2)`` array, and +# ``_search_bucket`` is the lookup that index is built for. +# +# NOTE: these are inlined into ``cache=True`` kernels in another module, and numba stamps its +# cache against the defining file alone, so editing them does not invalidate a caller's cached +# object. Clear ``uxarray/grid/__pycache__/*.nbi *.nbc`` after changing anything here. + +# Smallest bucket worth watching for pathological input. A bucket of ``size`` holds at most +# ``size * (size - 1) / 2`` inversions, so at or below this size it cannot exceed the shift +# budget below and the bookkeeping would never pay for itself +MIN_ADAPTIVE_SORT_SIZE = 16 + +# Shifts per edge an insertion sort may spend on a bucket before it is abandoned for a heap +# sort. Insertion sort costs ``O(size + shifts)``, so a constant budget per edge keeps the +# adaptive path linear while leaving ample room for the near-sorted input it is chosen for +MAX_SHIFTS_PER_EDGE = 8 + + +@njit(cache=True) +def _sort_bucket(end_node, payload, bucket_start, size): + """Orders one bucket by ``end_node``, picking the sort that suits its size.""" + if size > MIN_ADAPTIVE_SORT_SIZE: + # Large enough that a bad ordering would be worth catching, which only a + # collapsed pole or a similarly degenerate node reaches + _adaptive_sort_bucket(end_node, payload, bucket_start, size) + elif size > 1: + _insertion_sort_bucket(end_node, payload, bucket_start, size) + + +@njit(cache=True) +def _sift_down(end_node, payload, bucket_start, root, size): + """Restores the max-heap property at ``root`` for a bucket keyed on ``end_node``.""" + while True: + child = 2 * root + 1 + if child >= size: + break + + if ( + child + 1 < size + and end_node[bucket_start + child] < end_node[bucket_start + child + 1] + ): + child += 1 + + if end_node[bucket_start + root] >= end_node[bucket_start + child]: + break + + end_node[bucket_start + root], end_node[bucket_start + child] = ( + end_node[bucket_start + child], + end_node[bucket_start + root], + ) + payload[bucket_start + root], payload[bucket_start + child] = ( + payload[bucket_start + child], + payload[bucket_start + root], + ) + root = child + + +@njit(cache=True) +def _heap_sort_bucket(end_node, payload, bucket_start, size): + """Sorts a bucket by ``end_node`` in place, in ``O(size * log(size))`` and without + scratch space, for the rare bucket an insertion sort cannot finish cheaply.""" + for root in range(size // 2 - 1, -1, -1): + _sift_down(end_node, payload, bucket_start, root, size) + + for end in range(size - 1, 0, -1): + end_node[bucket_start], end_node[bucket_start + end] = ( + end_node[bucket_start + end], + end_node[bucket_start], + ) + payload[bucket_start], payload[bucket_start + end] = ( + payload[bucket_start + end], + payload[bucket_start], + ) + _sift_down(end_node, payload, bucket_start, 0, end) + + +@njit(cache=True) +def _insertion_sort_bucket(end_node, payload, bucket_start, size): + """Sorts a bucket by ``end_node`` in place, in ``O(size + inversions)``.""" + for i in range(bucket_start + 1, bucket_start + size): + key = end_node[i] + key_payload = payload[i] + + j = i - 1 + while j >= bucket_start and end_node[j] > key: + end_node[j + 1] = end_node[j] + payload[j + 1] = payload[j] + j -= 1 + end_node[j + 1] = key + payload[j + 1] = key_payload + + +@njit(cache=True) +def _adaptive_sort_bucket(end_node, payload, bucket_start, size): + """Sorts a large bucket by ``end_node`` in place, insertion sorting it unless it turns out + to be badly ordered, in which case the partial work is abandoned for a heap sort. + + This is ``_insertion_sort_bucket``'s loop with a shift meter around it. The duplication is + deliberate: metering every bucket instead of only the large ones measured ~5% slower + end-to-end, because typical buckets hold a handful of edges and the per-element bookkeeping + is a real fraction of that work. Keep the two in sync rather than merging them. + """ + budget = MAX_SHIFTS_PER_EDGE * size + shifts = 0 + + for i in range(bucket_start + 1, bucket_start + size): + key = end_node[i] + key_payload = payload[i] + + j = i - 1 + while j >= bucket_start and end_node[j] > key: + end_node[j + 1] = end_node[j] + payload[j + 1] = payload[j] + j -= 1 + end_node[j + 1] = key + payload[j + 1] = key_payload + + shifts += i - 1 - j + if shifts > budget: + _heap_sort_bucket(end_node, payload, bucket_start, size) + return + + +@njit(cache=True) +def _count_unique_in_bucket(end_node, bucket_start, bucket_end): + """Number of distinct keys in an already sorted bucket, where equal keys are adjacent.""" + n_unique = 0 + previous_end_node = INT_FILL_VALUE + + for i in range(bucket_start, bucket_end): + if end_node[i] != previous_end_node: + n_unique += 1 + previous_end_node = end_node[i] + + return n_unique + + +@njit(cache=True) +def _is_lexicographically_sorted(pairs): + """Whether every row of an ``(n, 2)`` array is ascending and the rows are themselves in + nondecreasing lexicographic order.""" + previous_low = INT_FILL_VALUE + previous_high = INT_FILL_VALUE + + for i in range(pairs.shape[0]): + low = pairs[i, 0] + high = pairs[i, 1] + + if low > high: + return False + if low < previous_low or (low == previous_low and high < previous_high): + return False + + previous_low = low + previous_high = high + + return True + + +@njit(cache=True) +def _count_pairs_per_bucket(pairs, n_bucket): + """Bucket offsets keyed on each row's lower value: bucket ``a`` will occupy + ``[bucket_offset[a], bucket_offset[a + 1])``.""" + bucket_offset = np.zeros(n_bucket + 1, dtype=INT_DTYPE) + + for i in range(pairs.shape[0]): + bucket_offset[min(pairs[i, 0], pairs[i, 1]) + 1] += 1 + + for a in range(n_bucket): + bucket_offset[a + 1] += bucket_offset[a] + + return bucket_offset + + +@njit(cache=True) +def _build_pair_index(pairs, n_bucket): + """Indexes an ``(n, 2)`` array of integer pairs so a pair can be looked up by value. + + Each row is bucketed on its lower value and each bucket ordered by its higher value, which is + what lets :func:`_search_bucket` find a row with a single binary search. Rows that are already + canonically ordered are indexed without being sorted again. + + Returns ``(bucket_offset, high, row)``: bucket ``a`` occupies + ``[bucket_offset[a], bucket_offset[a + 1])``, ``high`` holds each entry's higher value, and + ``row`` the index of ``pairs`` it came from. ``high`` is a copy rather than a column view, both + to keep the search off a strided array and so that either path below returns the same arrays. + """ + bucket_offset = _count_pairs_per_bucket(pairs, n_bucket) + + n_pair = pairs.shape[0] + high = np.empty(n_pair, dtype=INT_DTYPE) + row = np.empty(n_pair, dtype=INT_DTYPE) + + if _is_lexicographically_sorted(pairs): + # Already grouped by lower value, ascending within each group, so the buckets are the + # runs the counting pass just measured and no sorting is needed + for i in range(n_pair): + high[i] = pairs[i, 1] + row[i] = i + return bucket_offset, high, row + + for i in range(n_pair): + pair_low = pairs[i, 0] + pair_high = pairs[i, 1] + if pair_low > pair_high: + pair_low, pair_high = pair_high, pair_low + + slot = bucket_offset[pair_low] + high[slot] = pair_high + row[slot] = i + bucket_offset[pair_low] = slot + 1 + + # The scatter left each entry at its bucket's end, one slot right of where the convention + # above wants it. One backward pass puts it back. + for a in range(n_bucket, 0, -1): + bucket_offset[a] = bucket_offset[a - 1] + bucket_offset[0] = 0 + + for a in range(n_bucket): + bucket_start = bucket_offset[a] + _sort_bucket(high, row, bucket_start, bucket_offset[a + 1] - bucket_start) + + return bucket_offset, high, row + + +@njit(cache=True) +def _search_bucket(high, row, bucket_start, bucket_end, key): + """The ``row`` entry whose key is ``key`` within a sorted bucket, or ``INT_FILL_VALUE`` + when the bucket does not hold it.""" + low = bucket_start + stop = bucket_end + + while low < stop: + mid = (low + stop) // 2 + if high[mid] < key: + low = mid + 1 + else: + stop = mid + + if low < bucket_end and high[low] == key: + return row[low] + return INT_FILL_VALUE