Chunk an array

General Usage
,
2

You can replace size(arr)[2] by size(arr, 2) (this is more of a style improvement). You can pre-allocate the vector for the chunks.

function chunk_array_2(arr::AbstractMatrix, N::Integer)
    n_cols = size(arr, 2)
    n_chunks = ceil(Int, n_cols / N)
    chunks = Vector{typeof(arr)}(undef, n_chunks)
    for i in 1:n_chunks
        from = N * (i - 1) + 1
        to = min(i * N, n_cols)
        chunks[i] = arr[:, from:to]
    end
    return chunks
end

If it suits your needs, you can use a view when slicing the matrix.

function chunk_array_3(arr::AbstractMatrix, N::Integer)
    n_cols = size(arr, 2)
    n_chunks = ceil(Int, n_cols / N)
    chunks = Vector{AbstractMatrix}(undef, n_chunks)
    for i in 1:n_chunks
        from = N * (i - 1) + 1
        to = min(i * N, n_cols)
        chunks[i] = @view arr[:, from:to]
    end
    return chunks
end

The timings are

arr = rand(Float64, (10, 100))
@btime chunk_array($arr, 30);
  2.271 μs (6 allocations: 8.50 KiB)
@btime chunk_array_2($arr, 30);
  2.168 μs (5 allocations: 8.45 KiB)
@btime chunk_array_3($arr, 30);
  94.155 ns (5 allocations: 336 bytes)
1 Like