10

I want map a numpy.array from NxM to NxMx3, where a vector of three elements is a function of the original entry:

lambda x: [f1(x), f2(x), f3(x)]

However, things like numpy.vectorize do not allow to change dimensions. Sure, I can create an array of zeros and make a loop (and it is what I am doing by now), but it does not sound neither Pythonic nor efficient (as every looping in Python).

Is there a better way to perform an elementwise operation on numpy.array, producing a vector for each entry?

Community
  • 1
  • 1
Piotr Migdal
  • 11,864
  • 9
  • 64
  • 86
  • If `N` and `M` are significantly larger than 3, the looping over the third dimension will have an insignificant effect on performance. And there's nothing un-pythonic in using for loops! What isn't very numpythonic or efficient is using `np.vectorize`. You could try to convert `f1`, `f2` and `f3` into a single function that took arrays and returned arrays. Without knowing what your functions are doing, it is not possible to know if this approach would suit your problem. – Jaime Jun 15 '13 at 14:07
  • @Jaime I am looping over N and M, not 3. The problem is conversion of complex numbers in three floats [R, G, B], so I can plot a complex function (see the link in the question). – Piotr Migdal Jun 16 '13 at 13:05

2 Answers2

4

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 axis
to make a single array. Rebuilds arrays divided by `dsplit`.
This is a simple way to stack 2D arrays (images) into a single
3D 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]]])
root
  • 76,608
  • 25
  • 108
  • 120
4

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.

Jaime
  • 65,696
  • 17
  • 124
  • 159
  • I didn't knew that it is possible to reference subarrays as `Y[..., i]` or `Y[i, ...]`. Thanks! Is there a particular name for this operation? (I was googling *vectorization*, but it returns things like `np.vectorize`.) – Piotr Migdal Jun 16 '13 at 17:06
  • 1
    It´s the _ellipsis_ notation. Corrected the typo. – Jaime Jun 17 '13 at 07:22
  • I see. Now I see some notes on it: http://stackoverflow.com/a/773472/907575 and http://stackoverflow.com/questions/118370/how-do-you-use-the-ellipsis-slicing-syntax-in-python. – Piotr Migdal Jun 17 '13 at 10:41