从返回每个(N，N)子数组的函数构建(...，N，N)numpy数组

我想使用一个返回每个（N，N）块的函数来构建一个形状为（...，N，N）的numpy数组。总的来说，构建（...，N，N）数组的速度要比在嵌套的for循环中迭代第一个（...）索引快得多。
N=2的最小示例如下所示：

a = 1.
b = 2 + 3.j
def h0(pi):
    return np.array([
        [a,        b * pi ],
        [b * pi,   a      ]
    ])

字符串
例如，我想传递pi的（M，M）数组值，并得到一个（M，M，N，N）数组（称之为y）作为输出。我想构建它，使y[i,j,k,l] = h0(pi[i,j])[k,l]不只是循环i和j。
很明显我做的是错的-我祈祷得到了我应得的。我有一种直觉，当我传递一个数组给pi时，我试图以一种不连贯的方式嵌套数组。但是我没有足够的实践/知识来看到这种事情的优雅解决方案。
任何建议都很感激！
我想做的有点类似于this question，但是建议使用列表解析会违背目的（避免任何类型的python循环）。
编辑：问题是，我需要显式地设置元素a，其中没有pi，以采用与pi相同的形状。当我弄清楚如何干净地做到这一点时，将最后一次更新。（感谢@chrslg帮助我意识到这是问题所在。）
编辑与决议：实际上有两个步骤来解决这个问题。首先，我必须广播a到正确的大小。@chrslg在下面提供了一种方法。另一种方法是在函数中定义one = np.ones_like(pi)，并将其乘以数组定义中出现的a。然后我只需使用np.transpose将（N，N）轴放在最后。正确的代码如下所示

a = 1.
b = 2 + 3.j
def h0(pi):
    one = np.ones_like(pi)
    return np.array([
        [a * one, b * pi ],
        [b * pi,  a * one]
    ]).transpose((2,3,0,1))

型
当pi是一个二维数组时。

只是为了澄清，因为注解不适合提供代码，这里是我的工作代码

import numpy as np
import timeit

a=3
b=2+3.j

def h0(pi):
    return np.array([
        [a * pi, b * pi.conjugate()],
        [b.conjugate() * pi, a * pi]
    ])

nK = 51
cutoff = 0.5
def f1():
    return np.fromfunction(
        lambda jX, jY: cutoff * ((-1 + 2*jX/(nK-1)) + 1.j * (-1 + 2*jY/(nK-1))),
        (nK, nK),
        dtype=complex
    )

piff = f1()
t1=timeit.timeit(f1, number=10000)

def f2():
    jX=np.arange(nK)[:, None]
    jY=np.arange(nK)[None, :]
    return cutoff * ((-1 + 2*jX/(nK-1)) + 1.j * (-1 + 2*jY/(nK-1)))

pi2=f2()
t2=timeit.timeit(f2, number=10000)

print("fromfunc timing:", t1, "broadcast timing", t2)
print("result difference", np.abs(piff-pi2).sum())

h0grid=h0(pi2)
print("h0grid shape", h0grid.shape)

字符串
我也（因为我在你改变问题之前就开始输入了）包括对什么是fromfunc的澄清。基本上fromfunc所做的就是我在f2中所做的。（不只是同样的结果。同样的方法：为jX和jY创建数组，并使用这些数组作为参数调用lambda。如果lambda函数尚未“向量化”，f1和f2都将失败，它既可以使用标量参数，也可以使用数组参数）。
你可以看到，在时间上，f2更快，（在我的机器上快3倍。所以值得一试）。这仍然是一个日志，足以避免使用fromfunc时，你可以做我的方式。但是，然而，这不是一个很大因素，因为如果fromfunc是一个函数，它调用lambda 51x51次来创建一个51x51结果的数组，（通常当人们做这样的事情时，使用for循环，或使用vectorize，或使用applyalongaxis，或.，我们给予他们一种方法来使用向量化函数，时间因子更多的是1000的数量级，而不是3。所以“只有3个因子”证明了fromfunc确实不比我的f2多。至于为什么它更慢，不确定，但我会说这可能是因为它为jX和jY构建了2个51x51数组，而我使用51x1和1x51数组，并让广播做它的事情。并不是说构建51x51数组花费了很多。但由于广播并不发生在所有操作中，只有当jX和jy合并时，对于计算的某些部分，它是51倍的操作）
所以，大括号，但因为我已经写了代码...
现在，无论你用什么方法创建51x51网格，只要你的函数将标量Map到标量的2x2矩阵是矢量化的，它也将51x51矩阵Map到51x51矩阵的2x2矩阵。也就是2x2x51x51数组。
如果它在你的例子中不起作用，那一定是因为你的实际代码和你发布的代码之间的差异。因为，这是你的代码。它起作用了。

从返回每个(N，N)子数组的函数构建(...，N，N)numpy数组

1条答案

相关问题

热门标签

最新问答