Multidimensional array comprehension

General Usage
,
1

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.

2

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
3

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?

5

Minor side comment:

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

6

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

7

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.

1 Like
8

Alternately,

julia> Tuple.(reshape(CartesianIndices((3,2,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)
9

There are other cute games you can play with this particular use case:

julia> reinterpret(Float64, [cispi(2*Ts*o) for o in 1:2, Ts in 0:0.01:1])'
101×4 adjoint(reinterpret(Float64, ::Matrix{ComplexF64})) with eltype 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.968583  -0.24869    0.876307  -0.481754
 0.982287  -0.187381   0.929776  -0.368125
 0.992115  -0.125333   0.968583  -0.24869
 0.998027  -0.0627905  0.992115  -0.125333
 1.0        0.0        1.0        0.0

(The axis reordering and transposing was necessary here so that the real/imag values would go into alternating columns after the reinterpret)

10

That is a good question! For bigger problems it might be worth instantiating them “on the fly” for each batch. Right now they fit nicely into memory (even on the GPU).

1 Like
11

Good to know - thanks!

12

A comprehension solution (but not efficient):

[(a, d...) for a=1:3, d in Tuple((b, c) for b=1:2, c=1:3)]

This other comprehension is efficient but not as nice looking:

[(a, b, c) for a=1:3, (b,c) in vec(collect(Iterators.product(1:2, 1:3)))]
1 Like
13

Great answers! Thanks to all of you. Hard to choose any particular one as “the” solution :slight_smile:

1 Like
14

Sounds interesting. Do you happen to have a pointer or two to relevant documentation?

15

The array interface is super-simple to implement, eg

struct BasisMatrix <: AbstractMatrix{NTuple{3,Int}}
    nrow::Int
    ncol::Int
end

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

Base.size(bm::BasisMatrix) = (bm.nrow, bm.ncol)

function Base.getindex(bm::BasisMatrix, r, c)
    (; nrow, ncol) = bm
    (1 ≤ r ≤ nrow && 1 ≤ c ≤ ncol) || throw(BoundsError(bm, (r, c)))
    basis_index(r, c)
end

gives you

julia> BasisMatrix(3,4)
3×4 BasisMatrix:
 (1, 1, 1)  (1, 2, 1)  (1, 1, 2)  (1, 2, 2)
 (2, 1, 1)  (2, 2, 1)  (2, 1, 2)  (2, 2, 2)
 (3, 1, 1)  (3, 2, 1)  (3, 1, 2)  (3, 2, 2)
4 Likes