python - 一个函数需要什么才能使用堆叠的 NumPy 参数？

Question

使用如下定义的网格进行绘图时:

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

例如这个功能

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

可以处理 polynomial(X,Y)但不是

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

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

而另一方面，例如SciPy.stats multivariate_normal 的 pdf

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

不能处理

normal(X,Y)

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

但它可以处理 normal(stack) .两者都是两个变量的函数，但它们接受参数的方式显然非常不同。
我必须对多项式进行哪些更改才能接受堆叠参数，例如 normal能够？

Answer 1

查看数组形状:

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)

那是2个数组。在新的尾随轴上加入它们:

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

另一种方法，使用新的引导轴:

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

ogrid制作 sparse一对:

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)

您的 function可以处理这些，因为它们的行为与 X 相同和 Y关于 + 等广播运营商:

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

堆叠阵列可以通过以下方式使用:

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

该函数采用 2 个数组 - 这在定义中是明确的。
签名 pdf是(来自文档)

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

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

它不接受第二个位置参数，但您可以指定添加的关键字参数。它如何处理 x 的维度是函数内部的(不是其签名的一部分)，但大概是 dstack网格作品，具有“2个组件”。
每个函数都有自己的签名。您不能假定一种模式适用于另一种模式。保留 文档 在眼前!
进一步挖掘， pdf通过 x虽然一个函数

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

此函数定义应处理两种形式(未测试):

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