Search code examples
pythonarraysnumpymatrixarray-broadcasting

Applying function to 2 vectors to obtain a 2D matrix

My question is exactly the same as this one but in Python instead of R. I couldn't find a similar question for python and so I'm asking this one.

I have 2 vectors (lists actually)

a = [1,2,3]
b = [5,6,7]

A defined function

def calc(i, j):
    return (i+j)/2

I want to apply the calc function to all the pairs of values in the vectors so that i obtain a matrix (preferably using numpy package).

The expected result in this case is

result = [
           [3,   3.5,   4],
           [3.5,  4,  4.5], 
           [4,   4.5,   5]
         ]

I would like to avoid list comprehensions and for loops to achieve this result.

Solution

  • You can use meshgrid

    import numpy as np

a = [1, 2, 3]
b = [5, 6, 7]

def calc(i, j):
    return (i + j) / 2

A, B = np.meshgrid(a, b)

result = calc(A, B)

print(result)

    output:

    [[3.  3.5 4. ]
 [3.5 4.  4.5]
 [4.  4.5 5. ]]