Skip to content

Optimized connectivity: Simultaneous face_edges and edge_nodes - #1560

Open
cmdupuis3 wants to merge 19 commits into
UXARRAY:mainfrom
cmdupuis3:cmd/merge-OFE
Open

Optimized connectivity: Simultaneous face_edges and edge_nodes#1560
cmdupuis3 wants to merge 19 commits into
UXARRAY:mainfrom
cmdupuis3:cmd/merge-OFE

Conversation

@cmdupuis3

@cmdupuis3 cmdupuis3 commented Jul 10, 2026

Copy link
Copy Markdown
Collaborator

Would close #1138, #1196

Related to #1180

Supercedes #1195

Overview

This set of changes optimizes face_edge, edge_node, and face_face connectivity. face_edge and edge_node connectivity are combined into one routine, while face_face is optimized stand-alone.

PR Checklist

General

  • An issue is linked created and linked
  • Add appropriate labels
  • Filled out Overview and Expected Usage (if applicable) sections

Testing

  • Adequate tests are created if there is new functionality
  • Tests cover all possible logical paths in your function
  • Tests are not too basic (such as simply calling a function and nothing else)

Documentation

  • Docstrings have been added to all new functions
  • Docstrings have updated with any function changes
  • Internal functions have a preceding underscore (_) and have been added to docs/internal_api/index.rst

@cmdupuis3 cmdupuis3 self-assigned this Jul 10, 2026
@cmdupuis3 cmdupuis3 added the improvement Improvements on existing features or infrastructure label Jul 10, 2026
@cmdupuis3
cmdupuis3 requested a review from hongyuchen1030 July 10, 2026 22:41
@cmdupuis3

Copy link
Copy Markdown
Collaborator Author

I think all the hard parts of the merge are done.

At this point, the only failures seems to be sorting issues.

@cmdupuis3 cmdupuis3 added scalability Related to scalability & performance efforts and removed improvement Improvements on existing features or infrastructure labels Jul 10, 2026
@hongyuchen1030

Copy link
Copy Markdown
Contributor

I think all the hard parts of the merge are done.

At this point, the only failures seems to be sorting issues.

@cmdupuis3 Thanks for your work, I will review it as soon as possible.

And do you have any idea why the CIs are all failing?

@cmdupuis3

Copy link
Copy Markdown
Collaborator Author

The CI is failing because the new algorithm returns the results in a different order.

I was thinking we can add some sorting mechanism, at least for legacy behavior. Phillip was evidently aware of this issue too. It depends on if you need to support that... Personally I'd be okay with just changing it, but I think you would be a better judge of the situation.

Comment thread uxarray/grid/connectivity.py Outdated
pass

edge_nodes, face_edges = _build_edge_node_connectivity(
grid.face_node_connectivity.values, grid.n_nodes_per_face.values

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please think about if you can get rid of .values calls to be chunked-array compatible.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On it, I think Philip was trying to scalarize here but it introduced this regression.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So the best I could come up with is

    face_nodes, n_nodes_per_face = dask.compute(
        grid.face_node_connectivity.data, grid.n_nodes_per_face.data
    )

    edge_nodes, face_edges = _build_edge_node_connectivity(
        face_nodes, n_nodes_per_face, grid.n_node
    )

...but I'm unclear on if we're committing to dask or not. Other parts of the repo imply that dask is still sort of an optional load, and doing this would pretty much require a dask import for most grids, unless we have some kind of numpy/dask conditional.

@cmdupuis3 cmdupuis3 Jul 29, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Alright, I hacked it a bit to get this:

    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
    )

@cmdupuis3

Copy link
Copy Markdown
Collaborator Author

Altering njit and .values here runs the risk of conflicting with #1583. I'm thinking we can scope it so that .values cleanup in connectivity.py is allowed on this PR and everywhere except connectivity.py belongs to #1583.

cmdupuis3 and others added 4 commits July 22, 2026 21:46
The optimized edge builder deduped half edges with a numba hash map, which
numbered edges in first-encounter order. Edges had previously been numbered
lexicographically by their (min_node, max_node) pair, as a side effect of the
np.unique(..., axis=0) the hash map replaced.

Global edge index is a public identity: it indexes edge_lon/edge_lat, edge
centered data variables, and edge_node_distances, so renumbering silently
re-pairs user data with different physical edges. It also broke the five
TestQuadHexagon connectivity tests, which assert on edge_node, face_edge,
node_edge, edge_face and face_face -- all the same renumbering cascading
through the derived connectivities.

Sort as the dedup mechanism instead of hashing. Node indices are dense
integers in [0, n_node), so a counting sort buckets the half edges by their
first node without any comparisons, and sorting each bucket by its second node
leaves the duplicates adjacent -- the dedup then falls out of the same walk.
Buckets hold one entry per edge incident to a node, so on a real mesh they are
tiny (node degree, typically under ten) and an insertion sort finishes them.
A bucket above MAX_INSERTION_SORT_SIZE is heap sorted so that a degenerate
mesh cannot degrade the build quadratically; np.argsort is deliberately not
used there, as numba's implementation degrades badly on structured input.

Half edges are identified throughout by their flat face_node_connectivity
index, which is also the face_edge_connectivity slot they are written back to,
so the sort needs a single permutation array and no mapping back.

This is faster and leaner than the hash map it replaces. On a synthetic one
million face quad mesh, measured by peak RSS rather than tracemalloc, which
does not observe numba's typed dict allocations:

    dict build     397.3 ms    239.5 MB
    bucket sort     85.1 ms     91.6 MB

The five legacy tests now pass unchanged. Adds order invariant coverage for
the canonical ordering and the face_edge positional contract, plus high degree
nodes either side of the insertion sort threshold.

Also casts n_nodes_per_face back to INT_DTYPE, so that the builder is not
compiled a second time for int64, and restores a blank line dropped between
two top level functions.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@cmdupuis3 cmdupuis3 added the run-benchmark Run ASV benchmark workflow label Jul 29, 2026
@github-actions

github-actions Bot commented Jul 29, 2026

Copy link
Copy Markdown

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change Before [852a0c0] After [ca770ae] Ratio Benchmark (Parameter)
- 588±7μs 12.0±0.2μs 0.02 bench_connectivity.Connectivity.time_n_nodes_per_face('480km')
- 520M 336M 0.65 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
- 631M 335M 0.53 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
- 833±4ms 1.09±0.01ms 0 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
- 53.3±0.3ms 592±10μs 0.01 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
- 451M 330M 0.73 mpas_ocean.Gradient.peakmem_gradient('480km')

Benchmarks that have stayed the same:

Change Before [852a0c0] After [ca770ae] Ratio Benchmark (Parameter)
10.3±0.09μs 10.7±0.3μs 1.04 bench_connectivity.Connectivity.time_edge_face('120km')
10.9±0.1μs 10.7±0.2μs 0.98 bench_connectivity.Connectivity.time_edge_face('480km')
10.3±0.06μs 10.7±0.2μs 1.04 bench_connectivity.Connectivity.time_edge_node('120km')
11.0±0.1μs 11.1±0.04μs 1.01 bench_connectivity.Connectivity.time_edge_node('480km')
10.4±0.1μs 10.9±0.2μs 1.04 bench_connectivity.Connectivity.time_face_edge('120km')
11.0±0.1μs 11.1±0.1μs 1 bench_connectivity.Connectivity.time_face_edge('480km')
10.3±0.2μs 10.7±0.09μs 1.05 bench_connectivity.Connectivity.time_face_face('120km')
10.9±0.1μs 11.1±0.3μs 1.03 bench_connectivity.Connectivity.time_face_face('480km')
21.0±0.1μs 21.4±0.2μs 1.02 bench_connectivity.Connectivity.time_face_node('120km')
21.8±0.3μs 22.2±0.3μs 1.02 bench_connectivity.Connectivity.time_face_node('480km')
11.4±0.3μs 12.4±0.1μs 1.09 bench_connectivity.Connectivity.time_n_nodes_per_face('120km')
10.4±0.09μs 10.7±0.1μs 1.03 bench_connectivity.Connectivity.time_node_edge('120km')
10.6±0.4μs 11.2±0.3μs 1.05 bench_connectivity.Connectivity.time_node_edge('480km')
10.4±0.2μs 10.6±0.2μs 1.02 bench_connectivity.Connectivity.time_node_face('120km')
10.9±0.2μs 11.0±0.09μs 1.01 bench_connectivity.Connectivity.time_node_face('480km')
334M 334M 1 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
363M 365M 1.01 face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
22.4±0.09μs 23.6±0.1μs 1.05 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
9.97±0.1μs 10.0±0.03μs 1.01 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
10.2±0.01ms 10.5±0.06ms 1.03 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
2.11±0.02ms 2.13±0.01ms 1.01 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
940±2ns 944±7ns 1 geometry_kernels.AccucrossKernels.time_accucross
2.48±0.03μs 2.46±0.01μs 0.99 geometry_kernels.AccucrossKernels.time_accucross_pair
298±4ns 294±0.4ns 0.98 geometry_kernels.EFTPrimitives.time_acc_sqrt_re
280±0.1ns 280±1ns 1 geometry_kernels.EFTPrimitives.time_diff_of_products
233±1ns 234±0.6ns 1 geometry_kernels.EFTPrimitives.time_two_prod
238±0.9ns 236±0.4ns 0.99 geometry_kernels.EFTPrimitives.time_two_sum
1.21±0.02μs 1.20±0μs 0.99 geometry_kernels.GCAConstLatIntersection.time_accux_constlat_kernel
885±1ns 863±7ns 0.97 geometry_kernels.GCAConstLatIntersection.time_gca_const_lat_intersection
1.61±0μs 1.65±0.03μs 1.03 geometry_kernels.GCAConstLatIntersection.time_try_gca_const_lat_intersection
1.33±0μs 1.34±0.01μs 1 geometry_kernels.GCAGCAIntersection.time_accux_gca_kernel
1.12±0μs 1.14±0.02μs 1.01 geometry_kernels.GCAGCAIntersection.time_gca_gca_intersection
1.91±0.09μs 1.85±0μs 0.97 geometry_kernels.GCAGCAIntersection.time_try_gca_gca_intersection
50.3±0.4μs 49.6±0.8μs 0.99 geometry_kernels.OrientPredicates.time_on_minor_arc
523±20ns 508±7ns 0.97 geometry_kernels.OrientPredicates.time_orient3d_on_sphere
2.71±0.1ms 2.58±0.01ms 0.95 geometry_samebody.SameBodyConstLat.time_accux_dispatch
1.17±0ms 1.17±0ms 1 geometry_samebody.SameBodyConstLat.time_accux_kernel
1.83±0.1ms 1.72±0ms 0.94 geometry_samebody.SameBodyConstLat.time_fp64_dispatch
146±0.4μs 145±0.2μs 1 geometry_samebody.SameBodyConstLat.time_fp64_kernel
33.4±0.9ms 32.6±0.4ms 0.98 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_dispatch
10.2±0.02ms 10.2±0.04ms 1.01 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_kernel
26.8±0.3ms 26.4±0.03ms 0.98 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_dispatch
4.86±0.01ms 4.90±0.01ms 1.01 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_kernel
813±5ms 806±5ms 0.99 import.Imports.timeraw_import_uxarray
900±3ns 940±20ns 1.05 mpas_ocean.CheckNorm.time_check_norm('120km')
885±10ns 884±2ns 1 mpas_ocean.CheckNorm.time_check_norm('480km')
14.5±0.1μs 15.4±0.3μs 1.06 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
14.5±0.06μs 15.3±0.06μs 1.06 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
4.83±0.01ms 5.18±0.02ms 1.07 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('120km')
3.37±0.01ms 3.48±0.01ms 1.03 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('480km')
3.52±0.08s 3.43±0.01s 0.98 mpas_ocean.ConstructFaceLatLon.time_welzl('120km')
223±0.5ms 223±2ms 1 mpas_ocean.ConstructFaceLatLon.time_welzl('480km')
18.1±0.01ms 18.2±0.02ms 1 mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
1.04±0.01ms 1.04±0.02ms 1 mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
10.6±0.01ms 10.6±0.04ms 1 mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')
702±7μs 704±9μs 1 mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
698±1ms 704±2ms 1.01 mpas_ocean.CrossSections.time_const_lat('120km', 1)
355±3ms 351±2ms 0.99 mpas_ocean.CrossSections.time_const_lat('120km', 2)
183±0.7ms 182±2ms 0.99 mpas_ocean.CrossSections.time_const_lat('120km', 4)
533±2ms 540±5ms 1.01 mpas_ocean.CrossSections.time_const_lat('480km', 1)
272±0.6ms 273±2ms 1 mpas_ocean.CrossSections.time_const_lat('480km', 2)
139±0.5ms 141±0.9ms 1.01 mpas_ocean.CrossSections.time_const_lat('480km', 4)
24.4±0.07ms 24.1±0.2ms 0.99 mpas_ocean.DualMesh.time_dual_mesh_construction('120km')
3.12±0.04ms 3.13±0.04ms 1 mpas_ocean.DualMesh.time_dual_mesh_construction('480km')
946±10ms 944±1ms 1 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
52.4±2ms 56.5±0.5ms 1.08 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
83.4±0.4ms 82.7±0.6ms 0.99 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
5.55±0.07ms 5.76±0.2ms 1.04 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
350M 350M 1 mpas_ocean.Gradient.peakmem_gradient('120km')
175±0.9ms 174±1ms 0.99 mpas_ocean.Gradient.time_gradient('120km')
12.1±0.05ms 12.2±0.08ms 1.01 mpas_ocean.Gradient.time_gradient('480km')
227±2μs 225±1μs 0.99 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
130±0.7μs 131±0.5μs 1.01 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
349M 349M 1 mpas_ocean.Integrate.peakmem_integrate('120km')
329M 328M 1 mpas_ocean.Integrate.peakmem_integrate('480km')
218±2μs 214±1μs 0.98 mpas_ocean.Integrate.time_integrate('120km')
194±0.2μs 199±0.7μs 1.03 mpas_ocean.Integrate.time_integrate('480km')
181±4ms 183±2ms 1.01 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
187±1ms 183±0.6ms 0.98 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
183±2ms 182±0.7ms 1 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
13.5±0.1ms 14.0±0.3ms 1.03 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
13.4±0.07ms 14.0±0.3ms 1.04 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
13.6±0.08ms 13.7±0.1ms 1.01 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
346±2μs 347±2μs 1 mpas_ocean.PointInPolygon.time_face_search_lonlat('120km')
349±1μs 351±1μs 1.01 mpas_ocean.PointInPolygon.time_face_search_lonlat('480km')
328±3μs 331±1μs 1.01 mpas_ocean.PointInPolygon.time_face_search_xyz('120km')
329±3μs 334±2μs 1.02 mpas_ocean.PointInPolygon.time_face_search_xyz('480km')
243±3ms 242±0.6ms 0.99 mpas_ocean.RemapDownsample.time_bilinear_remapping
283±2ms 286±1ms 1.01 mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
4.24±0.02ms 4.26±0.03ms 1 mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
1.42±0.01s 1.41±0.01s 0.99 mpas_ocean.RemapUpsample.time_bilinear_remapping
35.9±0.7ms 36.3±0.2ms 1.01 mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
9.23±0.1ms 9.44±0.1ms 1.02 mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
25.2±0.2ms 26.3±0.4ms 1.04 mpas_ocean.ZonalAverage.time_zonal_average('120km')
5.63±0.04ms 5.69±0.05ms 1.01 mpas_ocean.ZonalAverage.time_zonal_average('480km')
326M 326M 1 quad_hexagon.QuadHexagon.peakmem_open_dataset
324M 324M 1 quad_hexagon.QuadHexagon.peakmem_open_grid
6.95±0.1ms 6.89±0.2ms 0.99 quad_hexagon.QuadHexagon.time_open_dataset
5.84±0.07ms 5.86±0.08ms 1 quad_hexagon.QuadHexagon.time_open_grid

@cmdupuis3
cmdupuis3 marked this pull request as ready for review July 29, 2026 19:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-benchmark Run ASV benchmark workflow scalability Related to scalability & performance efforts

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Poor performancewhen constructing edge_node, face_edge and face_face connectivity Improve performance of edge_node_connectivity construction

4 participants