[FIXED] What is necessary for a function to be able to take stacked NumPy arguments?

Issue

While plotting with a meshgrid defined like this:

Y,X = np.mgrid[-5:5:20j, -5:5:20j]

For example this function

def polynomial(x,y):
    return x**2 + y**2

can handle polynomial(X,Y) but not

stack = np.dstack((X,Y))
polynomial(stack)

--->
Type error: polynomial missing 1 required positional argument: 'y'

while on the other hand e.g. the pdf of SciPy.stats multivariate_normal

mu = [0,0]
sigma = [[3,2]
         [2,3]]
normal = st.multivariate_normal(mu, sigma)
normal = normal.pdf

can't handle

normal(X,Y)

--->
Type error: pdf() takes 2 positional arguments but 3 were given

but it can handle normal(stack). Both are functions of two variables but the way they accept arguments is apparantly very different.

What changes would I have to make to polynomial such that it can accept stacked arguments like normal can?

Solution

Look at the array shapes:

In [165]: Y,X = np.mgrid[-5:5:20j, -5:5:20j]
In [166]: Y.shape
Out[166]: (20, 20)
In [167]: X.shape
Out[167]: (20, 20)

That's 2 arrays. Joining them on a new trailing axis:

In [168]: stack = np.dstack((X,Y))
In [169]: stack.shape
Out[169]: (20, 20, 2)

An alternative way, with a new leading axis:

In [170]: stack1 = np.stack((X,Y))
In [171]: stack1.shape
Out[171]: (2, 20, 20)

ogrid makes a sparse pair:

In [172]: y,x = np.ogrid[-5:5:20j, -5:5:20j]
In [173]: y.shape
Out[173]: (20, 1)
In [174]: x.shape
Out[174]: (1, 20)

Your function can handle these, since they behave the same as X and Y with respect to broadcasting operators like +:

In [175]: def polynomial(x,y):
     ...:     return x**2 + y**2
     ...: 
In [176]: polynomial(x,y).shape
Out[176]: (20, 20)

The stacked arrays can be used via:

In [177]: polynomial(stack[...,0],stack[...,1]).shape
Out[177]: (20, 20)
In [178]: polynomial(stack1[0],stack1[1]).shape
Out[178]: (20, 20)

The function takes 2 arrays - that's explicit in the definition.

The signature of pdf is (from the docs)

pdf(x, mean=None, cov=1, allow_singular=False)

x : array_like
    Quantiles, with the last axis of `x` denoting the components.

It doesn't accept a second positional argument, though you can specify added keyword ones. How it handles the dimensions of x is internal to the function (not part of its signature), but presumably the dstack grid works, with "2 components".

Each function has its own signature. You can't presume that the pattern for one applies to another. Keep the docs at hand!

digging further, pdf passes the x though a function that

Adjust quantiles array so that last axis labels the components of
    each data point.

This function definition should handle both the forms (not tested):

def polynomial(x,y=None):
    if y is None:
        x, y = x[...,0], x[...,1]
    return x**2 + y**2

Answered By - hpaulj