Rolling/Increasing dimensionality of a NumPy array

Question

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

Divakar · Accepted Answer · 2018-07-30 10:34:56Z

2

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]

edited Jul 30, 2018 at 10:34

answered Jul 30, 2018 at 7:38

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

6 Comments

chuckstables Over a year ago

Awesome! Great speedups I'm getting. Bad news; I'm now limited to a max lattice size of about 33 in the three dimensions (i'm simulating a lattice quantum field theory), as I get a memory error if I go above that as 33^6 is a very large number. Any suggested ways to get around that?

chuckstables Over a year ago

It doesn't really work. I think the actual memory sink would be when I'm doing np.multiply to element wise multiply the new six dimensional matrix with the original matrix, as I'm trying to compute Euclidean (imaginary time) green functions for the interacting theory on a lattice. Regardless using your method I can probe lattice sizes that would be unfeasible for me previously, so thank you!

Divakar Over a year ago

@chuckstables By - It doesn't really work. , do yo mean approach2 doesn't work? If so, do you mean its slow or is it giving wrong output? I have checked for both performance and correctness at my end.

chuckstables Over a year ago

I mean approach 2. It's giving the correct output, and seems to be as fast as the first method, but I'm still hitting memory limits. I'll run some more tests when I have some time and I'll let you know what I find.

Divakar Over a year ago

@chuckstables Just curious, are you hitting the memory limits with your roll method too? If so, I think we don't have any other way to go around.

|

Collectives™ on Stack Overflow

Rolling/Increasing dimensionality of a NumPy array

1 Answer 1

6 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

6 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related