CODEWITHSUNDEEP
pythonnumpy
JP1
JP1

Reputation: 751

split ndarray into chunks - whilst maintaining order

I am trying to implement a function, which can split a 3 dimensional numpy array in to 8 pieces, whilst keeping the order intact. Essentially I need the splits to be:

G[:21, :18,:25]
G[21:, :18,:25]
G[21:, 18:,:25]
G[:21, 18:,:25]
G[:21, :18,25:]
G[21:, :18,25:]
G[21:, 18:,25:]
G[:21, 18:,25:]

Where the original size of this particular matrix would have been 42, 36, 50. How is it possible to generalise these 8 "slices" so I do not have to hardcode all of them? essentially move the :in every possible position.

Thanks!

Upvotes: 0

Views: 545

Answers (1)

hpaulj
hpaulj

Reputation: 231355

You could apply the 1d slice to successive (lists of) dimensions.

With a smaller 3d array

In [147]: X=np.arange(4**3).reshape(4,4,4)

A compound list comprehension produces a nested list. Here I'm using the simplest double split

In [148]: S=[np.split(z,2,0) for y in np.split(X,2,2) for z in np.split(y,2,1)]

In this case, all sublists have the same size, so I can convert it to an array for convenient viewing:

In [149]: SA=np.array(S)

In [150]: SA.shape
Out[150]: (4, 2, 2, 2, 2)

There are your 8 subarrays, but grouped (4,2). 

In [153]: SAA = SA.reshape(8,2,2,2)

In [154]: SAA[0]
Out[154]: 
array([[[ 0,  1],
        [ 4,  5]],

       [[16, 17],
        [20, 21]]])

In [155]: SAA[1]
Out[155]: 
array([[[32, 33],
        [36, 37]],

       [[48, 49],
        [52, 53]]])

Is the order right? I can change it by changing the axis in the 3 split operations.

Another approach is write your indexing expressions as tuples

In [156]: x,y,z = 2,2,2   # define the split points

In [157]: ind = [(slice(None,x), slice(None,y), slice(None,z)),     
                 (slice(x,None), slice(None,y), slice(None,z)),]  
                 # and so on

In [158]: S1=[X[i] for i in ind]

In [159]: S1[0]
Out[159]: 
array([[[ 0,  1],
        [ 4,  5]],

       [[16, 17],
        [20, 21]]])

In [160]: S1[1]
Out[160]: 
array([[[32, 33],
        [36, 37]],

       [[48, 49],
        [52, 53]]])

Looks like the same order I got before.

That ind list of tuples can be produced with some sort of iteration and/or list comprehension. Maybe even using itertools.product or np.mgrid to generate the permutations.

An itertools.product version could look something like

In [220]: def foo(i):
    return [(slice(None,x) if j else slice(x,None)) 
            for j,x in zip(i,[2,2,2])]

In [221]: SAA = np.array([X[foo(i)] for i in 
    itertools.product(range(2),range(2),range(2))])

In [222]: SAA[-1]
Out[222]: 
array([[[ 0,  1],
        [ 4,  5]],

       [[16, 17],
        [20, 21]]])

product iterates the last value fastest, so the list is reversed (compared to your target).

To generate a particular order it may be easier to list the tuples explicitly, e.g.:

In [227]: [X[foo(i)] for i in [(1,1,1),(0,1,1),(0,0,1)]]
Out[227]: 
[array([[[ 0,  1],
         [ 4,  5]],

        [[16, 17],
         [20, 21]]]), array([[[32, 33],
         [36, 37]],

        [[48, 49],
         [52, 53]]]), array([[[40, 41],
         [44, 45]],

        [[56, 57],
         [60, 61]]])]

This highlights the fact that there are 2 distinct issues - generating the iteration pattern, and splitting the array based on this pattern.

Upvotes: 1

Related Questions