Shorthand for ((i-1)*k+1):(i*k)

General Usage
,
1

I often run into the problem of partitioning data to into equal “chunks”, eg a vector of N*K elements into N chunks of K, then doing something to each chunk. As an MWE, consider

z = zeros(N*K)
for i in 1:N
    z[((i-1)*k+1):(i*k)] = some_calculation(i)
end

Typing this up every time is prone to various coding errors.

In some contexts, I can use reshape, as in

_z = reshape(z, K, N)
for i in 1:N
    _z[:, i] = some_calculation(i)
end
# subsequent calculations use z as a result

but this only works if z is accepted by reshape as a view, and it relies on the structure and representation of matrices.

I also considered making a calculator for this, as in

struct Chunks
    K::Int
end

Base.getindex(c::Chunks, i::Int) = (e = c.K * i; (e - c.K + 1):e)

that I can reuse:

julia> for i in 1:3
       @show Chunks(7)[i]
       end
(Chunks(7))[i] = 1:7
(Chunks(7))[i] = 8:14
(Chunks(7))[i] = 15:21

but it is too tiny to make into a package (does any package have this already?)

How do others deal with this kind of problem?

2 Likes
2

I often do

inds = 1:K
for i = ...
   z[inds] = ...
   inds = inds .+ K
end

this is allocation free since adding a constant to a range keeps inds a range. It’s also easy to handle the case when inds doesn’t cover the entire chunk, that is, inds is shorter than the step length between iterations

4 Likes
3

Isn’t that what ChunkSplitters.jl is for?

1 Like
4

Exactly! I just didn’t know about it.

(The Julia community is spoiling me. Every time I run into a problem, it turns out that someone made a high-quality package that solves it.)

2 Likes
5

Or even Iterators.partition: that’ll efficiently spit back out ranges. I often find its API to feel backwards — you specify the length of each inner chunk rather than the total number of chunks, but in this case it looks like it’s exactly what you want!

Just grab the axis, get your list of k-length partitions, and iterate over that (or enumerate it)

2 Likes
6

You also have partition

julia> K = 3; N = 5;

julia> Iterators.partition(1:K*N, K) |> collect
5-element Vector{UnitRange{Int64}}:
 1:3
 4:6
 7:9
 10:12
 13:15

Looks like Matt beat me to it :slight_smile:

3 Likes