Skip to content

Compute face_node_angles from uxarray.Grid? #1566

Closed

Description

@Sevans711

Proposed new feature or change:

This would have the same shape as face_node_connectivity: for every face, for every node in that face, this tells the angle at that node.

Computing angles is necessary for computing one common measure of grid skewness. Related: #1305.

There are a few different possible ways to get the result, so I'm not sure if this should be a cached property or a function with a few optional flags:

  • should the result be in radians or degrees?
  • should the computation use 'euclidean' / flat geometry (sum of angles of a triangle is always 180 degrees) or respect the 'spherical' geometry (sum of angles of a triangle is always larger than 180 degrees)?

What are the best defaults for those flags?

Also, perhaps grid.face_node_angles(geometry=...) should dispatch to face_node_angles_flat or face_node_angles_spherical properties (whose results should be cached)?

Activity

  1. self-assigned this
    on Jul 14, 2026
  2. Sevans711 commented on Jul 14, 2026

    @Sevans711
    CollaboratorAuthor

    Pinging @rajeeja in case you have thoughts or advice on this!

  3. rajeeja commented on Jul 14, 2026

    @rajeeja
    Contributor

    Hi @Sevans711 — looked into this. Answers to your questions first, then a numerics point and a question for @hongyuchen1030.

    Your questions

    • Radians vs degrees: default radians (matches uxarray's internal convention across arcs.py / coordinates.py), with a degrees=True flag for convenience.
    • Flat vs spherical: default spherical — it is the true geometry of these grids and the one that matters for skewness; keep euclidean as a cheaper option.
    • Best defaults: geometry="spherical", degrees=False.
    • Property vs function: go with grid.face_node_angles(geometry=..., degrees=...) dispatching to cached _face_node_angles_spherical / _face_node_angles_euclidean internals. Your instinct is right — a bare property cannot take flags, and the dispatch-to-cached-variant pattern gives both flexibility and caching.

    Compute lazily, never eagerly. These angles are a derived, optional quantity most workflows never need, and computing them is a per-face/per-node pass over the whole mesh. Follow the face_areas pattern: compute on first request, cache, and never touch it at grid construction. Cache per geometry variant; keep the cache in radians and convert to degrees on return (cheap post-scale, no need to cache separately).

    Geometry, so you don't compute the wrong thing. The angle at a node is the interior angle between the two edges meeting there, not the angle between node position vectors. At node B with face-neighbors A and C: euclidean is the angle between chords (A−B) and (C−B); spherical is the angle between the great-circle arcs BA and BC, from the tangent directions at B (project A and C onto the tangent plane at B, i.e. remove the component along B, then take the angle between the tangents) — equivalently the angle between the edge-plane normals n1 = B×A and n2 = B×C.

    Reuse existing primitives: uxarray/grid/utils.py::_small_angle_of_2_vectors uses the stable 2*atan2(|u−v|,|u+v|) form — prefer it over anything arccos-based, since skewed / near-degenerate cells (your regime) are where arccos loses precision. node_x/y/z, face_node_connectivity, and n_nodes_per_face are cached and ready; mask INT_FILL_VALUE in the connectivity for variable-sided faces or the padding slots produce junk angles.

    Free validation: for a spherical triangle, sum(interior angles) − π equals the spherical excess, which equals the face area — so you can cross-check against grid.face_areas in tests (generalizes to (n−2)π + area for n-gons; flat sum is exactly (n−2)π).

    On accuracy / #1513. The spherical angle is built from cross products of the edge vectors (n1 = B×A, n2 = B×C). Naive cross products are difference-of-products and lose precision to cancellation exactly when edges are short or nearly collinear — which is the skewed-cell regime skewness analysis cares about most. The compensated cross-product kernels in the AccuSphGeom port (#1513) are built for precisely that, so I would structure the cross-product step so it can use those kernels once #1513 merges, rather than blocking on it now. #1513 is still open (not yet merged).

    @hongyuchen1030 — a question for you, referencing this issue (#1566): AccuSphGeom's EFT is aimed at robust predicates (sign / orientation). Computing a node angle is a magnitude problem — accurate normals via accucross, then atan2(|n1×n2|, n1·n2). Does the compensated cross product meaningfully improve the accuracy of the resulting angle magnitude in the near-degenerate regime, or is the final error there dominated by the atan2 conditioning near 0/π? Trying to gauge whether wiring the #1513 kernels into this buys real angle accuracy or mainly helps the orientation/sign side.

    Happy to review once you have a draft.

  4. Sevans711 commented on Jul 14, 2026

    @Sevans711
    CollaboratorAuthor

    @rajeeja Thank you for your detailed response! My thoughts below:

    Best defaults: geometry="spherical", degrees=False

    Sounds good.

    go with grid.face_node_angles(geometry=..., degrees=...) dispatching to cached _face_node_angles_spherical / _face_node_angles_euclidean internals.

    Compute lazily, never eagerly. … keep the cache in radians.

    Agreed.

    Geometry [misc notes]

    I think that all makes sense. I might realize I have questions once I start the implementation; I'll follow up if I get stuck.

    Reuse existing primitives: uxarray/grid/utils.py::_small_angle_of_2_vectors

    Thank you for pointing me to this! This helped me discover _angle_of_2_vectors, which is probably what I will end up using, to account for the possibility of non-convex faces where an internal angle can be >180 degrees. Unless… do uxarray grids actually require all faces to be convex?

    Free validation: cross-check against grid.face_areas

    I agree (and also, Claude happened to mention that idea as well when I initially asked about computing skewness). I found #1567 and #1571 while looking into face_areas because of this.

    On accuracy: … [precision issues in the regime we care about, might be fixed but AccuSphGeom port]

    This makes sense, and I will keep in mind that we might replace with more accurate cross products, or maybe a flag to swap between modes, in the near future.

    Happy to review once you have a draft.

    Thank you! I will request you as a reviewer once the PR is ready to go.

  5. erogluorhan commented on Jul 15, 2026

    @erogluorhan
    Member

    Some thoughts:

    • Computation of this should be straight forward and inexpensive since we already have face lats/lons and node lats/lons
    • Calculations should be on spherical geometry only I believe
    • Caching would be costly and redundant though - I think we should avoid it
      • We had some initial discussions about similar things and ended up deciding to cache only the necessary-to-store properties and deferring to adhoc calculations for other properties like this. That approach still makes sense to me from a memory utilization standpoint.
  6. Sevans711 commented on Jul 15, 2026

    @Sevans711
    CollaboratorAuthor

    @erogluorhan thank you for sharing these thoughts as well. I'll focus on spherical geometry first, and disable caching by default. I might still add a flag like grid.face_node_angles(..., cache=False) to enable the user to opt-in to caching, in case they know ahead of time that they will be doing multiple computations related to angles, if that sounds reasonable?

  7. erogluorhan commented on Jul 15, 2026

    @erogluorhan
    Member

    @erogluorhan thank you for sharing these thoughts as well. I'll focus on spherical geometry first

    Great; thank you!

    , and disable caching by default. I might still add a flag like grid.face_node_angles(..., cache=False) to enable the user to opt-in to caching, in case they know ahead of time that they will be doing multiple computations related to angles, if that sounds reasonable?

    This is worth running some large data tests to see the performance and memory usage on NCAR's HPC clusters, maybe using 3.75 km data. Memory usage associated with the caching of it could be roughly calculated actually. Assuming you'll use float, it would (for our sample 3.75 km MPAS grid) be around (~42M faces * 6 nodes * 8 bytes) = ~2GB of extra memory usage.

    The rest should be:

    • If execution time is negligibly short (I am anticipating this to be the case), we should definitely avoid caching
    • Else, let's further discuss expectations vs costs with caching
  8. Sevans711 commented on Jul 15, 2026

    @Sevans711
    CollaboratorAuthor

    Small update: I managed to finish an initial implementation (see commit ecba4f1). Still need to add tests and examples before opening a PR.

    The non-spherical geometry seems simple to support (it doesn't add any operations, it just skips the "project onto the tangent plane" step) so I included it as an option for now; it only took a few extra minutes to implement, to document the relevant kwarg and add "if geometry=..." checks inside the code.

    I definitely agree that caching should be disabled by default, unless the calculation actually becomes a bottleneck. I did still include the caching option though, because it only took a few extra minutes to add it. This should help us later test in real workflows whether it saves non-negligible amounts of time:

    1. Run workflow/analysis without ever touching the cache kwarg. This will avoid caching entirely.
    2. Run same workflow/analysis but call grid.face_node_angles(..., cache=True) once at the top. This will lead to all calls using the cached value.

    If tests show there's never any point in supporting caching and we want to remove that option to avoid the overhead of slightly longer code / more for users to read about and consider, it should be easy enough to remove the relevant pieces of the code.

  9. erogluorhan commented on Jul 16, 2026

    @erogluorhan
    Member

    Re caching: Please see my previous comment and complete the tests accordingly and let us know the results before introducing functionality into API. (Once introduced into user API, even if it is disabled by default, removing it later would break user workflows if it gets some usage)

    Re flat vs sphere: @rajeeja @hongyuchen1030 do you think it is worth to introduce any flat geometry calculations or should we focus only on spherical geometry?

  10. erogluorhan commented on Jul 16, 2026

    @erogluorhan
    Member

    Small update: I managed to finish an initial implementation (see commit ecba4f1). Still need to add tests and examples before opening a PR.

    Oh, I just had a look into the code now, and I believe it could be and has to be vectorized (instead of nested for loops) to run it in virtually "no time" for varying grid sizes. That way, I believe we'd have negligible execution times with it.

  11. Sevans711 commented on Jul 16, 2026

    @Sevans711
    CollaboratorAuthor

    Re caching: Please see my previous comment and complete the tests accordingly and let us know the results before introducing functionality into API. (Once introduced into user API, even if it is disabled by default, removing it later would break user workflows if it gets some usage)

    Ah, this really helped clarify for me, thank you. Now I understand your concern about including it in the API in the first place. I will make sure to run some initial tests so we can decide whether to remove the caching option before merging into main. I will reach out to you separately about getting HPC access if I have questions!

    Similarly, I now understand it might make more sense to exclude the "geometry" option entirely. I am happy to follow whatever guidance people provide on that. Since it is in my initial implementation already, I will wait for others to chime in before removing it, and will also make sure we've come to a decision on this before merging into main.

    I believe it could be and has to be vectorized

    My understanding is that numba speeds are comparable to (or better than) vectorized numpy computations. I believe vectorizing here is also not particularly elegant, because n_nodes_per_face varies from face to face. Am I misunderstanding something?

    I didn't want to spend too much time on premature optimizations, but it did seem worthwhile to get some initial tests, so I asked claude for an initial/naive vectorized implementation without numba. I skimmed it for correctness and tested that the results are consistent with the numba results for all the uxarray tutorial grids and the grids below. I pushed it to a new branch for now (face_node_angles_nonumba; see 0feade2).

    Now we can compare the timing (tests run locally on my machine, in Python 3.14):

    import time
    import numpy as np
    import uxarray as ux
    for key in ['mpas-QU-480', 'outCSne30-vortex', 'scrip-ne30pg2']:  # some of the larger tutorial datasets
        grid = ux.tutorial.open_grid(key)
        gridsizestr = f"n_node={grid['n_node']}, n_max_face_nodes={grid['n_max_face_nodes']}"
        node_per_face_minmax = (grid.n_nodes_per_face.min().item(), grid.n_nodes_per_face.max().item())
        grid.node_x  # <-- make sure node_x, node_y, node_z are all computed beforehand
        # time numba implementation
        now = time.time()
        res0 = grid.compute_face_node_angles(assume_convex=True, numba=True)
        t0 = time.time()-now
        print(f"{key:>20s} (numba): {t0:.4f} s, {gridsizestr}, node_per_face={node_per_face_minmax}")
        # time numpy implementation
        now = time.time()
        res1 = grid.compute_face_node_angles(assume_convex=True, numba=False)
        t1 = time.time()-now
        print(f"{key:>20s} (numpy): {t1:.4f} s, {gridsizestr}, node_per_face={node_per_face_minmax}")
        assert np.max(np.abs(res1-res0)) < 1e-5  # ensure both implementations give similar results.
        print()

    This prints:

             mpas-QU-480 (numba): 0.0018 s, n_node=3947, n_max_face_nodes=6, node_per_face=(5, 6)
             mpas-QU-480 (numpy): 0.0024 s, n_node=3947, n_max_face_nodes=6, node_per_face=(5, 6)
    
        outCSne30-vortex (numba): 0.0035 s, n_node=5402, n_max_face_nodes=4, node_per_face=(4, 4)
        outCSne30-vortex (numpy): 0.0041 s, n_node=5402, n_max_face_nodes=4, node_per_face=(4, 4)
    
           scrip-ne30pg2 (numba): 0.0096 s, n_node=21727, n_max_face_nodes=4, node_per_face=(4, 4)
           scrip-ne30pg2 (numpy): 0.0128 s, n_node=21727, n_max_face_nodes=4, node_per_face=(4, 4)
    

    (Well, actually, the first time it is run, it takes closer to 1.8 seconds for the very first numba computation, because numba takes time to compile relevant functions. This is what it looked like the second time.)

    I also ran similar code for some larger grids I happen to have downloaded on my machine, and got:

          oQU120.grid.nc (numba): 0.0252 s, n_node=59329, n_max_face_nodes=6, node_per_face=(5, 6)
          oQU120.grid.nc (numpy): 0.0225 s, n_node=59329, n_max_face_nodes=6, node_per_face=(5, 6)
    
    conus12km/conus12km.invariant.nc_L60_GFS (numba): 0.0789 s, n_node=317280, n_max_face_nodes=6, node_per_face=(5, 6)
    conus12km/conus12km.invariant.nc_L60_GFS (numpy): 0.1157 s, n_node=317280, n_max_face_nodes=6, node_per_face=(5, 6)
    

    Conclusions: numba seems maybe faster, and at least comparable, to a numpy vectorized approach. Or, maybe there is some huge inefficiency that needs to be solved in the numpy implementation (see 0feade2)?

  12. Sevans711 commented on Jul 16, 2026

    @Sevans711
    CollaboratorAuthor

    Actually, I also have a few more quick follow-up questions:

    1. I am considering adding a way to check whether faces are actually convex. The current face_node_angles implementation will fail if any face is non-convex. Is there any reason that would be a bad idea? Are there other parts of uxarray which implicitly assume faces are all convex (does anything in face_areas assume convex faces, @rajeeja)?
    2. Do you have something in mind already for what to compare the time to?

    For 2, my initial thought is to compare to face_areas computation time, since uxarray already supports caching for face_areas. On main right now, face_areas takes much longer than face_node_angles. But, on the face_areas API cleanup and speedup branch (see #1571), the timing is actually quite similar:

    for key in ['mpas-QU-480', 'outCSne30-vortex', 'scrip-ne30pg2']:
        grid = ux.tutorial.open_grid(key)
        grid.node_x  # <-- make sure node_x, node_y, node_z are all computed beforehand
        now = time.time()
        res0 = grid.compute_face_areas()
        t0 = time.time()-now
        print(f"{key:>20s} (face_area): {t0:.4f} s, n_node={grid['n_node']}")

    This, and similar code for the other two grid files tested above, produces:

             mpas-QU-480 (face_area): 0.0061 s, n_node=3947
        outCSne30-vortex (face_area): 0.0035 s, n_node=5402
           scrip-ne30pg2 (face_area): 0.0094 s, n_node=21727
          oQU120.grid.nc (face_area): 0.0251 s, n_node=59329
    conus12km/conus12km.invariant.nc_L60_GFS (face_area): 0.0895 s, n_node=317280
    

    Of course, these tests aren't as helpful as large data tests on HPC clusters, but hopefully they provide a bit more context in the meantime before I am set up on HPC.

  13. 7 remaining items

  14. hongyuchen1030 commented on Jul 20, 2026

    @hongyuchen1030
    Contributor

    @Sevans711 Yes, we already implicitly assume convex/star-shaped faces: face_areas triangulates each face as a fan from node 0 (exact for convex, silently wrong for non-convex - it won't error), and the barycentric helpers in geometry.py/neighbors.py are documented convex-only. So a convexity check is a good idea - I'd make it a shared is_face_convex mask so face_areas/remapping can reuse it, and just fail loudly (with a clear message) in face_node_angles for now.

    To my knowledge, none of the model outputs we are supporting might come natively with concave faces, and our implicit assumption of convex faces only grids seems to have worked so far. If that is the case, we might not need an explicit convex face check. Thoughts? @rajeeja @rljacob @hongyuchen1030 @brianpm @zarzycki

    I don't think we need the "convex" check. Given how hard it is to determine a face is actually "convex" on the sphere and the fact that most of our face are not very ill-conditioned. Our algorithms should cover all convex cases and mildly concave faces as well.

    But I think it might be useful to have a mask to indicate if the given grids are triangles or not, if possible. Triangulated faces sometimes have tailored algorithms for better performance.

  15. Sevans711 commented on Aug 4, 2026

    @Sevans711
    CollaboratorAuthor

    Finally got set up to use HPC and was able to test the timing. Here's what the timing looks like for large MPAS grids:

    grid numba 4 threads numpy numba 1 thread ux.open_grid()
    MPAS 30km 0.53 s 1.23 s 1.79 s ~2.7 s
    MPAS 15km 1.94 s 4.99 s 7.07 s ~18 s
    MPAS 7.5km 7.64 s 12.31 s 29.03 s ~65 s
    MPAS 3.75km 29.4 s 48.8 s 115.56 s ~250 s

    (See previous comment about timing for more details on the code being run: #1566 (comment))

    The last column shows how long it took to run ux.open_grid(gridfile) for each grid file, for reference.

    I used numba.set_num_threads() to adjust the number of threads. Getting a job with 4 threads through the HPC job submission queues was very fast on Casper. I think it is reasonable to expect machines to be capable of handling at least 4 threads if they are also able to handle data processing on a 20 GB grid (~MPAS 3.75km). I included the 1 thread case for reference, to showcase the importance of using prange and @njit(..., parallel=True) in the implementation.

    I believe both the numba and numpy implementations could probably be further refined to improve efficiency, but I think we should save that for future PRs? These tests feel sufficient to answer the two primary questions we've been asking in this thread:

    1. Should we use numpy or numba?
    • Answer: numba.
      • The two are close enough that either would probably be fine, but numba seems consistently faster.
      • Numba allows using more threads; using more than 4 threads presumably might see even more speedup. This makes numba appealing even if further optimizations are possible in the numpy case.
      • The numba code feels easier to read and write, as utilizing connectivity variables directly with numpy alone requires some smart indexing tricks due to non-constant n_nodes_per_face.
    1. Does the computation run in virtually "no time" / have negligible execution time? I.e., is the compute time short enough that we shouldn't even provide a cache option?
    • Answer: no, I don't think so.
      • Considering the 3.75km grid case, we're trying to weigh repeating 30 seconds of compute time against roughly 2 GB of extra memory usage. Just calling ux.open_grid() already uses roughly 20GB, and that doesn't factor in any of the memory of actual data variables living on the grid itself. So, while 2 GB sounds large, relatively it's not actually that huge. These numbers make me think "cache=True should be an option but not the default."
      • Also, as per messages above, timing is similar to face_areas computations. face_areas results are cached, presumably because their computation time was considered to be non-negligible. There, caching is enabled by default, but I wouldn't argue for cache=True by default here because face_areas is only a fraction of the size (e.g., 1/6th the size of face_node_angles for the large MPAS grids where n_max_face_nodes=6).
  16. Sevans711 commented on Aug 4, 2026

    @Sevans711
    CollaboratorAuthor

    With those tests completed, and considering the other comments in this thread, I would propose the following plan for implementation:

    1. Remove the geometry option. Always assume spherical geometry.
    2. Remove the assume_convex flag. Always assume convex faces.
    3. Keep the degrees option as-is, with degrees=False by default. (Technically people could just convert afterwards, but it's a simple and convenient option to keep, and it helps to clarify the result is in radians, just by glancing at the function signature.)
    4. Keep the cache option untouched, with caching disabled by default.
    5. Keep the implementation in numba. Do not introduce a numpy implementation.
    6. For consistency with the cleaned-up face_areas API (see Clean up and vectorize the face_areas API #1577), keep the name Grid.compute_face_node_angles(), and add a property Grid.face_node_angles which returns the result of this method. Clarify in the Grid.face_node_angles property docstring that results are recomputed each time by default, but user can call Grid.compute_face_node_angles(cache=True) once to cache the results.
  17. moved this from 📚 Backlog to 🏗 In progress in UXarray Developmenton Aug 4, 2026
  18. erogluorhan commented on Aug 4, 2026

    @erogluorhan
    Member

    Will look into this and review your suggestions soon

  19. erogluorhan commented on Aug 4, 2026

    @erogluorhan
    Member

    AI Disclosure: I had an instinct that an optimized numpy vectorization could have outperformed numba and played with Claude but it optimized numba as well as numpy and further opened the gap between the two in favor of numba :)

    @Sevans711 — happy to share the full benchmark script if useful.
    @cmdupuis3 - These might be of interest to you

    The following is from Claude:

    Followed up on the numba-vs-numpy benchmarking question above — wrote a vectorized numpy implementation and compared it against the current numba kernel (ecba4f1, branch face_node_angles) and the rough numpy version already tried (0feade2, branch face_node_angles_nonumba_option).

    Setup: mpas-QU-480 (real, 1,791 faces) plus the same grid tiled to ~500K and ~2M faces as a stand-in for the 3.75km/42M-face HPC test above (didn't have that grid file available; relative trends below should still hold). All outputs agree to float64 precision (max|diff| ~ 1e-15) across every implementation.

    New numpy implementation — groups faces by node-count (most grids have only 1–3 distinct arities: MPAS is pentagons+hexagons, cubed-sphere is all-quads) and vectorizes each group with np.roll for prev/next-node lookup instead of a global modulo + take_along_axis gather, so there's no wasted work on padded slots:

    def _compute_face_node_angles_convex_numpy(
        node_x, node_y, node_z, face_node_connectivity, n_nodes_per_face
    ):
        n_faces, n_max_face_nodes = face_node_connectivity.shape
        result = np.full((n_faces, n_max_face_nodes), np.nan, dtype=np.float64)
        for n in np.unique(n_nodes_per_face):
            n = int(n)
            rows = np.nonzero(n_nodes_per_face == n)[0]
            if rows.size == 0 or n < 3:
                continue
            conn = face_node_connectivity[rows, :n]
            prev, nxt = np.roll(conn, 1, axis=1), np.roll(conn, -1, axis=1)
    
            xhere, yhere, zhere = node_x[conn], node_y[conn], node_z[conn]
            xprev, yprev, zprev = node_x[prev], node_y[prev], node_z[prev]
            xnext, ynext, znext = node_x[nxt], node_y[nxt], node_z[nxt]
    
            v1 = np.stack([xprev - xhere, yprev - yhere, zprev - zhere], axis=-1)
            v2 = np.stack([xnext - xhere, ynext - yhere, znext - zhere], axis=-1)
    
            here = np.stack([xhere, yhere, zhere], axis=-1)
            normal = here / np.linalg.norm(here, axis=-1, keepdims=True)
            v1 = v1 - np.sum(v1 * normal, axis=-1, keepdims=True) * normal
            v2 = v2 - np.sum(v2 * normal, axis=-1, keepdims=True) * normal
    
            norm_v1 = np.linalg.norm(v1, axis=-1, keepdims=True)
            norm_v2 = np.linalg.norm(v2, axis=-1, keepdims=True)
            vec_minus = norm_v2 * v1 - norm_v1 * v2
            vec_sum = norm_v2 * v1 + norm_v1 * v2
            angles = 2 * np.arctan2(
                np.linalg.norm(vec_minus, axis=-1), np.linalg.norm(vec_sum, axis=-1)
            )
            result[rows[:, None], np.arange(n)[None, :]] = angles
        return result
    grid size numba (as in branch) numpy naive (0feade2) numpy grouped (above)
    1,791 faces 1.0 ms 1.1 ms 1.1 ms
    ~500K faces 1,755 ms (1 thr) / 528 ms (4 thr) / 223 ms (12 thr) 459 ms 424 ms
    ~2M faces 7,037 ms (1 thr) / 2,107 ms (4 thr) / 847 ms (12 thr) 1,873 ms 1,711 ms

    So: single-threaded, the current numba kernel is ~4x slower than plain numpy, and only edges past numpy once given several threads.

    Why: the current kernel builds a fresh np.array([...]) for v1, v2, and the surface normal on every loop iteration — a well-known numba anti-pattern that blocks scalar/register optimization.

    Same math, rewritten with plain scalars and no array literals in the hot loop:

    @njit(cache=True, parallel=True, fastmath=True)
    def _compute_face_node_angles_convex_scalar(
        node_x, node_y, node_z, face_node_connectivity, n_nodes_per_face
    ):
        n_faces, n_max_face_nodes = face_node_connectivity.shape
        result = np.full(face_node_connectivity.shape, np.nan, dtype=np.float64)
        for i in prange(n_faces):
            n_nodes = n_nodes_per_face[i]
            for j in range(n_nodes):
                ihere = face_node_connectivity[i, j]
                iprev = face_node_connectivity[i, (j - 1) % n_nodes]
                inext = face_node_connectivity[i, (j + 1) % n_nodes]
                xhere, yhere, zhere = node_x[ihere], node_y[ihere], node_z[ihere]
                v1x, v1y, v1z = node_x[iprev]-xhere, node_y[iprev]-yhere, node_z[iprev]-zhere
                v2x, v2y, v2z = node_x[inext]-xhere, node_y[inext]-yhere, node_z[inext]-zhere
                nnorm = (xhere*xhere + yhere*yhere + zhere*zhere) ** 0.5
                nx, ny, nz = xhere/nnorm, yhere/nnorm, zhere/nnorm
                d1 = v1x*nx + v1y*ny + v1z*nz
                v1x, v1y, v1z = v1x - d1*nx, v1y - d1*ny, v1z - d1*nz
                d2 = v2x*nx + v2y*ny + v2z*nz
                v2x, v2y, v2z = v2x - d2*nx, v2y - d2*ny, v2z - d2*nz
                norm_v1 = (v1x*v1x + v1y*v1y + v1z*v1z) ** 0.5
                norm_v2 = (v2x*v2x + v2y*v2y + v2z*v2z) ** 0.5
                mx, my, mz = norm_v2*v1x - norm_v1*v2x, norm_v2*v1y - norm_v1*v2y, norm_v2*v1z - norm_v1*v2z
                sx, sy, sz = norm_v2*v1x + norm_v1*v2x, norm_v2*v1y + norm_v1*v2y, norm_v2*v1z + norm_v1*v2z
                norm_m = (mx*mx + my*my + mz*mz) ** 0.5
                norm_s = (sx*sx + sy*sy + sz*sz) ** 0.5
                result[i, j] = 2 * np.arctan2(norm_m, norm_s)
        return result
    grid size numba scalar (1 thr) numba scalar (4 thr) numba scalar (12 thr) numpy grouped
    ~500K faces 68 ms 19 ms 11 ms 424 ms
    ~2M faces 283 ms 77 ms 42 ms 1,711 ms

    Once written idiomatically, numba is ~6x faster than the best numpy vectorization even single-threaded, and 20–40x faster with threads — numpy has to materialize a full intermediate array in memory for every op (v1, v2, projected normal, dot
    products, norms...), while the scalar loop fuses all of that per-face-node with no intermediate allocation. This is the expected shape for numba vs numpy: ragged data with tiny, fixed-size per-element math is numba's sweet spot.

    Suggestion: keep numba, just fix the array-allocation pattern in the kernel — it's a small, mechanical change with a large payoff, not a rewrite.

    (Benchmarking and both implementations above worked out with Claude's help.)

  20. Sevans711 commented on Aug 5, 2026

    @Sevans711
    CollaboratorAuthor

    Thank you for looking into this in detail!

    Can you re-run the benchmarks without using fastmath=True? I think uxarray cares a lot about numerical precision, and my understanding is that that flag is happy to silently mess up your accuracy in strange edge cases where you wouldn't expect any issues (even if no such edge cases appeared in the cases you tried here).

    I also suspected the numba implementation was sub-optimal for constructing so many np.arrays internally; I have the intuition that constructing a numpy array is "expensive". But, I wanted to use the _small_angle_of_2_vectors primitive, which also uses numpy functions internally already. Using _small_angle_of_2_vectors inside _compute_face_node_angles_convex feels more idiomatic for this PR, rather than in-lining all of its logic.

    Seeing that we care enough to optimize efficiency like this on the first implementation of this feature, I would maybe add to the plan: rewrite _small_angle_of_2_vectors to avoid building numpy arrays unnecessarily (when the inputs aren't numpy arrays yet).

    I also suspected the numpy implementation was sub-optimal but didn't know how much to focus on efficiency here. I appreciate your (and Claude's) efforts to optimize both and compare.

  21. hongyuchen1030 commented on Aug 5, 2026

    @hongyuchen1030
    Contributor

    Thank you for looking into this in detail!

    Can you re-run the benchmarks without using fastmath=True? I think uxarray cares a lot about numerical precision, and my understanding is that that flag is happy to silently mess up your accuracy in strange edge cases where you wouldn't expect any issues (even if no such edge cases appeared in the cases you tried here).

    I also suspected the numba implementation was sub-optimal for constructing so many np.arrays internally; I have the intuition that constructing a numpy array is "expensive". But, I wanted to use the _small_angle_of_2_vectors primitive, which also uses numpy functions internally already. Using _small_angle_of_2_vectors inside _compute_face_node_angles_convex feels more idiomatic for this PR, rather than in-lining all of its logic.

    Seeing that we care enough to optimize efficiency like this on the first implementation of this feature, I would maybe add to the plan: rewrite _small_angle_of_2_vectors to avoid building numpy arrays unnecessarily (when the inputs aren't numpy arrays yet).

    I also suspected the numpy implementation was sub-optimal but didn't know how much to focus on efficiency here. I appreciate your (and Claude's) efforts to optimize both and compare.

    1. I agree with diabling fastmath = True. The concern includes both FMA contraction and the broader set of LLVM fast-math assumptions enabled by Numba (like assumption for nnan and ninf). UXarray is also not computationally expensive enough for fastmath=True to be necessary by default. If disabling it causes a significant performance regression, that likely indicates a real performance issue in one of the kernels that should be investigated separately.

    2. I also agree that we should use _small_angle_of_2_vectors consistently whenever calculating the angle between two vectors, rather than inlining the same logic in multiple places. It should not introduce meaningful overhead if it is implemented in the same plain-scalar, Numba-friendly style as the surrounding kernel. If it does, then the helper itself should be rewritten in that style.

    3. The updated implementation appears to mainly replace the original code with a more branch-free, SIMD-friendly formulation. That is generally the right direction for this type of Numba kernel. We are effectively writing C++-style SIMD code in Python, which gives us more direct control over branching, vectorization, and intermediate allocations than relying on higher-level NumPy operations.

    4. In my experience these LLM models still have fairly limited depth on this kind of low-level vectorization work, they can silently drop important edge cases or introduce things like fastmath=True without flagging the tradeoff (exactly what happened here). Worth treating their output as a first draft to review

  22. erogluorhan commented on Aug 5, 2026

    @erogluorhan
    Member

    @Sevans711 @hongyuchen1030 all great points that I 100% agree with. We should do our own assessment and finalize the implementation ourselves.

    Sam, dropping fastmath and doing small_angle_of_two_vecors, there didn't seem to be meaningful differences on execution times with Claude, but as we discussed offline with you, I think from this point it really needs you to do further tests on real data and see it from there.

  23. Sevans711 commented on Aug 11, 2026

    @Sevans711
    CollaboratorAuthor

    I added an optimized implementation on the face_node_angles branch. The diff from main looks like: 5ef66d8...09bafd7

    Also it looks like the "constructing many tiny numpy arrays inside numba methods" performance hit is occurring throughout more areas of uxarray; I opened #1648 to address that. I think that issue should be a relatively high priority in terms of speeding up uxarray methods. For example, it occurs inside _face_area_from_quadrature which gets called by all face_areas computations. This suggests to me that face_areas computations could be sped up dramatically.

    Updating the table from before to include the optimized numba methods:

    grid (optimized) numba 4 threads (optimized) numba 1 thread (unoptimized) numba 4 threads (unoptimized) numpy (unoptimized) numba 1 thread ux.open_grid()
    MPAS 30km 0.07 s 0.18 s 0.53 s 1.23 s 1.79 s ~2.7 s
    MPAS 15km 0.29 s 0.71 s 1.94 s 4.99 s 7.07 s ~18 s
    MPAS 7.5km 1.00 s 2.78 s 7.64 s 12.31 s 29.03 s ~65 s
    MPAS 3.75km 3.08 s 11.0 s 29.4 s 48.8 s 115.56 s ~250 s

    This order-of-magnitude improvement bringing MPAS 3.75 km down to roughly 3 seconds makes me think caching is no longer necessary to consider for now. I would rather wait 3 seconds than use an extra 2 GB of memory. It would be easy enough to add caching later if this ever becomes part of a hot path, but these speeds are now fast enough that it probably makes sense to ship compute_face_node_angles initially without a cache option.

    Additionally, let's keep the compute_face_node_angles() syntax for now, and not introduce a face_node_angles property unless it becomes clearly necessary or helpful. Other properties (e.g. face_areas) tend to be cached, so it might cause some confusion if a non-cached property gets introduced.

  24. erogluorhan commented on Aug 11, 2026

    @erogluorhan
    Member

    Looks great! No caching for now makes complete sense to me. When are you planning a PR on this?

  25. Sevans711 commented on Aug 12, 2026

    @Sevans711
    CollaboratorAuthor

    Looks great! No caching for now makes complete sense to me. When are you planning a PR on this?

    Just made it, see #1672! I was hoping to get to it yesterday but had to head out before I could finish setting up tests and writing up the PR overview / description.

  26. moved this from 🏗 In progress to ✅ Done in UXarray Developmenton Aug 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

new featureNew user-facing functionality

Type

No type

Fields

Priority

None yet

Projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions