DRAFT: Optimize Connectivity Construction

[philipc2]

Closes #1196 #1197

Overview

• Optimizes the construction of the following connectivity variables
  • edge_node_connectivity
  • face_edge_connectivity
  • face_face_connectivity
• Cleans up docstrings and setters in Grid

All timings below were taken on a single NCAR Derecho CPU node.

• AMD EPYC™ 7763 Milan processors
• Dual-socket nodes, 64 cores per socket
• 256 GB DDR4 memory per node

Connectivity Timing (Before)

Resolution	Nodes	Faces	Edges	edge_node & face_edge	face_face	edge_face
30km	1,310,720	655,362	1,966,080	9.37s	14.81s	0.05s
15km	5,242,880	2,621,442	7,864,320	39.99s	58.22s	0.23s
7.5km	5,242,880	2,621,442	7,864,320	99.31s	223.92s	0.50s

Connectivity Timings (After)

Resolution	Nodes	Faces	Edges	edge_node & face_edge	face_face	edge_face (unchanged)
30km	1,310,720	655,362	1,966,080	0.54	0.032s	0.05s
15km	5,242,880	2,621,442	7,864,320	2.79	0.123s	0.23s
7.5km	5,242,880	2,621,442	7,864,320	11.41	0.293s	0.50s

Connectivity Speedup

edge_node and face_edge connectivity

• 8.7x to 17.35x speedup

face_face

• 462x to 764x speedup

[cmdupuis3]

@hongyuchen1030 Hi Hongyu,

I'm trying to bring this PR up to date.

Currently, merging it into main is complex with significant merge conflicts. I've been working through some of them, but since I'm new I'm not completely sure what to prioritize.

I ran a speed comparison for the edge_node and face_face connectivity changes in Claude, which found similar speedups to what Philip was originally finding, so I think the core changes here are probably still valuable.

I'm still getting some value errors in the connectivity tests though.

[hongyuchen1030]

@hongyuchen1030 Hi Hongyu,

I'm trying to bring this PR up to date.

Currently, merging it into main is complex with significant merge conflicts. I've been working through some of them, but since I'm new I'm not completely sure what to prioritize.

I ran a speed comparison for the edge_node and face_face connectivity changes in Claude, which found similar speedups to what Philip was originally finding, so I think the core changes here are probably still valuable.

I'm still getting some value errors in the connectivity tests though.

Hi @cmdupuis3,

Thanks for looking into this PR. I think one of the easiest ways to speed up connectivity construction is to separate the logic by input grid type when possible, and only use the fully general default solution for uncategorized or irregular grids.

In many common cases, we are working with structured or semi-structured grids where we can safely use additional assumptions about the grid shape/connectivity. Those assumptions can make the implementation much easier to optimize and vectorize. The current general solution is useful because it can handle arbitrary grid shapes, but that also makes it less ideal for performance-sensitive paths.

So I think a good direction would be to provide faster specialized paths for the input grid types UXarray already supports, while keeping the existing general method as the fallback for arbitrary grids.

Also, if any of the terminology around spherical grids, connectivity, or grid geometry is confusing, you can refer to my paper here: https://egusphere.copernicus.org/preprints/2026/egusphere-2026-636/. It may be more detailed than what is needed for this specific PR, but I think it can serve as a useful guideline and reference for the concepts and terminology involved here.

[cmdupuis3]

In many common cases, we are working with structured or semi-structured grids where we can safely use additional assumptions about the grid shape/connectivity. Those assumptions can make the implementation much easier to optimize and vectorize. The current general solution is useful because it can handle arbitrary grid shapes, but that also makes it less ideal for performance-sensitive paths.

Yeah, I see what you mean. The edge_node/face_edge combined routine is nice for optimization, but it imposes branching logic at a low level depending on the grid state.

My idea is that we want to have a dataclass of booleans (and maybe add an enum layer for good API), mapping to CONNECTIVITY_NAMES that describes the connectivity state at the Grid level. Then instead of API calls to individual connectivity builders, we only allow a call to a general get_connectivity(state: ConnState) with the connectivity dataclass. Depending on what flags are chosen in the dataclass, that will give us a clear idea of what optimization possibilities are available in that specific case. That way, we don't have to have any branching logic inside the combined routines.

For now, I'm still working on the merge, but I just wanted to share this idea while it's fresh. I just noticed that the combined edge_node/face_edge routine can result in corner cases if one is already set in a grid, but not both.

[hongyuchen1030]

In many common cases, we are working with structured or semi-structured grids where we can safely use additional assumptions about the grid shape/connectivity. Those assumptions can make the implementation much easier to optimize and vectorize. The current general solution is useful because it can handle arbitrary grid shapes, but that also makes it less ideal for performance-sensitive paths.

Yeah, I see what you mean. The edge_node/face_edge combined routine is nice for optimization, but it imposes branching logic at a low level depending on the grid state.

My idea is that we want to have a dataclass of booleans (and maybe add an enum layer for good API), mapping to CONNECTIVITY_NAMES that describes the connectivity state at the Grid level. Then instead of API calls to individual connectivity builders, we only allow a call to a general get_connectivity(state: ConnState) with the connectivity dataclass. Depending on what flags are chosen in the dataclass, that will give us a clear idea of what optimization possibilities are available in that specific case. That way, we don't have to have any branching logic inside the combined routines.

For now, I'm still working on the merge, but I just wanted to share this idea while it's fresh. I just noticed that the combined edge_node/face_edge routine can result in corner cases if one is already set in a grid, but not both.

Yeah, I think this high-level dispatcher idea makes sense and is a good direction for improving vectorization.

Once this dispatcher structure is in place, we can then start adding specialized handling for different grid types, such as structured, semi-structured, or fully general unstructured grids.

[cmdupuis3]

Migrating this to #1560, since the fresh version is in my fork.