I often want to iterate either row-wise or column-wise though a matrix in Julia, so I've created a pair of functions to help:
function cols(x::Matrix)
function _it()
for ii in 1:size(x,2)
produce(x[:,ii])
end
end
Task(_it)
end
function rows(x::Matrix)
function _it()
for ii in 1:size(x,1)
produce(x[ii,:])
end
end
Task(_it)
end
Example of use:
AA = [1 2 3; 1 2 3]
println("Columns are:")
for ii in cols(AA)
println(ii)
end
println("\nRows are:")
for ii in rows(AA)
println(ii)
end
Outputs:
Columns are:
[1,1]
[2,2]
[3,3]
Rows are:
[1 2 3]
[1 2 3]
Possible issues:
- It seems like this could be one function, that just takes the dimension as a parameter.
- This might confound the JIT, which would mean losing out out a bunch of speed. (The Julia JIT can normally vectorize
for
loops). - I feel like there should be a inbuild function for this, but the closest I know of is mapslice.
1 Answer 1
I think the most performant solution here would be to create an iterator:
julia> immutable EachRow{T<:AbstractMatrix}
A::T
end
Base.start(::EachRow) = 1
Base.next(itr::EachRow, s) = (itr.A[s,:], s+1)
Base.done(itr::EachRow, s) = s > size(itr.A,1)
done (generic function with 48 methods)
julia> AA = [1 2 3; 1 2 3]
for ii in EachRow(AA)
println(ii)
end
[1 2 3]
[1 2 3]
EachCol
would be coded similarly.
This is advantageous over your task-based iterator design because Julia is not able to infer the return types from anonymous tasks. This means that within your loop, the iteration variable ii
is typed as Any
, so the compiler isn't able to emit efficient type-specific code. The custom iterator is type-stable, so that means that Julia knows that the iteration variable is always going to be an array, allowing it to compile specific and optimized instructions.
To answer a few of your questions: We could similarly make an EachDim
iterator which also stored the dimension parameter, but this is a little trickier to write in a type-stable manner. Since indexing into columns returns a 1-dimensional column vector and indexing rows returns a 2-dimensional row vector, simply using an integer argument 1
or 2
to specify which dimension would result in a type-instability. We'd need to ensure that the dimension information is stored in the type-domain. One way to do this would be to use a parametric type... but doing this nicely requires the call-overloading feature from the development 0.4 version.
Also note that Julia's JIT doesn't "vectorize" for loops in the way you're using the word. It can be thought of more as a just-barely-ahead-of-time compiler that compiles specialized machine code for each function based upon the types of the arguments. This is why type-stability is so important; when Julia can infer the types precisely the compiled code is typically very similar to what you'd get from equivalent C code.