In [1]:

module O

struct AACat{T, N, U} <: AbstractArray{T, N}
    A::U
end
AACat(A) = AACat{eltype(eltype(A)), ndims(eltype(A)) + ndims(A), typeof(A)}(A)

function Base.size(C::AACat)
    A = C.A
    (size(A[begin])..., size(A)...)
end

function Base.axes(C::AACat)
    A = C.A
    (axes(A[begin])..., axes(A)...)
end

function Base.getindex(C::AACat, I::Integer...)
    A = C.A
    n = length(I) - ndims(A)
    J = I[1:n]
    K = I[n+1:end]
    getindex(A[K...], J...)
end

function Base.setindex!(C::AACat, v, I::Integer...)
    A = C.A
    n = length(I) - ndims(A)
    J = I[1:n]
    K = I[n+1:end]
    setindex!(A[K...], v, J...)
end

function aacat!(C, A::AbstractVector)
    for i in keys(A)
        C[axes(A[i])..., i] .= A[i]
    end
    C
end

function aacat!(C, A)
    for i in keys(A)
        C[axes(A[i])..., i.I...] .= A[i]
    end
    C
end

function aacat(A; A1 = A[begin])
    C = similar(A1, axes(A1)..., axes(A)...)
    aacat!(C, A)
end

end

Out[1]:

Main.O

In [2]:

A = [rand(0:9, 2, 3) for _ in Iterators.product(1:3, 1:2)]

Out[2]:

3×2 Matrix{Matrix{Int64}}:
 [9 2 9; 8 6 5]  [6 6 9; 0 4 8]
 [2 5 7; 3 4 1]  [6 5 0; 2 6 3]
 [9 2 5; 8 8 1]  [6 4 7; 0 6 2]

In [3]:

AC = O.AACat(A)

Out[3]:

2×3×3×2 Main.O.AACat{Int64, 4, Matrix{Matrix{Int64}}}:
[:, :, 1, 1] =
 9  2  9
 8  6  5

[:, :, 2, 1] =
 2  5  7
 3  4  1

[:, :, 3, 1] =
 9  2  5
 8  8  1

[:, :, 1, 2] =
 6  6  9
 0  4  8

[:, :, 2, 2] =
 6  5  0
 2  6  3

[:, :, 3, 2] =
 6  4  7
 0  6  2

In [4]:

using OffsetArrays

b = [OffsetArray(rand(0:9, 2, 3), 0:1, 0:2) for _ in Iterators.product(1:3, 1:2)]
B = OffsetArray(b, 0:2, 0:1)

Out[4]:

3×2 OffsetArray(::Matrix{OffsetMatrix{Int64, Matrix{Int64}}}, 0:2, 0:1) with eltype OffsetMatrix{Int64, Matrix{Int64}} with indices 0:2×0:1:
 [1 2 0; 3 4 0]  [9 7 2; 0 8 2]
 [6 4 8; 3 6 5]  [7 2 6; 9 5 1]
 [3 1 7; 8 8 4]  [6 2 9; 2 6 4]

In [5]:

BC = O.AACat(B)

Out[5]:

2×3×3×2 Main.O.AACat{Int64, 4, OffsetMatrix{OffsetMatrix{Int64, Matrix{Int64}}, Matrix{OffsetMatrix{Int64, Matrix{Int64}}}}} with indices 0:1×0:2×0:2×0:1:
[:, :, 0, 0] =
 1  2  0
 3  4  0

[:, :, 1, 0] =
 6  4  8
 3  6  5

[:, :, 2, 0] =
 3  1  7
 8  8  4

[:, :, 0, 1] =
 9  7  2
 0  8  2

[:, :, 1, 1] =
 7  2  6
 9  5  1

[:, :, 2, 1] =
 6  2  9
 2  6  4

In [6]:

BC[1, 1, :, :]

Out[6]:

3×2 OffsetArray(::Matrix{Int64}, 0:2, 0:1) with eltype Int64 with indices 0:2×0:1:
 4  8
 6  5
 8  6

In [7]:

BC[:, :, 1, 1]

Out[7]:

2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:
 7  2  6
 9  5  1

In [8]:

BC[1, 1, 1, 1] = 99
BC[:, :, 1, 1]

Out[8]:

2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:
 7   2  6
 9  99  1

In [9]:

using SplitApplyCombine

c = [OffsetArray(rand(0:9, 2, 3), 0:1, 0:2) for _ in Iterators.product(1:3, 1:2)]
C = OffsetArray(c, 0:2, 0:1)

CC = combinedims(C)

Out[9]:

2×3×3×2 OffsetArray(::Array{Int64, 4}, 0:1, 0:2, 0:2, 0:1) with eltype Int64 with indices 0:1×0:2×0:2×0:1:
[:, :, 0, 0] =
 4  9  8
 4  9  1

[:, :, 1, 0] =
 3  9  9
 6  1  7

[:, :, 2, 0] =
 7  3  1
 5  1  5

[:, :, 0, 1] =
 6  1  5
 1  0  6

[:, :, 1, 1] =
 2  9  5
 6  5  6

[:, :, 2, 1] =
 4  5  0
 1  1  9

In [10]:

CC[1, 1, :, :]

Out[10]:

3×2 OffsetArray(::Matrix{Int64}, 0:2, 0:1) with eltype Int64 with indices 0:2×0:1:
 9  0
 1  5
 1  1

In [11]:

CC[:, :, 1, 1]

Out[11]:

2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:
 2  9  5
 6  5  6

In [12]:

CC[1, 1, 1, 1] = 99
CC[:, :, 1, 1]

Out[12]:

2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:
 2   9  5
 6  99  6

In [13]:

d = [OffsetArray(rand(0:9, 2, 3), 0:1, 0:2) for _ in Iterators.product(1:3, 1:2)]
D = OffsetArray(c, 0:2, 0:1)

DC = combinedimsview(D)

Out[13]:

2×3×3×2 CombineDimsArray{Int64, 4, 2, OffsetMatrix{OffsetMatrix{Int64, Matrix{Int64}}, Matrix{OffsetMatrix{Int64, Matrix{Int64}}}}} with indices 0:1×0:2×0:2×0:1:
[:, :, 0, 0] =
 4  9  8
 4  9  1

[:, :, 1, 0] =
 3  9  9
 6  1  7

[:, :, 2, 0] =
 7  3  1
 5  1  5

[:, :, 0, 1] =
 6  1  5
 1  0  6

[:, :, 1, 1] =
 2  9  5
 6  5  6

[:, :, 2, 1] =
 4  5  0
 1  1  9

In [14]:

DC[1, 1, :, :]

Out[14]:

3×2 OffsetArray(::Matrix{Int64}, 0:2, 0:1) with eltype Int64 with indices 0:2×0:1:
 9  0
 1  5
 1  1

In [15]:

DC[:, :, 1, 1]

Out[15]:

2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:
 2  9  5
 6  5  6

In [16]:

DC[1, 1, 1, 1] = 99

CanonicalIndexError: setindex! not defined for CombineDimsArray{Int64, 4, 2, OffsetMatrix{OffsetMatrix{Int64, Matrix{Int64}}, Matrix{OffsetMatrix{Int64, Matrix{Int64}}}}}

Stacktrace:
 [1] error_if_canonical_setindex(::IndexCartesian, ::CombineDimsArray{Int64, 4, 2, OffsetMatrix{OffsetMatrix{Int64, Matrix{Int64}}, Matrix{OffsetMatrix{Int64, Matrix{Int64}}}}}, ::Int64, ::Int64, ::Int64, ::Int64)
   @ Base .\abstractarray.jl:1349
 [2] setindex!(::CombineDimsArray{Int64, 4, 2, OffsetMatrix{OffsetMatrix{Int64, Matrix{Int64}}, Matrix{OffsetMatrix{Int64, Matrix{Int64}}}}}, ::Int64, ::Int64, ::Int64, ::Int64, ::Vararg{Int64})
   @ Base .\abstractarray.jl:1338
 [3] top-level scope
   @ In[16]:1
 [4] eval
   @ .\boot.jl:368 [inlined]
 [5] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
   @ Base .\loading.jl:1277

In [17]:

using RecursiveArrayTools

V = [rand(0:9, 2, 3) for _ in 1:4]
VC = VectorOfArray(V)

Out[17]:

VectorOfArray{Int64,3}:
4-element Vector{Matrix{Int64}}:
 [8 7 0; 2 9 9]
 [3 5 1; 2 9 1]
 [9 3 0; 7 1 7]
 [3 3 3; 0 0 8]

In [18]:

[
    VC[1,1,1] VC[1,2,1] VC[1,3,1]
    VC[2,1,1] VC[2,2,1] VC[2,3,1]
]

Out[18]:

2×3 Matrix{Int64}:
 8  7  0
 2  9  9

In [19]:

using BenchmarkTools

V = [rand(2, 3) for _ in 1:1000]

println("---------- O.aacat")
C3 = @btime O.aacat($V)
s3 = @btime sum($C3)
@show s3

println("---------- O.AACat")
C5 = @btime O.AACat($V)
s5 = @btime sum($C5)
@show s5

println("---------- SplitApplyCombine.combinedims")
C6 = @btime combinedims($V)
s6 = @btime sum($C6)
@show s6

println("---------- SplitApplyCombine.combinedimsview")
C7 = @btime combinedimsview($V)
s7 = @btime sum($C7)
@show s7

println("---------- RecursiveArrayTools.VectorOfArray")
C8 = @btime VectorOfArray($V)
s8 = @btime sum($C8)
@show s8

println("----------")
@show s3 ≈ s5 ≈ s6 ≈ s7 ≈ s8;

---------- O.aacat
  18.700 μs (2 allocations: 46.94 KiB)
  470.408 ns (0 allocations: 0 bytes)
s3 = 2993.2209616163786
---------- O.AACat
  2.300 ns (0 allocations: 0 bytes)
  8.700 μs (0 allocations: 0 bytes)
s5 = 2993.2209616163796
---------- SplitApplyCombine.combinedims
  98.500 μs (1003 allocations: 93.83 KiB)
  467.692 ns (0 allocations: 0 bytes)
s6 = 2993.2209616163786
---------- SplitApplyCombine.combinedimsview
  2.000 ns (0 allocations: 0 bytes)
  13.400 μs (2 allocations: 32 bytes)
s7 = 2993.2209616163796
---------- RecursiveArrayTools.VectorOfArray
  6.900 ns (1 allocation: 16 bytes)
  35.100 μs (2007 allocations: 180.02 KiB)
s8 = 2993.2209616163786
----------
s3 ≈ s5 ≈ s6 ≈ s7 ≈ s8 = true

In [20]:

using BenchmarkTools

V = [rand(2, 3) for _ in 1:10^6]

println("---------- O.aacat")
C3 = @btime O.aacat($V)
s3 = @btime sum($C3)
@show s3

println("---------- O.AACat")
C5 = @btime O.AACat($V)
s5 = @btime sum($C5)
@show s5

println("---------- SplitApplyCombine.combinedims")
C6 = @btime combinedims($V)
s6 = @btime sum($C6)
@show s6

println("---------- SplitApplyCombine.combinedimsview")
C7 = @btime combinedimsview($V)
s7 = @btime sum($C7)
@show s7

println("---------- RecursiveArrayTools.VectorOfArray")
C8 = @btime VectorOfArray($V)
s8 = @btime sum($C8)
@show s8

println("----------")
@show s3 ≈ s5 ≈ s6 ≈ s7 ≈ s8;

---------- O.aacat
  24.286 ms (2 allocations: 45.78 MiB)
  2.117 ms (0 allocations: 0 bytes)
s3 = 3.000001978459132e6
---------- O.AACat
  2.000 ns (0 allocations: 0 bytes)
  10.995 ms (0 allocations: 0 bytes)
s5 = 3.0000019784594364e6
---------- SplitApplyCombine.combinedims
  114.892 ms (1000003 allocations: 91.55 MiB)
  2.095 ms (0 allocations: 0 bytes)
s6 = 3.000001978459132e6
---------- SplitApplyCombine.combinedimsview
  2.200 ns (0 allocations: 0 bytes)
  15.276 ms (2 allocations: 32 bytes)
s7 = 3.0000019784594364e6
---------- RecursiveArrayTools.VectorOfArray
  7.000 ns (1 allocation: 16 bytes)
  66.881 ms (2000008 allocations: 175.48 MiB)
s8 = 3.000001978459132e6
----------
s3 ≈ s5 ≈ s6 ≈ s7 ≈ s8 = true

In [ ]: