numpy : assembling multiple slices into new array

Question

I have a 2-dimentional array, of which I need to extract sections (slices) into a new array:

original= numpy.ndarray( shape=(4,4) )
slices= numpy.ndarray( shape=(0,2) )
for x in range(3):
    slice= original[x:x+2,x:x+2] 
    slices=numpy.append(slices, slice,axis=0)

Is there a more efficient way to do this (getting rid of the python for cycle)?

----EDIT----

To clarify, I'm asking how to copy arbitrarily (but similarly) shaped 2D slices from arbitrary 2D indexes of an 2D array into another, vertically stacked - not particularly along the diagonal, or 2x2 sized.

Answer 1

Here is your solution:

v = np.arange(0,original.shape[0],0.5).astype(int)
result = np.c_[ original[v[1:-1],v[:-2]] , original[v[1:-1],v[2:]] ]

works for any size of the square input matrix ("original" as you called it).

The idea is to create a "helper array" v, which is simply [0,0,1,1,2,2,3,3,...], and then use the observation that the indices you need are always simple slices of v.

Enjoy!

Answer 2

There is a nifty trick with stride_tricks, you can find rolling window functions with different generality on SO and other (there currently is none in numpy itself), here is a version tailored to what you got:

def rolling_window(arr, window):
    """Very basic multi dimensional rolling window. window should be the shape of
    of the desired subarrays. Window is either a scalar or a tuple of same size
    as `arr.shape`.
    """
    shape = np.array(arr.shape*2)
    strides = np.array(arr.strides*2)
    window = np.asarray(window)
    shape[arr.ndim:] = window # new dimensions size
    shape[:arr.ndim] -= window - 1
    if np.any(shape < 1):
        raise ValueError('window size is too large')
    return np.lib.stride_tricks.as_strided(arr, shape=shape, strides=strides)

# Now:
view = rolling_window(arr, 2)
view[0,0] # first slice in your loop

Note that view holds the same data as the original array! Which can result in unexpected results. But you seem to want only the diagonal, you could do that with stride tricks as well to make sure you do not copy data if you want (next versions will create a view with diagonal, old ones always a copy):

diagonal = np.diagonal(view, axis1=0, axis2=1)
# unfortunatly now the first slice is diagonal[...,0], so just roll it to the start:
diagonal = np.rollaxis(diagonal, -1)

Now diagonal is the array you created in your for loop (on newer versions add a .copy() if you do not want a view).

Edit: Since the slices array is 2D and not 3D because you append, a reshape was missing here:

slices = diagonal.reshape(-1,2)

This might not be faster if you have such small arrays, but its constant (expect for the copying of the data in diagonal call) with the array size.