Mapping element-wise a NumPy array into an array of more dimensions

Now that I see your code, for most simple mathematical operations you can let numpy do the looping, what is often referred to as vectorization:

def complex_array_to_rgb(X, theme='dark', rmax=None):    '''Takes an array of complex number and converts it to an array of [r, g, b],    where phase gives hue and saturaton/value are given by the absolute value.    Especially for use with imshow for complex plots.'''    absmax = rmax or np.abs(X).max()    Y = np.zeros(X.shape + (3,), dtype='float')    Y[..., 0] = np.angle(X) / (2 * pi) % 1    if theme == 'light':        Y[..., 1] = np.clip(np.abs(X) / absmax, 0, 1)        Y[..., 2] = 1    elif theme == 'dark':        Y[..., 1] = 1        Y[..., 2] = np.clip(np.abs(X) / absmax, 0, 1)    Y = matplotlib.colors.hsv_to_rgb(Y)    return Y

This code should run much faster than yours.

If I understand your problem correctly, I suggest you use np.dstack:

Docstring:Stack arrays in sequence depth wise (along third axis).Takes a sequence of arrays and stack them along the third axisto make a single array. Rebuilds arrays divided by `dsplit`.This is a simple way to stack 2D arrays (images) into a single3D array for processing.

    In [1]: a = np.arange(9).reshape(3, 3)    In [2]: a    Out[2]:     array([[0, 1, 2],           [3, 4, 5],           [6, 7, 8]])    In [3]: x, y, z = a*1, a*2, a*3  # in your case f1(a), f2(a), f3(a)     In [4]: np.dstack((x, y, z))    Out[4]:     array([[[ 0,  0,  0],            [ 1,  2,  3],            [ 2,  4,  6]],           [[ 3,  6,  9],            [ 4,  8, 12],            [ 5, 10, 15]],           [[ 6, 12, 18],            [ 7, 14, 21],            [ 8, 16, 24]]])