Numpy: how to access axis 0 elements in a (2, N) matrix to be used as function arguments-CodePudding

I have a Numpy 2D matrix with shape (2, N). Let arr be:

arr = np.array([[2, 4, 6], [10, 20, 30]])
print(arr.shape)
arr
(2, 3)
array([[ 2,  4,  6],
       [10, 20, 30]])

If I have a function that, say, adds two numbers (yes, I know we can use broadcasting, but that's not the point here), I'd like to be able to access the axis 0 elements in arr to pass into the function.

For example:

def add_two(a, b):
    return a   b

How do I use iteration to access the pairs (2, 10), (4, 20), and (6, 30) to be used as arguments for add_two()?

For example:

for elem in arr.shape[1]:
    print('The sum is', add_two(arr...))

Desired output:

The sum is 12
The sum is 24
The sum is 36

Thanks!

CodePudding user response：

This is quite easy.

for (a,b) in arr.T:
    print('The sum is', add_two(a,b))

CodePudding user response：

Try this:

import numpy as np

arr = np.array([[2, 4, 6], [10, 20, 30]])

def add_two(a, b):
    return a   b
    
for vals in zip(*arr):
    print('The sum is', add_two(vals[0], vals[1]))

The crux of this is zip(*arr).

Explanation: zip takes an arbitrary number of arguments and "zips" them together. Ex: zip([1, 2, 3], [4, 5, 6]) outputs a zip object in the form [(1, 4), (2, 5), (3, 6)]. The * operator before arr "unpacks" the values of arr into zip, so zip(*arr) becomes zip([2, 4, 6], [10, 20, 30])

CodePudding user response：

Your column index iteration could be written as:

In [242]: for elem in range(arr.shape[1]):
     ...:     print('The sum is', add_two(arr[0,elem],arr[1,elem]))
     ...:     
The sum is 12
The sum is 24
The sum is 36

While iterating on indices is allowed in python, iteration on that values is more idiomatic (and using enumerate if you also need the indices). So:

In [243]: for i,j in zip(arr[0,:],arr[1,:]):
     ...:     print('The sum is', add_two(i,j))
     ...:     
The sum is 12
The sum is 24
The sum is 36

zip(a,b) is a list form of transpose.

Iteration on an array treats the array as a list (on the first dimension). Transpose makes the second dimension the iteration one:

In [244]: for i,j in arr.T:
     ...:     print('The sum is', add_two(i,j))
The sum is 12
The sum is 24
The sum is 36

But since you already have a 2d array, you can just use sum with the axis specification:

In [245]: arr.sum(axis=0)
Out[245]: array([12, 24, 36])