{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c1f520c8-8a2d-4b99-8fa8-1382bf6a2f0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Main.O"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "module O\n",
    "\n",
    "struct AACat{T, N, U} <: AbstractArray{T, N}\n",
    "    A::U\n",
    "end\n",
    "AACat(A) = AACat{eltype(eltype(A)), ndims(eltype(A)) + ndims(A), typeof(A)}(A)\n",
    "\n",
    "function Base.size(C::AACat)\n",
    "    A = C.A\n",
    "    (size(A[begin])..., size(A)...)\n",
    "end\n",
    "\n",
    "function Base.axes(C::AACat)\n",
    "    A = C.A\n",
    "    (axes(A[begin])..., axes(A)...)\n",
    "end\n",
    "\n",
    "function Base.getindex(C::AACat, I::Integer...)\n",
    "    A = C.A\n",
    "    n = length(I) - ndims(A)\n",
    "    J = I[1:n]\n",
    "    K = I[n+1:end]\n",
    "    getindex(A[K...], J...)\n",
    "end\n",
    "\n",
    "function Base.setindex!(C::AACat, v, I::Integer...)\n",
    "    A = C.A\n",
    "    n = length(I) - ndims(A)\n",
    "    J = I[1:n]\n",
    "    K = I[n+1:end]\n",
    "    setindex!(A[K...], v, J...)\n",
    "end\n",
    "\n",
    "function aacat!(C, A::AbstractVector)\n",
    "    for i in keys(A)\n",
    "        C[axes(A[i])..., i] .= A[i]\n",
    "    end\n",
    "    C\n",
    "end\n",
    "\n",
    "function aacat!(C, A)\n",
    "    for i in keys(A)\n",
    "        C[axes(A[i])..., i.I...] .= A[i]\n",
    "    end\n",
    "    C\n",
    "end\n",
    "\n",
    "function aacat(A; A1 = A[begin])\n",
    "    C = similar(A1, axes(A1)..., axes(A)...)\n",
    "    aacat!(C, A)\n",
    "end\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2cba7bae-c377-4683-819a-84387d079696",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×2 Matrix{Matrix{Int64}}:\n",
       " [9 2 9; 8 6 5]  [6 6 9; 0 4 8]\n",
       " [2 5 7; 3 4 1]  [6 5 0; 2 6 3]\n",
       " [9 2 5; 8 8 1]  [6 4 7; 0 6 2]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = [rand(0:9, 2, 3) for _ in Iterators.product(1:3, 1:2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5c87bddb-277f-4cff-9a2b-6a297cc05526",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3×3×2 Main.O.AACat{Int64, 4, Matrix{Matrix{Int64}}}:\n",
       "[:, :, 1, 1] =\n",
       " 9  2  9\n",
       " 8  6  5\n",
       "\n",
       "[:, :, 2, 1] =\n",
       " 2  5  7\n",
       " 3  4  1\n",
       "\n",
       "[:, :, 3, 1] =\n",
       " 9  2  5\n",
       " 8  8  1\n",
       "\n",
       "[:, :, 1, 2] =\n",
       " 6  6  9\n",
       " 0  4  8\n",
       "\n",
       "[:, :, 2, 2] =\n",
       " 6  5  0\n",
       " 2  6  3\n",
       "\n",
       "[:, :, 3, 2] =\n",
       " 6  4  7\n",
       " 0  6  2"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "AC = O.AACat(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4cefa303-e8e9-492e-8b90-3a528b613bcf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×2 OffsetArray(::Matrix{OffsetMatrix{Int64, Matrix{Int64}}}, 0:2, 0:1) with eltype OffsetMatrix{Int64, Matrix{Int64}} with indices 0:2×0:1:\n",
       " [1 2 0; 3 4 0]  [9 7 2; 0 8 2]\n",
       " [6 4 8; 3 6 5]  [7 2 6; 9 5 1]\n",
       " [3 1 7; 8 8 4]  [6 2 9; 2 6 4]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using OffsetArrays\n",
    "\n",
    "b = [OffsetArray(rand(0:9, 2, 3), 0:1, 0:2) for _ in Iterators.product(1:3, 1:2)]\n",
    "B = OffsetArray(b, 0:2, 0:1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9b618f6c-3c16-4860-a347-fe3b24c0cff6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "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:\n",
       "[:, :, 0, 0] =\n",
       " 1  2  0\n",
       " 3  4  0\n",
       "\n",
       "[:, :, 1, 0] =\n",
       " 6  4  8\n",
       " 3  6  5\n",
       "\n",
       "[:, :, 2, 0] =\n",
       " 3  1  7\n",
       " 8  8  4\n",
       "\n",
       "[:, :, 0, 1] =\n",
       " 9  7  2\n",
       " 0  8  2\n",
       "\n",
       "[:, :, 1, 1] =\n",
       " 7  2  6\n",
       " 9  5  1\n",
       "\n",
       "[:, :, 2, 1] =\n",
       " 6  2  9\n",
       " 2  6  4"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BC = O.AACat(B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1847d21b-ceb8-41ab-a8e5-383cc936e196",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×2 OffsetArray(::Matrix{Int64}, 0:2, 0:1) with eltype Int64 with indices 0:2×0:1:\n",
       " 4  8\n",
       " 6  5\n",
       " 8  6"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BC[1, 1, :, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e4caf847-b82a-4ccf-bb78-1361d4b70dc8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:\n",
       " 7  2  6\n",
       " 9  5  1"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BC[:, :, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1c198295-3916-46bf-8697-3d4fc5b7d46c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:\n",
       " 7   2  6\n",
       " 9  99  1"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BC[1, 1, 1, 1] = 99\n",
    "BC[:, :, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cc03388f-ab25-489b-88a4-a5e71b2211d1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "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:\n",
       "[:, :, 0, 0] =\n",
       " 4  9  8\n",
       " 4  9  1\n",
       "\n",
       "[:, :, 1, 0] =\n",
       " 3  9  9\n",
       " 6  1  7\n",
       "\n",
       "[:, :, 2, 0] =\n",
       " 7  3  1\n",
       " 5  1  5\n",
       "\n",
       "[:, :, 0, 1] =\n",
       " 6  1  5\n",
       " 1  0  6\n",
       "\n",
       "[:, :, 1, 1] =\n",
       " 2  9  5\n",
       " 6  5  6\n",
       "\n",
       "[:, :, 2, 1] =\n",
       " 4  5  0\n",
       " 1  1  9"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using SplitApplyCombine\n",
    "\n",
    "c = [OffsetArray(rand(0:9, 2, 3), 0:1, 0:2) for _ in Iterators.product(1:3, 1:2)]\n",
    "C = OffsetArray(c, 0:2, 0:1)\n",
    "\n",
    "CC = combinedims(C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4d5b900d-4586-45a5-88d7-a5265e06277a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×2 OffsetArray(::Matrix{Int64}, 0:2, 0:1) with eltype Int64 with indices 0:2×0:1:\n",
       " 9  0\n",
       " 1  5\n",
       " 1  1"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CC[1, 1, :, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4455c6e0-4d79-4398-8418-025c0fac5743",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:\n",
       " 2  9  5\n",
       " 6  5  6"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CC[:, :, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "64148b2c-05b4-49cd-bda1-e68764307226",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:\n",
       " 2   9  5\n",
       " 6  99  6"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CC[1, 1, 1, 1] = 99\n",
    "CC[:, :, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b5ca2a7c-c128-41c0-89cf-94cfa5b8eb0e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "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:\n",
       "[:, :, 0, 0] =\n",
       " 4  9  8\n",
       " 4  9  1\n",
       "\n",
       "[:, :, 1, 0] =\n",
       " 3  9  9\n",
       " 6  1  7\n",
       "\n",
       "[:, :, 2, 0] =\n",
       " 7  3  1\n",
       " 5  1  5\n",
       "\n",
       "[:, :, 0, 1] =\n",
       " 6  1  5\n",
       " 1  0  6\n",
       "\n",
       "[:, :, 1, 1] =\n",
       " 2  9  5\n",
       " 6  5  6\n",
       "\n",
       "[:, :, 2, 1] =\n",
       " 4  5  0\n",
       " 1  1  9"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = [OffsetArray(rand(0:9, 2, 3), 0:1, 0:2) for _ in Iterators.product(1:3, 1:2)]\n",
    "D = OffsetArray(c, 0:2, 0:1)\n",
    "\n",
    "DC = combinedimsview(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "fbed8c6c-8846-4e95-aa04-81f8ae59620c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×2 OffsetArray(::Matrix{Int64}, 0:2, 0:1) with eltype Int64 with indices 0:2×0:1:\n",
       " 9  0\n",
       " 1  5\n",
       " 1  1"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DC[1, 1, :, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3c250370-a28c-44e0-8884-33106775dd29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3 OffsetArray(::Matrix{Int64}, 0:1, 0:2) with eltype Int64 with indices 0:1×0:2:\n",
       " 2  9  5\n",
       " 6  5  6"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DC[:, :, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7dcd73c8-04e4-4bd4-8391-f6cdf788239c",
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "CanonicalIndexError: setindex! not defined for CombineDimsArray{Int64, 4, 2, OffsetMatrix{OffsetMatrix{Int64, Matrix{Int64}}, Matrix{OffsetMatrix{Int64, Matrix{Int64}}}}}",
     "output_type": "error",
     "traceback": [
      "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"
     ]
    }
   ],
   "source": [
    "DC[1, 1, 1, 1] = 99"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a321c224-2125-4724-a0fb-cfc20f2bf9ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "VectorOfArray{Int64,3}:\n",
       "4-element Vector{Matrix{Int64}}:\n",
       " [8 7 0; 2 9 9]\n",
       " [3 5 1; 2 9 1]\n",
       " [9 3 0; 7 1 7]\n",
       " [3 3 3; 0 0 8]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using RecursiveArrayTools\n",
    "\n",
    "V = [rand(0:9, 2, 3) for _ in 1:4]\n",
    "VC = VectorOfArray(V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "18bef1e8-2bad-4911-9a67-621126157bfc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×3 Matrix{Int64}:\n",
       " 8  7  0\n",
       " 2  9  9"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[\n",
    "    VC[1,1,1] VC[1,2,1] VC[1,3,1]\n",
    "    VC[2,1,1] VC[2,2,1] VC[2,3,1]\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0f5ecdcd-8f43-4d21-bd44-93764411b5e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- O.aacat\n",
      "  18.700 μs (2 allocations: 46.94 KiB)\n",
      "  470.408 ns (0 allocations: 0 bytes)\n",
      "s3 = 2993.2209616163786\n",
      "---------- O.AACat\n",
      "  2.300 ns (0 allocations: 0 bytes)\n",
      "  8.700 μs (0 allocations: 0 bytes)\n",
      "s5 = 2993.2209616163796\n",
      "---------- SplitApplyCombine.combinedims\n",
      "  98.500 μs (1003 allocations: 93.83 KiB)\n",
      "  467.692 ns (0 allocations: 0 bytes)\n",
      "s6 = 2993.2209616163786\n",
      "---------- SplitApplyCombine.combinedimsview\n",
      "  2.000 ns (0 allocations: 0 bytes)\n",
      "  13.400 μs (2 allocations: 32 bytes)\n",
      "s7 = 2993.2209616163796\n",
      "---------- RecursiveArrayTools.VectorOfArray\n",
      "  6.900 ns (1 allocation: 16 bytes)\n",
      "  35.100 μs (2007 allocations: 180.02 KiB)\n",
      "s8 = 2993.2209616163786\n",
      "----------\n",
      "s3 ≈ s5 ≈ s6 ≈ s7 ≈ s8 = true\n"
     ]
    }
   ],
   "source": [
    "using BenchmarkTools\n",
    "\n",
    "V = [rand(2, 3) for _ in 1:1000]\n",
    "\n",
    "println(\"---------- O.aacat\")\n",
    "C3 = @btime O.aacat($V)\n",
    "s3 = @btime sum($C3)\n",
    "@show s3\n",
    "\n",
    "println(\"---------- O.AACat\")\n",
    "C5 = @btime O.AACat($V)\n",
    "s5 = @btime sum($C5)\n",
    "@show s5\n",
    "\n",
    "println(\"---------- SplitApplyCombine.combinedims\")\n",
    "C6 = @btime combinedims($V)\n",
    "s6 = @btime sum($C6)\n",
    "@show s6\n",
    "\n",
    "println(\"---------- SplitApplyCombine.combinedimsview\")\n",
    "C7 = @btime combinedimsview($V)\n",
    "s7 = @btime sum($C7)\n",
    "@show s7\n",
    "\n",
    "println(\"---------- RecursiveArrayTools.VectorOfArray\")\n",
    "C8 = @btime VectorOfArray($V)\n",
    "s8 = @btime sum($C8)\n",
    "@show s8\n",
    "\n",
    "println(\"----------\")\n",
    "@show s3 ≈ s5 ≈ s6 ≈ s7 ≈ s8;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5cfa9bea-60e7-46f8-b1b2-75c3c97b6764",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------- O.aacat\n",
      "  24.286 ms (2 allocations: 45.78 MiB)\n",
      "  2.117 ms (0 allocations: 0 bytes)\n",
      "s3 = 3.000001978459132e6\n",
      "---------- O.AACat\n",
      "  2.000 ns (0 allocations: 0 bytes)\n",
      "  10.995 ms (0 allocations: 0 bytes)\n",
      "s5 = 3.0000019784594364e6\n",
      "---------- SplitApplyCombine.combinedims\n",
      "  114.892 ms (1000003 allocations: 91.55 MiB)\n",
      "  2.095 ms (0 allocations: 0 bytes)\n",
      "s6 = 3.000001978459132e6\n",
      "---------- SplitApplyCombine.combinedimsview\n",
      "  2.200 ns (0 allocations: 0 bytes)\n",
      "  15.276 ms (2 allocations: 32 bytes)\n",
      "s7 = 3.0000019784594364e6\n",
      "---------- RecursiveArrayTools.VectorOfArray\n",
      "  7.000 ns (1 allocation: 16 bytes)\n",
      "  66.881 ms (2000008 allocations: 175.48 MiB)\n",
      "s8 = 3.000001978459132e6\n",
      "----------\n",
      "s3 ≈ s5 ≈ s6 ≈ s7 ≈ s8 = true\n"
     ]
    }
   ],
   "source": [
    "using BenchmarkTools\n",
    "\n",
    "V = [rand(2, 3) for _ in 1:10^6]\n",
    "\n",
    "println(\"---------- O.aacat\")\n",
    "C3 = @btime O.aacat($V)\n",
    "s3 = @btime sum($C3)\n",
    "@show s3\n",
    "\n",
    "println(\"---------- O.AACat\")\n",
    "C5 = @btime O.AACat($V)\n",
    "s5 = @btime sum($C5)\n",
    "@show s5\n",
    "\n",
    "println(\"---------- SplitApplyCombine.combinedims\")\n",
    "C6 = @btime combinedims($V)\n",
    "s6 = @btime sum($C6)\n",
    "@show s6\n",
    "\n",
    "println(\"---------- SplitApplyCombine.combinedimsview\")\n",
    "C7 = @btime combinedimsview($V)\n",
    "s7 = @btime sum($C7)\n",
    "@show s7\n",
    "\n",
    "println(\"---------- RecursiveArrayTools.VectorOfArray\")\n",
    "C8 = @btime VectorOfArray($V)\n",
    "s8 = @btime sum($C8)\n",
    "@show s8\n",
    "\n",
    "println(\"----------\")\n",
    "@show s3 ≈ s5 ≈ s6 ≈ s7 ≈ s8;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c12e806-2d72-4c44-99ff-a027055fdbc6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md"
  },
  "kernelspec": {
   "display_name": "Julia 1.8.0-beta1",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}