Optimized connectivity: Simultaneous face_edges and edge_nodes

[cmdupuis3]

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]

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

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

[hongyuchen1030]

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]

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.

[erogluorhan]

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

[cmdupuis3]

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

[cmdupuis3]

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]

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
    )

[Sevans711]

Question for @erogluorhan, does uxarray have any plans to support chunked grids? I was under the impression that data variables can be chunked, but that it is always safe to assume the grid itself can be fully loaded into memory. Is that incorrect?

That said, using more dask-compatible syntax like this doesn't seem like it would have a downside (no need to block merging on this question).

[erogluorhan]

Per @cmdupuis3 's latest comment:

This looks great, and it is more beneficial than just "hacking" I believe (please document it accordingly wherever possible in your code):

• Using xr.Dataset.compute() guarantees two major advantages:

  • Bundled instead of separate .compute() calls on a couple DataArrays - uses same upstream dask graph and gets rid of duplicate overhead
  • Doesn't break optional Dask - If given chunked arrays, uses Dask; otherwise, skips compute under the xarray's hood and respects numpy arrays

[erogluorhan]

Per @Sevans711's latest comment:

The intent with Grid's I/O design was already to have it support chunked features; however; full end-to-end "Grid supports chunkedness" is probably impossible - there are global-topology and spatial-index operations like what we see here in this PR that could probably never avoid materializing the whole array.

1. Grid I/O is already made chunked - there might be things we can further fix though
2. Some of the operations can be made chunk-compatible

I think a smart way forward would be to (1) keep (or make) the chunk-local geometric operations lazy, and (2) leave ( and explicitly document) the global-topology/KDTree/validation operations as materialized (whole array computations) rather than attempting to "fix" every .values or .compute() - this PR shows there is still room for optimization for even those though

[cmdupuis3]

Added this comment:

# 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

Is this bit resolved?

[cmdupuis3]

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]

pre-commit.ci autofix

[cmdupuis3]

pre-commit.ci autofix

[erogluorhan]

On this PR, just doing grid=ux.tutorial.open_grid("outCSne30-vortex"); grid.edge_node_connectivity leads to 648036 bytes. (The extra ~40000 bytes are from n_nodes_per_face, though that is basically just a rounding error when compared to the ~200000 extra bytes from storing face_edge_connectivity.) Of course, the numbers here are very small; this is a tiny example. But, they demonstrate that the concern about memory costs may be somewhat plausible, at least. (Unless everybody who asks for edge_node_connectivity also wants face_edge_connectivity?)

Yeah, so I was wondering about that too, which is why I proposed having a dispatcher for combined connectivity routines #1559. However, in practice, I think the most likely situation with the grids we support is that both face/edge and edge/node connectivity are the missing connectivities, which is why they were interested in doing this. There could definitely be a discussion on the design, but I think that's out of scope for this PR.

Ah, yes that makes sense, and seems like a good solution. I also agree a full design like that should probably be scoped to a different PR. Before this PR merges I would still be interested to hear from other uxarray developers/users to confirm whether the increased memory cost is an acceptable tradeoff for the possible huge speedup in the meantime. E.g., @erogluorhan, @rajeeja, @rljacob, if you get a chance to look into this, what are your thoughts?

I realize I never responded to this.

• Dispatcher: Agreed. Out of scope for this PR but we could look into it separately (Dispatcher for optimized connectivity routines #1559).
• Memory usage:
  • The extra bytes in the static output: They seem acceptable to me in this context where we have extreme performance optimization, and the models we support seem to have these two connectivities either going together or lacking together (per @cmdupuis3 's observation) make this combined approavh even more logical.
  • Also, the peak-mem results @cmdupuis3 provided seem like the approach here even settles peak memory usage. The extra bytes are measured for stored output size, but the peak memory during combined construction seems significantly improved.

All that said, there seems like a big win on both performance and memory utilization in this PR.

[erogluorhan]

@cmdupuis3 a quick observation:

• I believe this PR is not documenting this new way of combined constructions, so please mention it in both face_edge and edge_node docstrings for the user's reference.

[Sevans711]

Took a close look at this one just now. It does look like it is getting closer to being ready overall, I have a few small questions and I left some inline comments accordingly.

The main blockers for me would be a couple bigger questions which I do not think are resolved yet. These are: "why does this PR introduce lexicographic ordering?" and "do the changes here actually speed up n_nodes_per_face for large grids?" (See replies to older comment threads for more details).

[Sevans711]

(keeping this comment unresolved until at least one check shows faster n_nodes_per_face performance on a large grid)

[cmdupuis3]

pre-commit.ci autofix