ijustlovemath
ijustlovemath

Reputation: 948

How do I split a 2D array into smaller 2D arrays, of variable size?

I've seen this: Slice 2d array into smaller 2d arrays

However, I have images that cannot be evenly split into some N by N blocks. Ideally, I'd like for the following behavior to occur:

>>> import numpy as np
>>> a = np.arange(70*85).reshape([70,85])
>>> arrlist = split_2d(a, [16, 16])
>>> dd = dict()
>>> for arr in arrlist:
...     dd[arr.shape] = dd.get(arr.shape, 0) + 1
>>> print(dd)
    {(16, 16): 16, (10, 16): 4, (16, 5): 4, (10, 5): 1}

Which reflects splitting a (70x85) array into the following:

(16x16) (16x16) (16x16) (16x16) (10x16)
(16x16) (16x16) (16x16) (16x16) (10x16)
(16x16) (16x16) (16x16) (16x16) (10x16)
(16x16) (16x16) (16x16) (16x16) (10x16)
(16x16) (16x16) (16x16) (16x16) (10x16)
(16x5)  (16x5)  (16x5)  (16x5)  (10x5)

The return value can either be a list of 2D arrays, or a 2D array of 2D arrays.

Upvotes: 1

Views: 2571

Answers (1)

piRSquared
piRSquared

Reputation: 294586

This builds the generator you need.

from itertools import product
import numpy as np

n, m = 70, 85
a = np.arange(n * m).reshape(n, m)

def igen(a, n, m):
    i_ = np.arange(a.shape[0]) // n
    j_ = np.arange(a.shape[1]) // m
    for i, j in product(np.unique(i_), np.unique(j_)):
        yield (i, j), a[i_ == i][:, j_ == j]

dict_of_arrays = dict(igen(a, 16, 16))

another alternative
pad with np.nan and reshape + transpose

def cover_multiple(current_length, multiple):
    return ((current_length - 1) // multiple + 1) * multiple

def slicer(a, chunk_i, chunk_j, two_d=True):
    n = cover_multiple(a.shape[0], chunk_i)
    m = cover_multiple(a.shape[1], chunk_j)
    c = np.empty((n, m))
    c.fill(np.nan)
    c[:a.shape[0], :a.shape[1]] = a
    c = c.reshape(n // chunk_i, chunk_i, m // chunk_j, chunk_j)
    c = c.transpose(0, 2, 1, 3)
    if not two_d:
        c = c.reshape(-1, chunk_i, chunk_j)
    return c

demo

a = np.arange(64).reshape(8, 8)
a

[[ 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]]

print(slicer(a, 3, 5))

[[[[  0.   1.   2.   3.   4.]
   [  8.   9.  10.  11.  12.]
   [ 16.  17.  18.  19.  20.]]

  [[  5.   6.   7.  nan  nan]
   [ 13.  14.  15.  nan  nan]
   [ 21.  22.  23.  nan  nan]]]


 [[[ 24.  25.  26.  27.  28.]
   [ 32.  33.  34.  35.  36.]
   [ 40.  41.  42.  43.  44.]]

  [[ 29.  30.  31.  nan  nan]
   [ 37.  38.  39.  nan  nan]
   [ 45.  46.  47.  nan  nan]]]


 [[[ 48.  49.  50.  51.  52.]
   [ 56.  57.  58.  59.  60.]
   [ nan  nan  nan  nan  nan]]

  [[ 53.  54.  55.  nan  nan]
   [ 61.  62.  63.  nan  nan]
   [ nan  nan  nan  nan  nan]]]]

Upvotes: 2

Related Questions