Threaded sum for large numpy arrays

[bnikolic]

Using this makes runtime of mean and other similar calculations about 5-7 times faster on large images such as 24k x 24k that I use for the SKA reference run. Overall it improves performance by ~~15% in the current run (consistent with expectations given time in these functions and that the submaps do not gain from this optimisation)

[darafferty]

Very nice! These changes work well in my tests (I didn't do any benchmarking though). @gmloose what do you think?

[gmloose]

I asked ChatGPT to review this code, and it came up with a slightly different implementation. Most importantly related to the chunking (covering the edge case where a.size is an exact multiple of N), and some other changes (e.g., replace the lambda with a local function for readability).

def sum_threaded(a, N=10_000_000):
    if a.size < N:
        return a.sum()

    n = (a.size + N - 1) // N

    def partial_sum(i):
        return a[i*N:(i+1)*N].sum()

    with ThreadPoolExecutor() as te:
        res = te.map(partial_sum, range(n))
        return numpy.sum(list(res))


def mean_threaded(a):
    return sum_threaded(a) / a.size

So, choose your poison 😄. But I would at least keep the improved chunking.

[bnikolic]

This looks fine and it saves a thread invocation in the edge case (although for a sum on empty array so not a big saving!). Need to test it though just in case, I can have a look at that when I clear the urgent things accumulated over the holidays.