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

Question

I would like to produce a 4D array from a 2D one by periodic shifts, in a way that can be summarized by the following:

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

This is easiest to illustrate with a "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)]

Here, the periodicity condition pc just brings the index back into the allowed range. The output for Npts=4 is:

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

So that each value is shifted slightly. For the "4D from 2D" case, I could obviously use:

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)]

However, using four loops is going to be slow, as I will be doing this multiple times for much larger arrays. How can I do this more efficiently?

Please note that, although the MWE example could be reproduced by applying the square function to a 2D array, that would not be a helpful solution. Using the MWE to illustrate, the goal is to apply the function as few times as possible (ie only on the 1D array) and then to create the 2D array without for loops. Ultimately, I will need to do this to generate a 4D array from a 2D array. How can I do this?

Answer 1

You can replicate the 2D array and then extract the shifted 2D sub-arrays (avoiding modulus and conditionals). Here is how to do that:

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]

With Npts=64 , this solution is about 1000 times faster .