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?

2
  • 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. Commented Jun 15, 2013 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). Commented Jun 16, 2013 at 13:05

2 Answers 2

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.

Sign up to request clarification or add additional context in comments.

3 Comments

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.)
It´s the ellipsis notation. Corrected the typo.
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]]])

2 Comments

I would avoid defining x,y,z before to save memory, calling the functions inside np.dstack()
+1 Nice and works. I accepted Jaime's answer as it is more convenient for my purpose and for some generalizations.

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.