Explore our sites

Stack Exchange
tour help careers 2.0
Take the 2-minute tour ×
Stack Overflow is a question and answer site for professional and enthusiast programmers. It's 100% free, no registration required.
up vote 4 down vote favorite

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?

share|improve this question
 
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 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 at 13:05
add comment

2 Answers

up vote 2 down vote accepted

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.

share|improve this answer
 
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 at 17:06
1  
It´s the ellipsis notation. Corrected the typo. –  Jaime Jun 17 at 7:22
 
I see. Now I see some notes on it: stackoverflow.com/a/773472/907575 and stackoverflow.com/questions/118370/…. –  Piotr Migdal Jun 17 at 10:41
add comment
up vote 3 down vote

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]]])
share|improve this answer
 
I would avoid defining x,y,z before to save memory, calling the functions inside np.dstack() –  Saullo Castro Jun 15 at 18:45
 
@sgpc -- You are right. I only did it for greater clarity. –  root Jun 15 at 18:52
 
+1 Nice and works. I accepted Jaime's answer as it is more convenient for my purpose and for some generalizations. –  Piotr Migdal Jun 16 at 17:09
add comment

Your Answer

 
discard

By posting your answer, you agree to the privacy policy and terms of service.

Not the answer you're looking for? Browse other questions tagged or ask your own question.