Multidimensional array comprehension

I wonder it is possible to get the following matrix with a single array comprehension that’s also efficient to evaluate:

(1, 1, 1)  (1, 2, 1) (1, 1, 2) (1, 2, 2) (1, 1, 3) (1, 2, 3)....
(2, 1, 1)  (2, 2, 1) (2, 1, 2) (2, 2, 2) (2, 1, 3) (2, 2, 3)
(3, 1, 1)  (3, 2, 1) (3, 1, 2) (3, 2, 2) (3, 1, 3) (3, 2, 3)
.
.
.

Note the alternating second entry and the third entry increment every second column. I sometimes need this kind of comprehension to quickly build up a fourier basis. Something like this works:

julia> [(a, mod(b-1, 2)+1, div(b+1, 2)) for a = 1:3, b = 1:6]
3×6 Matrix{Tuple{Int64, Int64, Int64}}:
 (1, 1, 1)  (1, 2, 1)  (1, 1, 2)  (1, 2, 2)  (1, 1, 3)  (1, 2, 3)
 (2, 1, 1)  (2, 2, 1)  (2, 1, 2)  (2, 2, 2)  (2, 1, 3)  (2, 2, 3)
 (3, 1, 1)  (3, 2, 1)  (3, 1, 2)  (3, 2, 2)  (3, 1, 3)  (3, 2, 3)

Or more explicitly for this use-case:

function fourier_basis(Ts, order)
    [sin(2f0 * pi .* (div(o+1, 2) .* t) .+ (mod(o, 2) * pi/2f0)) for t in Ts, o in 1:(2*order)]
end

julia> fourier_basis(0:0.01:1, 2)
101×4 Matrix{Float64}:
 1.0        0.0          1.0         0.0
 0.998027   0.0627905    0.992115    0.125333
 0.992115   0.125333     0.968583    0.24869
 0.982287   0.187381     0.929776    0.368125
 0.968583   0.24869      0.876307    0.481754
 0.951057   0.309017     0.809017    0.587785
 0.929776   0.368125     0.728969    0.684547
 0.904827   0.425779     0.637424    0.770513
 0.876307   0.481754     0.535827    0.844328
 0.844328   0.535827     0.425779    0.904827
 0.809017   0.587785     0.309017    0.951057
 0.770513   0.637424     0.187381    0.982287
 0.728969   0.684547     0.0627905   0.998027
 0.684547   0.728969    -0.0627906   0.998027
 0.637424   0.770513    -0.187381    0.982287
 ⋮
 0.637424  -0.770513    -0.187381   -0.982287
 0.684547  -0.728969    -0.0627902  -0.998027
 0.728969  -0.684547     0.0627908  -0.998027
 0.770513  -0.637424     0.187382   -0.982287
 0.809017  -0.587785     0.309017   -0.951056
 0.844328  -0.535827     0.42578    -0.904827
 0.876307  -0.481754     0.535827   -0.844328
 0.904827  -0.425779     0.637424   -0.770513
 0.929777  -0.368124     0.728969   -0.684547
 0.951057  -0.309017     0.809017   -0.587785
 0.968583  -0.24869      0.876307   -0.481753
 0.982287  -0.187381     0.929777   -0.368124
 0.992115  -0.125333     0.968583   -0.24869
 0.998027  -0.0627903    0.992115   -0.125333
 1.0        1.74846e-7   1.0         3.49691e-7

but it’s a little unwieldy. Of course I could go the complex number route, but then I would need to destructure the complex entries in another step which might be unnecessary :slight_smile:

I tried different combinations of the two iteration primitives in array comprehensions for a in A, b in B and for a in A for b in B and I couldn’t find one that did the above in a single step.

OK, sometimes you just need to talk about it to then find a solution. In this case I can just do:

julia> reshape([(a, b, c) for a = 1:3, b = 1:2, c = 1:3], 3, 6)
3×6 Matrix{Tuple{Int64, Int64, Int64}}:
 (1, 1, 1)  (1, 2, 1)  (1, 1, 2)  (1, 2, 2)  (1, 1, 3)  (1, 2, 3)
 (2, 1, 1)  (2, 2, 1)  (2, 1, 2)  (2, 2, 2)  (2, 1, 3)  (2, 2, 3)
 (3, 1, 1)  (3, 2, 1)  (3, 1, 2)  (3, 2, 2)  (3, 1, 3)  (3, 2, 3)
1 Like

I would just define something like

 basis_index(a, b) = ((d, r) = divrem(b + 1, 2); (a, r + 1, d))

and then use comprehension; but are you sure that you actually need to instantiate this matrix?

(post deleted by author)

Minor side comment:

This is the same as fldmod1(b, 2) although in reverse order (assuming b isn’t negative).

Yes, it is used as an input to a neural network. So sadly it has to “exist” :wink:

Is the calculation in Julia though? If yes, you can code up an <:AbstractArray that is never instantiated, just calculated on demand when Base.getindex is called.