Iterating over the columns of a matrix

I want to iterate a function whose argument is a vector. The vectors to iterate on are the columns of a matrix. I was “slicing” the matrix into columns following Converting a matrix into an array of arrays and then iterate over the new array whose entries are the vectors I want.

I think it is not optimal and that is my first question: If all I want is to iterate over the columns of a matrix, then shouldn’t I define an iterator (instead of copying the whole matrix)?
So, I came up with the following code which does the trick, but I am not sure whether this is the right way to define such iterator (e.g. are the last two lines strictly necessary? defining the length gives me problems with the length function…)

struct AbstractMatrixColumn
    M::AbstractMatrix
end

function Base.iterate(M::AbstractMatrixColumn)
    state = 1
    return Base.iterate(M,state)
end

function Base.iterate(M::AbstractMatrixColumn, state)
    if state > size(M.M,2)
        next = nothing
    else
        next = M.M[:,state], state+1
    end
    return  next
end

Base.length(M::AbstractMatrixColumn) = size(M.M,2)
Base.eltype(M::AbstractMatrixColumn) = typeof(M.M[:,1])

Coming back to Converting a matricx into an array of arrays, now we can use this iterator to collect the columns of a matrix M with collect(AbstractMatrixColumns(M)), but it is much slower. Is it supposed to be or it is a problem with my implementation?

function slicematrix_column(M::AbstractMatrix{T}) where T
    return [M[:,i] for i in 1:size(M,2)]
end

n = 1000;
N = rand(n,n);
@btime collect(AbstractMatrixColumn($N));
@btime slicematrix_column($N);

> julia
  2.083 ms (6945 allocations: 7.87 MiB)
  1.329 ms (1003 allocations: 7.76 MiB)

This is what eachcol, eachrow and eachslice are for:

julia> A = rand(0:9, 4,5)
4×5 Array{Int64,2}:
 9  2  8  7  6
 2  4  9  0  6
 4  4  4  5  3
 3  0  7  3  1

julia> for col in eachcol(A)
       println(col)
       end
[9, 2, 4, 3]
[2, 4, 4, 0]
[8, 9, 4, 7]
[7, 0, 5, 3]
[6, 6, 3, 1]

These iterators return views, in case that’s relevant, while your code makes copies.

1 Like

BTW, a performance issue in your code is that you have an abstract type signature in the field of AbstractMatrixColumn. If you want to experiment furter with this, you should make that field concrete or parametric.

Thank you so much @DNF. Yes, I want to experiment further :smile: