Test code for a function that applies functions in a sequence over arbitrary selections from an array.

indexed_function_application_test.py

import timeit

setup = '''
import numpy as np

def apply_indexed_fast(array, func_indices, func_table):
    func_argsort = func_indices.argsort()
    func_ranges = list(np.searchsorted(func_indices[func_argsort], range(len(func_table))))
    func_ranges.append(None)
    out = np.zeros_like(array)
    for f, start, end in zip(func_table, func_ranges, func_ranges[1:]):
        ix = func_argsort[start:end]
        out[ix] = f(array[ix])
    return out

def apply_indexed_med(array, funciton_indices, func_table):
    idx_funcsort = np.argsort(function_indices)
    unique_funcs, unique_func_indices = np.unique(function_indices[idx_funcsort], return_index=True)
    desired_output = np.zeros_like(array)
    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])
    return desired_output

def apply_indexed_slow(array, function_indices, func_table):
    desired_output = np.zeros_like(array)
    for func_index in set(function_indices):
        idx = np.where(function_indices==func_index)[0]
        desired_output[idx] = func_table[func_index](array[idx])
    return desired_output

def trivial_functional(i): 
    return lambda x : i*x

k = 250
func_table = [trivial_functional(j) for j in range(k)]
func_table = np.array(func_table) # possibly unnecessary

Npts = 1e6
abcissa_array = np.random.random(Npts)
function_indices = np.random.random_integers(0,len(func_table)-1,Npts)
func_array = func_table[function_indices]
'''

stmt_fast = 'a = apply_indexed_fast(abcissa_array, function_indices, func_table)'
stmt_med = 'b = apply_indexed_med(abcissa_array, function_indices, func_table)'
stmt_slow = 'c = apply_indexed_slow(abcissa_array, function_indices, func_table)'

print "apply_indexed_fast, 50 iterations:"
print timeit.timeit(setup=setup, stmt=stmt_fast, number=10)
print "apply_indexed_med, 50 iterations:"
print timeit.timeit(setup=setup, stmt=stmt_med, number=10)
print "apply_indexed_slow, 5 iterations:"
print timeit.timeit(setup=setup, stmt=stmt_slow, number=10)

exec(setup)
exec(stmt_fast)
exec(stmt_med)
exec(stmt_slow)

print "apply_indexed_fast sample results:"
print a[0:10]
print "apply_indexed_med sample results:"
print b[0:10]
print "apply_indexed_slow sample results:"
print c[0:10]