Craig Odell
Craig Odell

Reputation: 143

Python/Numpy rejoin matrix after splitting into sub-matrices

I have a numpy matrix that looks something like:

[[ 0  1  2  3  4  5  6  7]
 [ 8  9 10 11 12 13 14 15]
 [16 17 18 19 20 21 22 23]
 [24 25 26 27 28 29 30 31]
 [32 33 34 35 36 37 38 39]
 [40 41 42 43 44 45 46 47]
 [48 49 50 51 52 53 54 55]
 [56 57 58 59 60 61 62 63]]

where I break the matrix into "chunks" or sub-matrices like this:

[[[ 0  1]
  [ 8  9]]
 [[ 2  3]
  [10 11]]
 [[ 4  5]
  [12 13]]
 [[ 6  7]
  [14 15]]
 [[16 17]
  [24 25]]
 [[18 19]
  [26 27]]
 [[20 21]
  [28 29]]
...
 [[50 51]
  [58 59]]
 [[52 53]
  [60 61]]
 [[54 55]
  [62 63]]]

Where each 2x2 quandrant is taken left-to-right top-to-bottom. Generated with:

def view_as_blocks(arr, BSZ):
    # arr is input array, BSZ is block-size
    m, n = arr.shape
    M, N = BSZ
    return arr.reshape(m // M, M, n // N, N).swapaxes(1, 2).reshape(-1, M, N)

What I'm trying to do is re-construct the original matrix in a "numpy" sort of way using something like reshape, stack/vstack/hstack/etc. I am able to reconstruct the matrix using something like:

#Full code at the bottom

block_width = 2
block_size = 4
width = 8
blocks_per_row = width // block_width
for index in range(64):
    block_row = index // (block_size * blocks_per_row)
    normal_index = (index - (block_row * block_size * blocks_per_row))
    row = normal_index // width
    col = normal_index - row * width
    block_index = block_row * blocks_per_row + col // block_width
    block_col = col % block_width
    print(block_index, row, block_col, c[block_index][row][block_col], '=', index, c[block_index][row][block_col] == index)

It seems like there should be a way to reverse the splitting into "chunks" back to the original matrix, without having to walk each index, but my brain can't seem to figure it out. I can do the above and re-construct the original matrix, but I'm trying to stay in the "numpy world" for performance for when you have a much larger matrix. Thanks in advance to anyone.

Here's the full code:

import numpy as np

def view_as_blocks(arr, BSZ):
    # arr is input array, BSZ is block-size
    m, n = arr.shape
    M, N = BSZ
    return arr.reshape(m // M, M, n // N, N).swapaxes(1, 2).reshape(-1, M, N)


c = np.arange(64).reshape(8, 8)
print(c)
c = view_as_blocks(c, (2, 2))
print(c)

block_width = 2
block_size = 4
width = 8
blocks_per_row = width // block_width
for index in range(64):
    block_row = index // (block_size * blocks_per_row)
    normal_index = (index - (block_row * block_size * blocks_per_row))
    row = normal_index // width
    col = normal_index - row * width
    block_index = block_row * blocks_per_row + col // block_width
    block_col = col % block_width
    print(block_index, row, block_col, c[block_index][row][block_col], '=', index, c[block_index][row][block_col] == index)

Upvotes: 2

Views: 203

Answers (1)

Gabriel A.
Gabriel A.

Reputation: 641

Nice code, what are you doing with that chunks?

Anyways, for that you must know the original shape, because there are many ways to reconstruct a flattened matrix. I just reversed the order of operations that you did:

def reconstruct(arr, original_shape):
    _, M, N = arr.shape
    m, n = original_shape
    return arr.reshape(m // M, n // N, M, N).swapaxes(1, 2).reshape(m, n)

Upvotes: 3

Related Questions