CODEWITHSUNDEEP
pythonarraysnumpynumpy-ndarray
chuckstables
chuckstables

Reputation: 165

Rolling/Increasing dimensionality of a NumPy array

I'm currently trying to find an easy way to do the following operation to an N dimensional array in Python. For simplicity let's start with a 1 dimensional array of size 4.

X = np.array([1,2,3,4])

What I want to do is create a new array, call it Y, such that:

Y = np.array([1,2,3,4],[2,3,4,1],[3,4,1,2],[4,1,2,3])

So what I'm trying to do is create an array Y such that:

Y[:,i] = np.roll(X[:],-i, axis = 0)

I know how to do this using for loops, but I'm looking for a faster method of doing so. The actual array I'm trying to do this to is a 3 dimensional array, call it X. What I'm looking for is a way to find an array Y, such that:

Y[:,:,:,i,j,k] = np.roll(X[:,:,:],(-i,-j,-k),axis = (0,1,2))

I can do this using the itertools.product class using for loops, but this is quite slow. If anyone has a better way of doing this, please let me know. I also have CUPY installed with a GTX-970, so if there's a way of using CUDA to do this faster please let me know. If anyone wants some more context please let me know.

Here is my original code for computing the position space two point correlation function. The array x0 is an n by n by n real valued array representing a real scalar field. The function iterate(j,s) runs j iterations. Each iteration consists of generating a random float between -s and s for each lattice site. It then computes the change in the action dS and accepts the change with a probability of min(1,exp^(-dS))

def momentum(k,j,s):
global Gxa
Gx = numpy.zeros((n,n,t))
for i1 in range(0,k):
    iterate(j,s)
    for i2,i3,i4 in itertools.product(range(0,n),range(0,n),range(0,n)):
        x1 = numpy.roll(numpy.roll(numpy.roll(x0, -i2, axis = 0),-i3, axis = 1),-i4,axis = 2)
        x2 = numpy.mean(numpy.multiply(x0,x1))
        Gx[i2,i3,i4] = x2
    Gxa = Gxa + Gx
Gxa = Gxa/k

Upvotes: 4

Views: 285

Answers (1)

Divakar
Divakar

Reputation: 221524

Approach #1

We can extend this idea to our 3D array case here. So, simply concatenate with sliced versions along the three dims and then use np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to efficiently get the final output as the strided-view of the concatenated version, like so -

from skimage.util.shape import view_as_windows

X1 = np.concatenate((X,X[:,:,:-1]),axis=2)
X2 = np.concatenate((X1,X1[:,:-1,:]),axis=1)
X3 = np.concatenate((X2,X2[:-1,:,:]),axis=0)
out = view_as_windows(X3,X.shape)

Approach #2

For really large arrays, we might want to initialize the output array and then re-use X3 from earlier approach to assign with slicing it. This slicing process would be faster than the original-rolling. The implementation would be -

m,n,r = X.shape
Yout = np.empty((m,n,r,m,n,r),dtype=X.dtype)
for i in range(m):
    for j in range(n):
        for k in range(r):
            Yout[:,:,:,i,j,k] = X3[i:i+m,j:j+n,k:k+r]

Upvotes: 2

Related Questions