Reshaping a Numpy Array into lexicographical list of cubes of shape (n, n, n)

user_7365685 0
Viewed 109

In order to understand what I'm trying to achieve let's imagine an ndarray a with shape (8,8,8) from which I lexicographically take blocks of shape (4,4,4). So while iterating through such blocks the indexes would look as follows:

0: a[0:4, 0:4, 0:4]  
1: a[0:4, 0:4, 4:8]  
2: a[0:4, 4:8, 0:4]  
3: a[0:4, 4:8, 4:8]  
4: a[4:8, 0:4, 0:4]  
5: a[4:8, 0:4, 4:8]  
6: a[4:8, 4:8, 0:4]  
7: a[4:8, 4:8, 4:8]

It is these blocks of data which I'm trying to access. Obviously, this can be described by using an expression which converts the current iteration to the corresponding indexes. An example of that is given below.

a = np.ones((8,8,8))
f = 4

length = round(a.shape[0] * a.shape[1] * a.shape[2] / f**3)

x = a.shape[0] / f
y = a.shape[1] / f
z = a.shape[2] / f

for i in range(length):
    print(f"{i}: {round((int(i/(z*y))%x)*f)}:{round(f+(int(i/(z*y))%x)*f)}, {round((int(i/z)%y)*f)}:{round(f+(int(i/z)%y)*f)}, {round((i%z)*f)}:{round(f+(i%z)*f)}")

My apologies for having to do that to your eyes but it generates the following output:

0: 0:4, 0:4, 0:4
1: 0:4, 0:4, 4:8
2: 0:4, 4:8, 0:4
3: 0:4, 4:8, 4:8
4: 4:8, 0:4, 0:4
5: 4:8, 0:4, 4:8
6: 4:8, 4:8, 0:4
7: 4:8, 4:8, 4:8

So this does actually generate the right indexes, but it only allows you to access multiple blocks at once if they have the same index in the 0th and 1st axis, so no wrapping around. Ideally I would reshape this whole ndarray into an ndarray b with shape (4, 4, 32) and be ordered in such a way that b[:, :, :4] would return a[0:4, 0:4, 0:4], b[:, :, 4:12] returns an ndarray of shape (4, 4, 8) which contain a[0:4, 0:4, 4:8] and a[0:4, 4:8, 0:4] etc. I want this to be as fast as possible, so ideally, I keep the memory layout and just change the view on the array.
Lastly, if it helps to think about this conceptually, this is basically a variant of the ndarray.flatten() method but using blocks of shape (4, 4, 4) as "atomic size" if you will.

Hope this makes it clear enough!

scipy numpy python
2 Answers

It is a bit unclear what you want as output. Are you looking for this:

from skimage.util.shape import view_as_windows
b = view_as_windows(a,(f,f,f),f).reshape(-1,f,f,f).transpose(1,2,3,0).reshape(f,f,-1)

suggested by @Paul with similar result (I prefer this answer in fact):

N = 8
b = a.reshape(2,N//2,2,N//2,N).transpose(1,3,0,2,4).reshape(N//2,N//2,N*4)

output:

print(np.array_equal(b[:, :, 4:8],a[0:4, 0:4, 4:8]))
#True
print(np.array_equal(b[:, :, 8:12],a[0:4, 4:8, 0:4]))
#True
print(np.array_equal(b[:, :, 12:16],a[0:4, 4:8, 4:8]))
#True
user_49759810 
def flatten_by(arr, atomic_size):
    a, b, c = arr.shape
    x, y, z = atomic_size
    
    r = arr.reshape([a//x, x, b//y, y, c//z, z])
    r = r.transpose([0, 2, 4, 1, 3, 5])
    r = r.reshape([-1, x, y, z])
    
    return r

flatten_by(arr, [4,4,4]).shape

>>>   (8, 4, 4, 4)

EDIT:

the function applies C-style flattening to the array, as shown below

NOTE: this method and @Ehsan's method both produce "copies" NOT "views", im looking into it and would update the answer when if i find a solution

flattened = flatten_by(arr, [4,4,4])

required = np.array([
    arr[0:4, 0:4, 0:4],
    arr[0:4, 0:4, 4:8],  
    arr[0:4, 4:8, 0:4],  
    arr[0:4, 4:8, 4:8],  
    arr[4:8, 0:4, 0:4],  
    arr[4:8, 0:4, 4:8],  
    arr[4:8, 4:8, 0:4],  
    arr[4:8, 4:8, 4:8],  
])

np.array_equal(required, flattened)

>>> True
user_132505890
Related