Optimization: don't construct many tiny numpy arrays inside numba routines

[Sevans711]

Allocating lots of tiny numpy arrays is expensive compared to most other operations being performed in numba routines. Discovered during #1566, which also shows one example where this optimization caused >10x speedup: #1566 (comment). (Some part of the speedup quoted there might be from fastmath=True, which should generally be avoided in uxarray, but the majority of it is due to avoiding construction of many tiny numpy arrays.)

Here are a few examples where this occurs in uxarray right now:

• _compute_gradients_on_faces (in gradient.py) constructs np.zeros(3) and result of np.cross for many pairs of faces. Each of those is a tiny numpy array being constructed.
• _sort_points_by_angle (in zonal.py) constructs np.array for every point.
• _small_angle_of_2_vectors (in grid/utils.py) calls np.linalg.norm, which will create a numpy array if inputs weren't a numpy array already. It doesn't have a loop internally, but it is used inside loops in other functions (see thread of issue linked above), where the inputs are just vectors with 3 components, like [ux, uy, uz], with no extra dimensions (it isn't being used in a vectorized manner).

There may be more examples throughout the code. A complete fix to this issue should attempt to find and address all of them.

Suggested fix:

• Avoid constructing tiny numpy arrays inside numba functions. For example, work with ux, uy, uz components directly, using value = (ux**2 + uy**2 + uz**2)**0.5, instead of using value = np.linalg.norm(u).
• Perhaps add a few helper functions, e.g. value = _numba_norm3(ux, uy, uz) instead of value = (ux**2 + uy**2 + uz**2)**0.5. That could make the changes easier to implement (closer to existing syntax) and maintain (one source of truth for the norm function, easier to avoid easy-to-miss typos like (ux**2 + uy**2 + ux**2)**0.5).

Tagging @cmdupuis3 in case you want to claim this one after it gets triaged! (I'd also be happy to do it, but would certainly end up asking for you to review any related PRs.)

[cmdupuis3]

Yeah, this seems to be a common failure mode here, my advice would be to file PRs at a fine grain, so it's not too overwhelming. I can definitely review if you have something specific though.

[Sevans711]

Sounds good. Maybe the first PR could add some helper methods (I'm trying to think where they belong… my intuition would be to put them in a new file near top-level, something like uxarray/utils/numba_math.py?) and improve one of the examples here (such as _compute_gradients_on_faces) so it is easy to check for improvements in benchmarks.

Subsequent PRs shouldn't try to solve everything at once, but could maybe fix this for groups of related functions (e.g., _small_angle_of_2_vectors and the handful of functions that use it?) to avoid having too many tiny PRs, if that sounds like a reasonable plan?

[cmdupuis3]

I notice a lot of the times the fix is to transpose arrays-of-structs to structs-of-arrays (or similar since you probably don't need the struct at that point). If we want to be fancy, we could define index operator overloads so we keep the pointwise indexing bevavior, like x, y, z = list_of_points[i], but I would guess in most cases these funtions are low-level enough that you wouldn't want that either.

[Sevans711]

I notice a lot of the times the fix is to transpose arrays-of-structs to structs-of-arrays (or similar since you probably don't need the struct at that point). If we want to be fancy, we could define index operator overloads so we keep the pointwise indexing bevavior, like x, y, z = list_of_points[i], but I would guess in most cases these funtions are low-level enough that you wouldn't want that either.

Not sure if I understand this. If it's important to understand before starting work on this issue, would you be able to provide a more complete example or maybe rephrase?

[cmdupuis3]

An array of structs is like:

@dataclass
class Point:
    x: float
    y: float
    z: float

points = [Point(1.0, 2.0, 3.0), Point(4.0, 5.0, 6.0), ...]

This usually indicates a performance issue because structs as a whole occupy more storage than the individual values they contain. The transpose is to have all x, y, and z values in an array each:

class PointList:
    x: np.array
    y: np.array
    z: np.array

    def index(n):
        return x[n], y[n], z[n]

So here, if you know you'll index in order, you would have much better caching behavior.

[Sevans711]

Ah, that helped, thank you! Indeed, the numba implementations now are sort of like doing an array of structs. Confusingly, the "structs" in this analogy are actually the numpy arrays (because one is being constructed per point).

I agree with your original assessment that these functions are low-level enough that we probably also don't want to introduce a new thing like PointList here. (Although, I think the python equivalent would be something like grid["node_x"], grid["node_y"], grid["node_z"], which we already have where relevant, right?)