What is the correct way to iterate over "Kroneker Product" like terms?

I would like to write functions that iterate over objects that have a “Kroneker Product” like structure.

Consider for example:

using StaticArrays

struct MyKronekerType{D,N,T}
    data::SVector{T,N}
    function MyKronekerType(D::Int, data::SVector{T,N}) where {T,N}
        @assert D > 0
        new{D,N,T}(data)
    end
end

function Base.getindex(m::MyKronekerType{1}, i)
    return m.data[i]
end

function Base.getindex(m::MyKronekerType{2}, i, j)
    return m.data[i], m.data[j]
end

data = @SVector [1,2,3]
m1 = MyKronekerType(1, data)
m2 = MyKronekerType(2, data)

println("m1[1]   = ", m1[1])   # 1
println("m2[1,2] = ", m2[1,2]) # (1,2)

So MyKronekerType{D} represents a “D-dimensional” version of its data field. I wish to now iterate over all the entries of m::MyKronekerType{D,N,T}. For example, I wish to write code like:

m2 = MyKronekerType(2, data)
for d in m2
    # d is m2[1,1], m2[2,1], ... m2[3,3]
end

What is the best way to do this? I figured I can explicitly convert a linear index into a “D-dimensional” index, but I was hoping there is a more elegant way to do this. I thought maybe CartesianIndices might be what I require here but I don’t quite know how to use it for this purpose.

You can achieve this by defining Base.iterate for your type.

Base.iterate(d::MyKronekerType{D}) where {D} = iterate(Iterators.product((d.data for i=1:D)...))
Base.iterate(d::MyKronekerType{D},state) where {D} = iterate(Iterators.product((d.data for i=1:D)...),state)

With these defined you obtain

julia> for d in m2
   println(d)
   end
(1, 1)
(2, 1)
(3, 1)
(1, 2)
(2, 2)
(3, 2)
(1, 3)
(2, 3)
(3, 3)
3 Likes

This is perfect! Thanks!