I want to split a numpy array, along the first axis, into subarrays of unequal size. I have checked numpy.split but it seems that I can only pass indices instead of size (number of rows for each subarray).
for example:
arr = np.array([[1,2], [3,4], [5,6], [7,8], [9,10]])
should yield:
arr.split([2,1,2]) = [array([[1,2], [3,4]]), array([5,6]), array([[7,8], [9,10]])]
Let the cutting intervals be -
cut_intvs = [2,1,2]
Then, use np.cumsum to detect positions of cuts -
cut_idx = np.cumsum(cut_intvs)
Finally, use those indices with np.split to cut the input array along the first axis and ignore the last cut to get the desired output, like so -
np.split(arr,np.cumsum(cut_intvs))[:-1]
Sample run for explanation -
In [55]: arr # Input array
Out[55]:
array([[ 1, 2],
[ 3, 4],
[ 5, 6],
[ 7, 8],
[ 9, 10]])
In [56]: cut_intvs = [2,1,2] # Cutting intervals
In [57]: np.cumsum(cut_intvs) # Indices at which cuts are to happen along axis=0
Out[57]: array([2, 3, 5])
In [58]: np.split(arr,np.cumsum(cut_intvs))[:-1] # Finally cut it, ignore last cut
Out[58]:
[array([[1, 2],
[3, 4]]), array([[5, 6]]), array([[ 7, 8],
[ 9, 10]])]
