I have a package for doing all such reshapes and permutations. If A[i,j]
is the original matrix, then v[(i,j)]
is its notation for vec(A)
, combining the two indices without changing the order. And the one you ask for is B
, which involves two reshapes and a permutation:
julia> n = 16; d = 2;
julia> A = reshape(1:n, isqrt(n), isqrt(n)) .+ 0
4×4 Matrix{Int64}:
1 5 9 13
2 6 10 14
3 7 11 15
4 8 12 16
julia> using TensorCast
julia> @cast B[i,(k,j)] := A[(i,j),k] i in 1:d # := makes a new array, lazy when possible
2×8 reshape(transmute(::Array{Int64, 3}, (1, 3, 2)), 2, 8) with eltype Int64:
1 5 9 13 3 7 11 15
2 6 10 14 4 8 12 16
julia> @cast R[i,(j,k)] |= A[(i,j),k] i in 1:d # this is a normal reshape, |= collects
2×8 Matrix{Int64}:
1 3 5 7 9 11 13 15
2 4 6 8 10 12 14 16
julia> R == reshape(A, d, :)
true
julia> vec(A) == @cast v[(i,j)] := A[i,j]
true
julia> vec(transpose(A)) == @cast w[(j,i)] := A[i,j] # note the order of indices
true
You don’t have to depend on the package – you can also just copy the steps it figures out, sometimes this is easier than working out the permutations for yourself. Either by calling @pretty @cast B[i,(k,j)] := A[(i,j),k] i in 1:d
, or by just looking at the type above, you can get this answer:
julia> n = 36; d = 2;
julia> A = reshape(1:n, isqrt(n), isqrt(n));
julia> reshape(PermutedDimsArray(reshape(A, d, :, size(A,2)), (1,3,2)), d, :) # lazy
2×18 reshape(PermutedDimsArray(reshape(::UnitRange{Int64}, 2, 3, 6), (1, 3, 2)), 2, 18) with eltype Int64:
1 7 13 19 25 31 3 9 15 21 27 33 5 11 17 23 29 35
2 8 14 20 26 32 4 10 16 22 28 34 6 12 18 24 30 36
julia> reshape(permutedims(reshape(A, d, :, size(A,2)), (1,3,2)), d, :) # eager
2×18 Matrix{Int64}:
1 7 13 19 25 31 3 9 15 21 27 33 5 11 17 23 29 35
2 8 14 20 26 32 4 10 16 22 28 34 6 12 18 24 30 36