Paritition array into N chunks with Numpy
Try numpy.array_split
.
From the documentation:
>>> x = np.arange(8.0)
>>> np.array_split(x, 3)
[array([ 0., 1., 2.]), array([ 3., 4., 5.]), array([ 6., 7.])]
Identical to numpy.split
, but won't raise an exception if the groups aren't equal length.
If number of chunks > len(array) you get blank arrays nested inside, to address that - if your split array is saved in a
, then you can remove empty arrays by:
[x for x in a if x.size > 0]
Just save that back in a
if you wish.
Just some examples on usage of array_split
, split
, hsplit
and vsplit
:
n [9]: a = np.random.randint(0,10,[4,4])
In [10]: a
Out[10]:
array([[2, 2, 7, 1],
[5, 0, 3, 1],
[2, 9, 8, 8],
[5, 7, 7, 6]])
Some examples on using array_split
:
If you give an array or list as second argument you basically give the indices (before) which to 'cut'
# split rows into 0|1 2|3
In [4]: np.array_split(a, [1,3])
Out[4]:
[array([[2, 2, 7, 1]]),
array([[5, 0, 3, 1],
[2, 9, 8, 8]]),
array([[5, 7, 7, 6]])]
# split columns into 0| 1 2 3
In [5]: np.array_split(a, [1], axis=1)
Out[5]:
[array([[2],
[5],
[2],
[5]]),
array([[2, 7, 1],
[0, 3, 1],
[9, 8, 8],
[7, 7, 6]])]
An integer as second arg. specifies the number of equal chunks:
In [6]: np.array_split(a, 2, axis=1)
Out[6]:
[array([[2, 2],
[5, 0],
[2, 9],
[5, 7]]),
array([[7, 1],
[3, 1],
[8, 8],
[7, 6]])]
split
works the same but raises an exception if an equal split is not possible
In addition to array_split
you can use shortcuts vsplit
and hsplit
.
vsplit
and hsplit
are pretty much self-explanatry:
In [11]: np.vsplit(a, 2)
Out[11]:
[array([[2, 2, 7, 1],
[5, 0, 3, 1]]),
array([[2, 9, 8, 8],
[5, 7, 7, 6]])]
In [12]: np.hsplit(a, 2)
Out[12]:
[array([[2, 2],
[5, 0],
[2, 9],
[5, 7]]),
array([[7, 1],
[3, 1],
[8, 8],
[7, 6]])]
I believe that you're looking for numpy.split
or possibly numpy.array_split
if the number of sections doesn't need to divide the size of the array properly.