diff --git a/benchmarks/bench_connectivity.py b/benchmarks/bench_connectivity.py index de1bf3895..f4c05c637 100644 --- a/benchmarks/bench_connectivity.py +++ b/benchmarks/bench_connectivity.py @@ -54,6 +54,9 @@ def teardown(self, resolution, *args, **kwargs): class Connectivity(GridBenchmark): + 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 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..b28606754 100644 --- a/uxarray/grid/connectivity.py +++ b/uxarray/grid/connectivity.py @@ -1,9 +1,14 @@ 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 ( + MIN_ADAPTIVE_SORT_SIZE, + _adaptive_sort_bucket, + _insertion_sort_bucket, +) def close_face_nodes(face_node_connectivity, n_face, n_max_face_nodes): @@ -125,8 +130,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 +146,170 @@ 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.sizes: + # TODO: raise a warning or exception? + pass - edge_node_attrs = ugrid.EDGE_NODE_CONNECTIVITY_ATTRS - edge_node_attrs["inverse_indices"] = inverse_indices + # HACK: this is lieu of an xarray equivalent to `da.compute(a, b)` + 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, 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. 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) - - # 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) + 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 - # first index includes starting node up to non-padded value - edge_nodes[:, 0] = padded_face_nodes[:, :-1].ravel() + 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) - # second index includes second node up to padded value - edge_nodes[:, 1] = padded_face_nodes[:, 1:].ravel() + """ - # sorted edge nodes - edge_nodes.sort(axis=1) + n_face, n_max_face_nodes = face_node_connectivity.shape - # unique edge nodes - edge_nodes_unique, inverse_indices = np.unique( - edge_nodes, return_inverse=True, axis=0 - ) - # 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, + # Keep track of face_edge_connectivity + face_edge_connectivity = np.full_like( + face_node_connectivity, INT_FILL_VALUE, dtype=INT_DTYPE ) - # 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] - - # Update inverse_indices accordingly - indices_to_update = np.where(fill_value_mask)[0] + n_half_edge = np.sum(n_nodes_per_face) - remove_mask = np.isin(inverse_indices, indices_to_update) - inverse_indices[remove_mask] = INT_FILL_VALUE + if n_half_edge == 0: + return np.empty((0, 2), dtype=INT_DTYPE), face_edge_connectivity - # 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] - - return edge_nodes_unique, inverse_indices, fill_value_mask + # Count how many half edges fall into each ``start_node`` bucket, then prefix sum so that + # ``bucket_bounds[a]`` is where bucket ``a`` starts + bucket_bounds = np.zeros(n_node + 1, dtype=INT_DTYPE) + for face_idx in range(n_face): + n_edges = n_nodes_per_face[face_idx] + for current_node in range(n_edges): + start_node = face_node_connectivity[face_idx, current_node] + end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] + bucket_bounds[min(start_node, end_node) + 1] += 1 + for i in range(n_node): + bucket_bounds[i + 1] += bucket_bounds[i] + + # Scatter the half edges into their buckets. This advances each entry of + # ``bucket_bounds`` to the *end* of its bucket, so afterwards bucket ``a`` spans + # ``bucket_bounds[a - 1]`` up to ``bucket_bounds[a]``, with bucket 0 starting at 0 + order = np.empty(n_half_edge, dtype=INT_DTYPE) + end_node_keys = np.empty(n_half_edge, dtype=INT_DTYPE) + for face_idx in range(n_face): + n_edges = n_nodes_per_face[face_idx] + for current_node in range(n_edges): + start_node = face_node_connectivity[face_idx, current_node] + end_node = face_node_connectivity[face_idx, (current_node + 1) % n_edges] + + if start_node > end_node: + end_node, start_node = start_node, end_node + + slot = bucket_bounds[start_node] + order[slot] = face_idx * n_max_face_nodes + current_node + end_node_keys[slot] = end_node + bucket_bounds[start_node] = slot + 1 + + # Sort each bucket by ``node_b`` 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_bounds[n - 1] if n > 0 else 0 + bucket_end = bucket_bounds[n] + + size = bucket_end - bucket_start + 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_keys, order, bucket_start, size) + elif size > 1: + _insertion_sort_bucket(end_node_keys, order, bucket_start, size) + + n_unique = 0 + prev_b = INT_FILL_VALUE + for i in range(bucket_start, bucket_end): + if end_node_keys[i] != prev_b: + n_unique += 1 + prev_b = end_node_keys[i] + + unique_per_bucket[n] = n_unique + + 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 + + # Duplicate half edges are now adjacent, so a single walk assigns each unique edge its + # index and populates the face edge connectivity. + edge_node_connectivity = np.empty((n_edge, 2), dtype=INT_DTYPE) + + for n in prange(n_node): + bucket_start = bucket_bounds[n - 1] if n > 0 else 0 + bucket_end = bucket_bounds[n] + + edge_idx = edge_offset[n] - 1 + prev_b = INT_FILL_VALUE + + for i in range(bucket_start, bucket_end): + flat_idx = order[i] + end_node = end_node_keys[i] + + if end_node != prev_b: + # Only store unique edges + edge_idx += 1 + edge_node_connectivity[edge_idx, 0] = n + edge_node_connectivity[edge_idx, 1] = end_node + prev_b = end_node + + face_edge_connectivity[ + flat_idx // n_max_face_nodes, flat_idx % n_max_face_nodes + ] = edge_idx + + return edge_node_connectivity, face_edge_connectivity def _populate_edge_face_connectivity(grid): @@ -252,8 +329,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 +351,10 @@ 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 - ): - _populate_edge_node_connectivity(grid) - - face_edges = _build_face_edge_connectivity( - grid.edge_node_connectivity.attrs["inverse_indices"], - grid.n_face, - grid.n_max_face_nodes, - ) + # TODO: Check if "edge_edge_connectivity" is already present - grid._ds["face_edge_connectivity"] = xr.DataArray( - data=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 + if "edge_node_connectivity" not in grid._ds: + _populate_edge_node_connectivity(grid) def _populate_node_face_connectivity(grid): @@ -304,7 +362,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 +373,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 +477,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 +488,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)} - - # 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) - - # 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 - ] +@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) + + 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 + + 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/utils.py b/uxarray/grid/utils.py index 9287bc325..f45288eab 100644 --- a/uxarray/grid/utils.py +++ b/uxarray/grid/utils.py @@ -443,3 +443,107 @@ def setter(self, value): self._ds[key] = value return setter + + +# Bucket sorts for the counting sort in ``connectivity._build_edge_node_connectivity``. Each +# sorts one contiguous ``[bucket_start, bucket_start + size)`` slice of ``node_b`` in place, +# applying the same permutation to ``order`` so the two stay aligned. + +# 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 _sift_down(node_b, order, bucket_start, root, size): + """Restores the max-heap property at ``root`` for a bucket keyed on ``node_b``.""" + while True: + child = 2 * root + 1 + if child >= size: + break + + if ( + child + 1 < size + and node_b[bucket_start + child] < node_b[bucket_start + child + 1] + ): + child += 1 + + if node_b[bucket_start + root] >= node_b[bucket_start + child]: + break + + node_b[bucket_start + root], node_b[bucket_start + child] = ( + node_b[bucket_start + child], + node_b[bucket_start + root], + ) + order[bucket_start + root], order[bucket_start + child] = ( + order[bucket_start + child], + order[bucket_start + root], + ) + root = child + + +@njit(cache=True) +def _heap_sort_bucket(node_b, order, bucket_start, size): + """Sorts a bucket by ``node_b`` 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(node_b, order, bucket_start, root, size) + + for end in range(size - 1, 0, -1): + node_b[bucket_start], node_b[bucket_start + end] = ( + node_b[bucket_start + end], + node_b[bucket_start], + ) + order[bucket_start], order[bucket_start + end] = ( + order[bucket_start + end], + order[bucket_start], + ) + _sift_down(node_b, order, bucket_start, 0, end) + + +@njit(cache=True) +def _insertion_sort_bucket(node_b, order, bucket_start, size): + """Sorts a bucket by ``node_b`` in place, in ``O(size + inversions)``.""" + for i in range(bucket_start + 1, bucket_start + size): + key = node_b[i] + flat_idx = order[i] + + j = i - 1 + while j >= bucket_start and node_b[j] > key: + node_b[j + 1] = node_b[j] + order[j + 1] = order[j] + j -= 1 + node_b[j + 1] = key + order[j + 1] = flat_idx + + +@njit(cache=True) +def _adaptive_sort_bucket(node_b, order, bucket_start, size): + """Sorts a large bucket by ``node_b`` 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. + """ + budget = MAX_SHIFTS_PER_EDGE * size + shifts = 0 + + for i in range(bucket_start + 1, bucket_start + size): + key = node_b[i] + flat_idx = order[i] + + j = i - 1 + while j >= bucket_start and node_b[j] > key: + node_b[j + 1] = node_b[j] + order[j + 1] = order[j] + j -= 1 + node_b[j + 1] = key + order[j + 1] = flat_idx + + shifts += i - 1 - j + if shifts > budget: + _heap_sort_bucket(node_b, order, bucket_start, size) + return