Unpack NumPy array by column

52.7k views Asked by At

If I have a NumPy array, for example 5x3, is there a way to unpack it column by column all at once to pass to a function rather than like this: my_func(arr[:, 0], arr[:, 1], arr[:, 2])?

Kind of like *args for list unpacking but by column.

3

There are 3 answers

0
Alex Riley On BEST ANSWER

You can unpack the transpose of the array in order to use the columns for your function arguments:

my_func(*arr.T)

Here's a simple example:

>>> x = np.arange(15).reshape(5, 3)
array([[ 0,  5, 10],
       [ 1,  6, 11],
       [ 2,  7, 12],
       [ 3,  8, 13],
       [ 4,  9, 14]])

Let's write a function to add the columns together (normally done with x.sum(axis=1) in NumPy):

def add_cols(a, b, c):
    return a+b+c

Then we have:

>>> add_cols(*x.T)
array([15, 18, 21, 24, 27])

NumPy arrays will be unpacked along the first dimension, hence the need to transpose the array.

0
Stephanie On

numpy.split splits an array into multiple sub-arrays. In your case, indices_or_sections is 3 since you have 3 columns, and axis = 1 since we're splitting by column.

my_func(numpy.split(array, 3, 1))
1
CookieMaster On

I guess numpy.split will not suffice in the future. Instead, it should be

my_func(tuple(numpy.split(array, 3, 1)))

Currently, python prints the following warning:

FutureWarning: Using a non-tuple sequence for multidimensional indexing is deprecated; use arr[tuple(seq)] instead of arr[seq]. In the future this will be interpreted as an array index, arr[np.array(seq)], which will result either in an error or a different result.