Numpy, how to break a list into multiple chunks

Question

I am trying to break a numpy array into chunks with a fixed size and pad the last one with 0. For example: [1,2,3,4,5,6,7] into chunks of 3 returns [[1,2,3],[4,5,6],[7,0,0]].

The function I wrote is:

def makechunk(lst, chunk):
    result = []
    for i in np.arange(0, len(lst), chunk):
        temp = lst[i:i + chunk]
        if len(temp) < chunk:
            temp = np.pad(temp, (0, chunk - len(temp)), 'constant')
        result.append(temp)
    return result

It works but when dealing with large size array it is very slow. What is a more numpy-ish and vectorized way of doing it?

Answer 1

Using the function resize() should do what you need :

l = np.array([1,2,3,4,5,6,7])
l.resize((3,3), refcheck=False)

(Edit: mea culpa, monday problem with reasignation)

@J: Resize boost the speed by about 5 times for np.arange(0,44100) into chunks of 512.

import math
def makechunk4(lst, chunk):
    l = lst.copy()
    l.resize((math.ceil(l.shape[0]/chunk),chunk), refcheck=False)
    return l

Answer 2

Time comparison of @Cedric Poulet's (all kudos to him, see his answer) solution (with added array splitting so it returns the result as desired) with another numpy approach I thought about at first (create array of zeros and insert data in-place):

import time

import numpy as np

def time_measure(func):
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        stop = time.time()
        print(f"Elapsed time: {stop-start}")
        return result

    return wrapper


@time_measure
def pad_and_chunk(array, chunk_size: int):
    padded_array = np.zeros(len(array) + (chunk_size - len(array) % chunk_size))
    padded_array[: len(array)] = array
    return np.split(padded_array, len(padded_array) / chunk_size)


@time_measure
def resize(array, chunk_size: int):
    array.resize(len(array) + (chunk_size - len(array) % chunk_size), refcheck=False)
    return np.split(array, len(array) / chunk_size)

@time_measure
def makechunk4(l, chunk):
    l.resize((math.ceil(l.shape[0] / chunk), chunk), refcheck=False)
    return l.reshape(chunk, -1)


if __name__ == "__main__":
    array = np.random.rand(1_000_000)

    ret = pad_and_chunk(array, 3)
    ret = resize(array, 3)
    ret = makechunk4(array, 3)

EDIT-EDIT

Gathering all possible answers it is indeed the case that np.split is horribly slow when compared to reshape.

Elapsed time: 0.3276541233062744
Elapsed time: 0.3169224262237549
Elapsed time: 1.8835067749023438e-05

Way of padding data is not essential, it's the split taking up most of the time.

Answer 3

in the itertools recipes there is a recipe for grouper:

from itertools import zip_longest
import numpy as np

array = np.array([1,2,3,4,5,6,7])

def grouper(iterable, n, fillvalue=None):
    args = [iter(iterable)] * n
    return zip_longest(*args, fillvalue=fillvalue)

res = list(grouper(array, 3, fillvalue=0))
# [(1, 2, 3), (4, 5, 6), (7, 0, 0)]

if you need the sublists to be lists and not tuples:

def grouper(iterable, n, fillvalue=None):
    args = [iter(iterable)] * n
    return (list(item) for item in zip_longest(*args, fillvalue=fillvalue))

Answer 4

A solution using numpy

I assume a chunk size of 3 and created a random array input of length 10 in x.

# Chunk size
chunk = 3
# Create array
x = np.arange(10)

First make sure to pad the array with zeros. Next you can use reshape to create an array of arrays.

# Pad array
x = np.pad(x, (0, chunk - (x.shape[0]%chunk)), 'constant')
# Divide into chunks
x = x.reshape(-1, chunk)

Optionally you can retrieve the numpy array as a list

x = x.tolist()