时间戳

首页 > 解决方案 > 从另一个具有移位索引的数组中创建一个 NumPy 数组

python - 从另一个具有移位索引的数组中创建一个 NumPy 数组

问题描述

我想通过周期性移位从 2D 数组生成 4D 数组，其方式可以总结如下：

uuvv[kx,ky,qx,qy] = uu[kx+qx,ky+qy]

这最容易用“2D from 1D”MWE 来说明：

def pc(idx):
    return idx - Npts*int(idx/Npts)

uu = np.square(np.arange(Npts))
uv = np.zeros((Npts,Npts))
for kx in np.arange(Npts):
    for qx in np.arange(Npts):
        uv[kx,qx] = uu[pc(kx+qx)]

在这里，周期性条件pc只是将索引带回到允许的范围内。的输出Npts=4是：

array([[0., 1., 4., 9.],
       [1., 4., 9., 0.],
       [4., 9., 0., 1.],
       [9., 0., 1., 4.]])

这样每个值都会稍微移动。对于“2D 中的 4D”案例，我显然可以使用：

def pbc(idx):
    return idx - Npts*int(idx/Npts)

uv = np.zeros((Npts,Npts,Npts,Npts))
for kx in np.arange(Npts):
    for ky in np.arange(Npts):
        for qx in np.arange(Npts):
            for qy in np.arange(Npts):
                uv[kx,ky,qx,qy] = uu[pbc(kx+qx),pbc(ky+qy)]

但是，使用四个循环会很慢，因为我将多次为更大的数组执行此操作。我怎样才能更有效地做到这一点？

请注意，尽管可以通过将 square 函数应用于 2D 数组来重现 MWE 示例，但这不是一个有用的解决方案。使用 MWE 来说明，目标是尽可能少地应用该函数（即仅在 1D 数组上），然后创建没有 for 循环的 2D 数组。最终，我需要这样做以从 2D 数组生成 4D 数组。我怎样才能做到这一点？

标签： pythonarraysperformancenumpy

解决方案

您可以复制 2D 数组，然后提取移位的 2D 子数组（避免模数和条件）。以下是如何做到这一点：

uuRep = np.tile(uu, (2,2))
uv = np.zeros((Npts,Npts,Npts,Npts))
for kx in np.arange(Npts):
    for ky in np.arange(Npts):
        uv[kx,ky,:,:] = uuRep[kx:kx+Npts,ky:ky+Npts]

有了Npts=64这个，这个解决方案就快了大约1000 倍。

推荐阅读