如何基于生成器创建数组？

Question

说我有一个大型数组：

A = 2*np.ones([100, 100, 100])

例如，我想对其进行一些计算：

def squared_elements(M):
    yield M**2

我选择使用生成器函数，因为我的数组很大，并且我不需要所有结果。事实上，我只需要在矩阵中心放置一个长度为10的立方体。

如果是正常功能，我可以写：

result = squared_elements(A)[45:55, 45:55, 45:55]

但是，生成器不可下标，因此最后一个表达式不起作用。

如何使用生成器函数获得相同的result？

Answer 1

解决方案

result = next(squared_elements(A))
result[45:55, 45:55, 45:55]

已解决问题的解决方案

您可以在平方之前对子集进行切片

def squared_elements(A, slice_):
    return A[slice_] ** 2

result = squared_elements(A, [slice(45, 55), slice(45, 55), slice(45, 55)])

注意

• next() here的更多信息。
• 有关生成器功能here的更多信息。

我希望这会有所帮助！