Function that takes in a meshgrid of values

user_14856539 0
Viewed 88

I want a function that takes in 2 numbers and a matrix and does a multiplication operation using them and then sums every element of the matrix after multiplication (see code below).

I am trying to do this operation for all possible combinations of x and y. What is the best way to do this without using loops, since I know I could loop over the function with different values of x and y but this doesn't seem like an efficient way to do it.

I tried using a meshgrid as input but that didnt work due to the way the input is broadcasted.

import numpy as np

def my_func(num1, num2, matrix2):
    return np.sum(matrix*num1*num2)

x = np.linspace(0,5)
y = np.linspace(0,1)

X, Y = np.meshgrid(x, y)
matrix = np.array([[1],[2],[3]])

a = my_func(X,Y,matrix)

I get the following error:

ValueError: operands could not be broadcast together with shapes (50,50) (3,1)

I would like a to equal a meshgrid of values where each value in the array corresponds to the output of my_func for every possible combination of x and y.

numpy python
2 Answers

The result of x * y * M when x and y are scalar is just M.shape. If you want this result for each value of x and y, you will want a result of shape x.shape + y.shape + M.shape. You can do this with broadcasting for the totally general case. The idea is that you need to reshape x to have trailing ones too fill in y.ndim + M.ndim dimensions and y to have M.ndim trailing dimensions.

For the sake of the summation, it's actually easier to ravel M, even though np.sum allows for multiple axes since version 1.7.0.

def my_func(x, y, matrix):
    x = np.reshape(x, x.shape + (1,) * (y.ndim + 1))
    y = np.reshape(y, y.shape + (1,))
    return (x * y * matrix.ravel()).sum(axis=-1)

If you want to input x and y that are already broadcasted together, you can adjust the calculation slightly:

def my_func(x, y, matrix):
    return ((x * y)[:, None] * matrix.ravel()).sum(-1)

The conceptual difference is that the first version accepts the linspaces you created directly, while the second version requires you to construct the meshgrid, or at least transpose one of the arrays.

user_29887300

It looks like you have to adjust the dimensions of your x and y array:

import numpy as np

def my_func(num1, num2, matrix2):
    return matrix*num1*num2

x = np.linspace(0,5, num=3)
y = np.linspace(0,1, num=3)

X, Y = np.meshgrid(x, y)
matrix = np.array([[1],[2],[3]])

a = my_func(X,Y,matrix)
print(a)
# [[ 0.   0.   0. ]
#  [ 0.   2.5  5. ]
#  [ 0.   7.5 15. ]]
user_56603150
Related