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.
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:
Not quite an answer, but a long comment with nice formatting of code to the other (correct) answers. If you try the following, you will see that what you are getting are views of the original array, not copies, and that was not the case for the accepted answer in the question you link. Be aware of the possible side effects!
>>> x = np.arange(9.0)
>>> a,b,c = np.split(x, 3)
>>> a
array([ 0., 1., 2.])
>>> a[1] = 8
>>> a
array([ 0., 8., 2.])
>>> x
array([ 0., 8., 2., 3., 4., 5., 6., 7., 8.])
>>> def chunks(l, n):
... """ Yield successive n-sized chunks from l.
... """
... for i in xrange(0, len(l), n):
... yield l[i:i+n]
...
>>> l = range(9)
>>> a,b,c = chunks(l, 3)
>>> a
[0, 1, 2]
>>> a[1] = 8
>>> a
[0, 8, 2]
>>> l
[0, 1, 2, 3, 4, 5, 6, 7, 8]
This can be achieved using as_strided of numpy. I have put a spin to answer by assuming that if chunk size is not a factor of total number of rows, then rest of the rows in the last batch will be filled with zeros.
from numpy.lib.stride_tricks import as_strided
def batch_data(test, chunk_count):
m,n = test.shape
S = test.itemsize
if not chunk_count:
chunk_count = 1
batch_size = m//chunk_count
# Batches which can be covered fully
test_batches = as_strided(test, shape=(chunk_count, batch_size, n), strides=(batch_size*n*S,n*S,S)).copy()
covered = chunk_count*batch_size
if covered < m:
rest = test[covered:,:]
rm, rn = rest.shape
mismatch = batch_size - rm
last_batch = np.vstack((rest,np.zeros((mismatch,rn)))).reshape(1,-1,n)
return np.vstack((test_batches,last_batch))
return test_batches
np.array_split will try to split "evenly", for example, if x.shape is 10, sections is 3, you will get splits with shape [3, 3, 2, 2] instead of [3, 3, 3, 1], a workaround is using spaced indices like snippet below
import math
import numpy as np
def split_evenly(x, chunk_size, axis=0):
return np.array_split(x, math.ceil(x.shape[axis] / chunk_size), axis=axis)
def split_reminder(x, chunk_size, axis=0):
indices = np.arange(chunk_size, x.shape[axis], chunk_size)
return np.array_split(x, indices, axis)
x = np.arange(10)
chunk_size = 3
print([i.shape[0] for i in split_evenly(x, chunk_size, 0)])
print([i.shape[0] for i in split_reminder(x, chunk_size, 0)])
# [3, 3, 2, 2]
# [3, 3, 3, 1]