# 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!