Efficient algorithm for evaluating a 1-d array of functions on a same-length 1d numpy array

后端 未结 3 1929
刺人心
刺人心 2021-01-19 18:18

I have a (large) length-N array of k distinct functions, and a length-N array of abcissa. I want to evaluate the functions at the abcissa to return a length-N array of ordin

3条回答
  •  夕颜
    夕颜 (楼主)
    2021-01-19 19:13

    Thanks to hpaulj for the suggestion to pursue a groupby approach. There are lots of canned routines out there for this operation, such as Pandas DataFrames, but they all come with the overhead cost of the data structure initialization, which is one-time-only, but can be costly if using for just a single calculation.

    Here is my pure numpy solution that is a factor of 13 faster than the original where loop I was using. The upshot summary is that I use np.argsort and np.unique together with some fancy indexing gymnastics.

    First we sort the function indices, and then find the elements of the sorted array where each new index begins

    idx_funcsort = np.argsort(function_indices)
    unique_funcs, unique_func_indices = np.unique(function_indices[idx_funcsort], return_index=True)
    

    Now there is no longer a need for blind lookups, since we know exactly which slice of the sorted array corresponds to each unique function. So we still loop over each called function, but without calling where:

    for func_index in range(len(unique_funcs)-1):
        idx_func = idx_funcsort[unique_func_indices[func_index]:unique_func_indices[func_index+1]]
        func = func_table[unique_funcs[func_index]]
        desired_output[idx_func] = func(abcissa_array[idx_func])
    

    That covers all but the final index, which somewhat annoyingly we need to call individually due to Python indexing conventions:

    func_index = len(unique_funcs)-1
    idx_func = idx_funcsort[unique_func_indices[func_index]:]
    func = func_table[unique_funcs[func_index]]
    desired_output[idx_func] = func(abcissa_array[idx_func])
    

    This gives identical results to the where loop (a bookkeeping sanity check), but the runtime of this loop is 0.027 seconds, a speedup of 13x over my original calculation.

提交回复
热议问题