CODEWITHSUNDEEP
pythonnumpy
NahPlsMan
NahPlsMan

Reputation: 33

Function that takes in a meshgrid of values

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.

Upvotes: 2

Views: 696

Answers (2)

Mad Physicist
Mad Physicist

Reputation: 114320

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.

Upvotes: 2

VirtualScooter
VirtualScooter

Reputation: 1888

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. ]]

Upvotes: 0

Related Questions