Multidimensional array comprehension

fps
Thu, 20 Nov 2025 09:25:18 +0000
General Usage

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

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.

14 posts - 6 participants