From 5fc092b36379dd6eb196f763907426638f3dfdfa Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Wed, 27 Nov 2024 16:35:09 +0100
Subject: [PATCH 01/14] just add the supertype to see what happens

---
 src/stack/stack.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/stack/stack.jl b/src/stack/stack.jl
index 68a645faf..84e067d4b 100644
--- a/src/stack/stack.jl
+++ b/src/stack/stack.jl
@@ -25,7 +25,7 @@ To extend `AbstractDimStack`, implement argument and keyword version of
 
 The constructor of an `AbstractDimStack` must accept a `NamedTuple`.
 """
-abstract type AbstractDimStack{K,T,N,L} end
+abstract type AbstractDimStack{K,T,N,L} <: AbstractArray{T,N} end
 const AbstractVectorDimStack = AbstractDimStack{K,T,1} where {K,T}
 const AbstractMatrixDimStack = AbstractDimStack{K,T,2} where {K,T}
 

From bd89aaa37bcea991d1bc31fecd42e0d63f399d10 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Wed, 27 Nov 2024 17:24:29 +0100
Subject: [PATCH 02/14] dfeine as AbstractBasicDimArray and fix some
 ambiguities

---
 src/stack/methods.jl |  6 ++++--
 src/stack/show.jl    |  2 +-
 src/stack/stack.jl   | 42 ++++++++----------------------------------
 test/stack.jl        |  3 ++-
 4 files changed, 15 insertions(+), 38 deletions(-)

diff --git a/src/stack/methods.jl b/src/stack/methods.jl
index 272cd2146..bf7a5fa1d 100644
--- a/src/stack/methods.jl
+++ b/src/stack/methods.jl
@@ -133,8 +133,10 @@ end
 
 # Methods with an argument that return a DimStack
 for fname in (:rotl90, :rotr90, :rot180)
-    @eval (Base.$fname)(s::AbstractDimStack, args...) =
-        maplayers(A -> (Base.$fname)(A, args...), s)
+    @eval (Base.$fname)(s::AbstractDimStack) =
+        maplayers(A -> (Base.$fname)(A), s)
+    @eval (Base.$fname)(s::AbstractDimStack, k::Integer) =
+        maplayers(A -> (Base.$fname)(A, k), s)
 end
 for fname in (:PermutedDimsArray, :permutedims)
     @eval function (Base.$fname)(s::AbstractDimStack, perm)
diff --git a/src/stack/show.jl b/src/stack/show.jl
index 2e841f80d..0a8e67926 100644
--- a/src/stack/show.jl
+++ b/src/stack/show.jl
@@ -20,7 +20,7 @@ function show_main(io, mime, stack::AbstractDimStack)
     _, blockwidth = print_metadata_block(io, mime, metadata(stack); displaywidth, blockwidth=min(displaywidth-2, blockwidth))
 end
 
-function show_after(io, mime, stack::AbstractDimStack)
+function show_after(io::IO, mime, stack::AbstractDimStack)
     blockwidth = get(io, :blockwidth, 0)
     print_block_close(io, blockwidth)
 end
diff --git a/src/stack/stack.jl b/src/stack/stack.jl
index 84e067d4b..d0369534d 100644
--- a/src/stack/stack.jl
+++ b/src/stack/stack.jl
@@ -25,7 +25,7 @@ To extend `AbstractDimStack`, implement argument and keyword version of
 
 The constructor of an `AbstractDimStack` must accept a `NamedTuple`.
 """
-abstract type AbstractDimStack{K,T,N,L} <: AbstractArray{T,N} end
+abstract type AbstractDimStack{K,T,N,L,D} <: AbstractBasicDimArray{T,N,D} end
 const AbstractVectorDimStack = AbstractDimStack{K,T,1} where {K,T}
 const AbstractMatrixDimStack = AbstractDimStack{K,T,2} where {K,T}
 
@@ -140,50 +140,24 @@ Base.:(==)(s1::AbstractDimStack, s2::AbstractDimStack) =
     data(s1) == data(s2) && dims(s1) == dims(s2) && layerdims(s1) == layerdims(s2)
 Base.read(s::AbstractDimStack) = maplayers(read, s)
 
-# Array-like
-Base.size(s::AbstractDimStack) = map(length, dims(s))
-Base.size(s::AbstractDimStack, dims::DimOrDimType) = size(s, dimnum(s, dims))
-Base.size(s::AbstractDimStack, dims::Integer) = size(s)[dims]
-Base.length(s::AbstractDimStack) = prod(size(s))
-Base.axes(s::AbstractDimStack) = map(first ∘ axes, dims(s))
-Base.axes(s::AbstractDimStack, dims::DimOrDimType) = axes(s, dimnum(s, dims))
-Base.axes(s::AbstractDimStack, dims::Integer) = axes(s)[dims]
-Base.similar(s::AbstractDimStack, args...) = maplayers(A -> similar(A, args...), s)
-Base.eltype(::AbstractDimStack{<:Any,T}) where T = T
-Base.ndims(::AbstractDimStack{<:Any,<:Any,N}) where N = N
-Base.CartesianIndices(s::AbstractDimStack) = CartesianIndices(dims(s))
-Base.LinearIndices(s::AbstractDimStack) = 
-    LinearIndices(CartesianIndices(map(l -> axes(l, 1), lookup(s))))
-Base.IteratorSize(::AbstractDimStack{<:Any,<:Any,N}) where N = Base.HasShape{N}()
-function Base.eachindex(s::AbstractDimStack)
-    li = LinearIndices(s)
-    first(li):last(li)
-end
-Base.firstindex(s::AbstractDimStack) = first(LinearIndices(s))
-Base.lastindex(s::AbstractDimStack) = last(LinearIndices(s))
-Base.first(s::AbstractDimStack) = s[firstindex((s))]
-Base.last(s::AbstractDimStack) = s[lastindex(LinearIndices(s))]
-Base.copy(s::AbstractDimStack) = modify(copy, s)
-# all of methods.jl is also Array-like...
+# Array interface methods
+Base.IndexStyle(A::AbstractDimStack) = Base.IndexStyle(first(layers(A)))
+
+Base.similar(s::AbstractDimStack; kw...) = maplayers(A -> similar(A; kw...), s)
+Base.similar(s::AbstractDimStack, ::Type{T}; kw...) where T = maplayers(A -> similar(A, T; kw...), s)
+Base.similar(s::AbstractDimStack, ::Type{T}, dt::DimTuple; kw...) where T = maplayers(A -> similar(A, T, dt; kw...), s)
 
 # NamedTuple-like
 @assume_effects :foldable Base.getproperty(s::AbstractDimStack, x::Symbol) = s[x]
 Base.haskey(s::AbstractDimStack{K}, k) where K = k in K
 Base.values(s::AbstractDimStack) = values(layers(s))
-Base.checkbounds(s::AbstractDimStack, I...) = checkbounds(CartesianIndices(s), I...)
-Base.checkbounds(T::Type, s::AbstractDimStack, I...) = checkbounds(T, CartesianIndices(s), I...)
 
 @inline Base.keys(s::AbstractDimStack{K}) where K = K
 @inline Base.propertynames(s::AbstractDimStack{K}) where K = K
 @inline Base.setindex(s::AbstractDimStack, val::AbstractBasicDimArray, name::Symbol) =
     rebuild_from_arrays(s, Base.setindex(layers(s), val, name))
 Base.NamedTuple(s::AbstractDimStack) = NamedTuple(layers(s))
-Base.collect(st::AbstractDimStack) = parent([st[D] for D in DimIndices(st)])
 Base.Array(st::AbstractDimStack) = collect(st)
-Base.vec(st::AbstractDimStack) = vec(collect(st))
-@propagate_inbounds Base.iterate(st::AbstractDimStack) = iterate(st, 1)
-@propagate_inbounds Base.iterate(st::AbstractDimStack, i) =
-    i > length(st) ? nothing : (st[DimIndices(st)[i]], i + 1)
 
 # `merge` for AbstractDimStack and NamedTuple.
 # One of the first three arguments must be an AbstractDimStack for dispatch to work.
@@ -377,7 +351,7 @@ julia> s[X(At(:a))] isa DimStack
 true
 ```
 """
-struct DimStack{K,T,N,L,D<:Tuple,R<:Tuple,LD,M,LM} <: AbstractDimStack{K,T,N,L}
+struct DimStack{K,T,N,L,D<:Tuple,R<:Tuple,LD,M,LM} <: AbstractDimStack{K,T,N,L,D}
     data::L
     dims::D
     refdims::R
diff --git a/test/stack.jl b/test/stack.jl
index ba5ef205a..d1918644a 100644
--- a/test/stack.jl
+++ b/test/stack.jl
@@ -1,5 +1,6 @@
 using DimensionalData, Test, LinearAlgebra, Statistics, ConstructionBase, Random
 
+using DimensionalData
 using DimensionalData: data
 using DimensionalData: Sampled, Categorical, AutoLookup, NoLookup, Transformed,
     Regular, Irregular, Points, Intervals, Start, Center, End,
@@ -73,7 +74,7 @@ end
     @test length(mixed) === 24
     @test firstindex(mixed) === 1
     @test lastindex(mixed) === 24
-    @test eachindex(mixed) === 1:24
+    @test eachindex(mixed) === Base.OneTo(24)
     @test axes(mixed) == (Base.OneTo(2), Base.OneTo(3), Base.OneTo(4))
     @test eltype(axes(mixed)) <: Dimensions.DimUnitRange
     @test dims(axes(mixed)) == dims(mixed)

From 1b5d81298cffdabbcd73e215589a8178c6c714a4 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Wed, 27 Nov 2024 18:56:41 +0100
Subject: [PATCH 03/14] clean up ambiguities

---
 src/array/array.jl    |  2 +-
 src/stack/indexing.jl | 43 ++++---------------------------------------
 src/stack/stack.jl    | 11 +++++++++--
 test/stack.jl         |  4 ++--
 4 files changed, 16 insertions(+), 44 deletions(-)

diff --git a/src/array/array.jl b/src/array/array.jl
index b4cc08a60..7a7930c50 100644
--- a/src/array/array.jl
+++ b/src/array/array.jl
@@ -475,7 +475,7 @@ function DimArray(A::AbstractBasicDimArray;
     data=parent(A), dims=dims(A), refdims=refdims(A), name=name(A), metadata=metadata(A)
 )
     newdata = collect(data)
-    DimArray(newdata, format(dims, newdata); refdims, name, metadata)
+    T(newdata, format(dims, newdata); refdims, name, metadata)
 end
 """
     DimArray(f::Function, dim::Dimension; [name])
diff --git a/src/stack/indexing.jl b/src/stack/indexing.jl
index e3d41eb0a..80fc7af1c 100644
--- a/src/stack/indexing.jl
+++ b/src/stack/indexing.jl
@@ -28,17 +28,6 @@ Base.@assume_effects :effect_free @propagate_inbounds function Base.getindex(s::
     # Use dimensional indexing
     Base.getindex(s, rebuild(only(dims(s)), i))
 end
-Base.@assume_effects :effect_free @propagate_inbounds function Base.getindex(
-    s::AbstractDimStack{<:Any,T}, i::Union{AbstractArray,Colon}
-) where {T}
-    ls = _maybe_extented_layers(s)
-    inds = to_indices(first(ls), (i,))[1]
-    out = similar(inds, T)
-    for (i, ind) in enumerate(inds)
-        out[i] = T(map(v -> v[ind], ls))
-    end
-    return out
-end
 @propagate_inbounds function Base.getindex(s::AbstractDimStack{<:Any,<:Any,N}, i::Integer) where N
     if N == 1 && hassamedims(s)
         # This is a few ns faster when possible
@@ -52,28 +41,10 @@ end
 @propagate_inbounds function Base.view(s::AbstractVectorDimStack, i::Union{AbstractVector{<:Integer},Colon,Integer})
     Base.view(s, DimIndices(s)[i])
 end
-@propagate_inbounds function Base.view(s::AbstractDimStack, i::Union{AbstractArray{<:Integer},Colon,Integer})
-    # Pretend the stack is an AbstractArray so `SubArray` accepts it.
-    Base.view(OpaqueArray(s), i)
-end
 
 for f in (:getindex, :view, :dotview)
     _dim_f = Symbol(:_dim_, f)
     @eval begin
-        @propagate_inbounds function Base.$f(s::AbstractDimStack, i)
-            Base.$f(s, to_indices(CartesianIndices(s), Lookups._construct_types(i))...)
-        end
-        @propagate_inbounds function Base.$f(s::AbstractDimStack, i::Union{SelectorOrInterval,Extents.Extent})
-            Base.$f(s, dims2indices(s, i)...)
-        end
-        @propagate_inbounds function Base.$f(s::AbstractVectorDimStack, i::Union{CartesianIndices,CartesianIndex})
-            I = to_indices(CartesianIndices(s), (i,))
-            Base.$f(s, I...)
-        end
-        @propagate_inbounds function Base.$f(s::AbstractDimStack, i::Union{CartesianIndices,CartesianIndex})
-            I = to_indices(CartesianIndices(s), (i,))
-            Base.$f(s, I...)
-        end
         @propagate_inbounds function Base.$f(s::AbstractDimStack, i1, i2, Is...)
             I = to_indices(CartesianIndices(s), Lookups._construct_types(i1, i2, Is...))
             # Check we have the right number of dimensions
@@ -95,12 +66,10 @@ for f in (:getindex, :view, :dotview)
             # Convert to Dimension wrappers to handle mixed size layers
             Base.$f(s, map(rebuild, dims(s), I)...)
         end
-        @propagate_inbounds function Base.$f(
-            s::AbstractDimStack, D::DimensionalIndices...; kw...
-        )
-            $_dim_f(s, _simplify_dim_indices(D..., kw2dims(values(kw))...)...)
-        end
         # Ambiguities
+        @propagate_inbounds Base.$f(A::AbstractDimStack, d1::DimensionalIndices, d2::DimensionalIndices, D::DimensionalIndices...; kw...) =
+            $_dim_f(A, _simplify_dim_indices(d1, d2, D..., kw2dims(values(kw))...)...)
+
         @propagate_inbounds function Base.$f(s::DimensionalData.AbstractVectorDimStack, 
             i::Union{AbstractVector{<:DimensionalData.Dimensions.Dimension},
             AbstractVector{<:Tuple{DimensionalData.Dimensions.Dimension, Vararg{DimensionalData.Dimensions.Dimension}}}, 
@@ -109,11 +78,10 @@ for f in (:getindex, :view, :dotview)
             $_dim_f(s, _simplify_dim_indices(i)...)
         end
 
-
         @propagate_inbounds function $_dim_f(
             A::AbstractDimStack, a1::Union{Dimension,DimensionIndsArrays}, args::Union{Dimension,DimensionIndsArrays}...
         )
-            return merge_and_index(Base.$f, A, (a1, args...))
+            return merge_and_index($_dim_f, A, (a1, args...))
         end
         # Handle zero-argument getindex, this will error unless all layers are zero dimensional
         @propagate_inbounds function $_dim_f(s::AbstractDimStack)
@@ -175,9 +143,6 @@ end
 
 @noinline _keysmismatch(K1, K2) = throw(ArgumentError("NamedTuple keys $K2 do not mach stack keys $K1"))
 
-# For @views macro to work with keywords
-Base.maybeview(A::AbstractDimStack, args...; kw...) = view(A, args...; kw...)
-
 function merge_and_index(f, s::AbstractDimStack, ds)
     ds, inds_arrays = _separate_dims_arrays(_simplify_dim_indices(ds...)...)
     # No arrays here, so abort (dispatch is tricky...)
diff --git a/src/stack/stack.jl b/src/stack/stack.jl
index d0369534d..fc9bf4849 100644
--- a/src/stack/stack.jl
+++ b/src/stack/stack.jl
@@ -25,12 +25,19 @@ To extend `AbstractDimStack`, implement argument and keyword version of
 
 The constructor of an `AbstractDimStack` must accept a `NamedTuple`.
 """
-abstract type AbstractDimStack{K,T,N,L,D} <: AbstractBasicDimArray{T,N,D} end
+abstract type AbstractDimStack{K,T<:NamedTuple,N,L,D} <: AbstractBasicDimArray{T,N,D} end
 const AbstractVectorDimStack = AbstractDimStack{K,T,1} where {K,T}
 const AbstractMatrixDimStack = AbstractDimStack{K,T,2} where {K,T}
 
-(::Type{T})(st::AbstractDimStack) where T<:AbstractDimArray =
+(::Type{T})(st::AbstractDimStack) where {T<:AbstractDimArray} =
     T([st[D] for D in DimIndices(st)]; dims=dims(st), metadata=metadata(st))
+DimArray(st::AbstractDimStack) =
+    DimArray(collect(st), dims(st); metadata=metadata(st))
+DimMatrix(st::AbstractMatrixDimStack) =
+    DimArray(collect(st), dims(st); metadata=metadata(st))
+DimVector(st::AbstractVectorDimStack) =
+    DimArray(collect(st), dims(st); metadata=metadata(st))
+
 
 data(s::AbstractDimStack) = getfield(s, :data)
 dims(s::AbstractDimStack) = getfield(s, :dims)
diff --git a/test/stack.jl b/test/stack.jl
index d1918644a..10a512405 100644
--- a/test/stack.jl
+++ b/test/stack.jl
@@ -83,8 +83,8 @@ end
     @test axes(mixed, 2) == Base.OneTo(3)
     @test lastindex(mixed, 2) == 3
     @test dims(axes(mixed, 2)) == dims(mixed, 2)
-    @test first(mixed) == mixed[Begin] 
-    @test last(mixed) == mixed[End]
+    @test first(mixed) == mixed[Begin()] 
+    @test last(mixed) == mixed[End()]
     @test NamedTuple(s) == (one=da1, two=da2, three=da3)
 end
 

From 1ba7d2b2dea4bafa8cbb0dfb9afebed37e451dd5 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Wed, 27 Nov 2024 19:06:23 +0100
Subject: [PATCH 04/14] fix broadcast_dims

---
 src/array/array.jl |  2 +-
 src/utils.jl       | 20 ++++++++++----------
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/array/array.jl b/src/array/array.jl
index 7a7930c50..b4cc08a60 100644
--- a/src/array/array.jl
+++ b/src/array/array.jl
@@ -475,7 +475,7 @@ function DimArray(A::AbstractBasicDimArray;
     data=parent(A), dims=dims(A), refdims=refdims(A), name=name(A), metadata=metadata(A)
 )
     newdata = collect(data)
-    T(newdata, format(dims, newdata); refdims, name, metadata)
+    DimArray(newdata, format(dims, newdata); refdims, name, metadata)
 end
 """
     DimArray(f::Function, dim::Dimension; [name])
diff --git a/src/utils.jl b/src/utils.jl
index fe3eb4172..fe25478c6 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -123,17 +123,17 @@ sliced by the dimensions of `B`.
 """
 function broadcast_dims(f, As::AbstractBasicDimArray...)
     dims = combinedims(As...)
-    T = Base.Broadcast.combine_eltypes(f, As)
-    broadcast_dims!(f, similar(first(As), T, dims), As...)
-end
-
-function broadcast_dims(f, As::Union{AbstractDimStack,AbstractBasicDimArray}...)
-    st = _firststack(As...)
-    nts = _as_extended_nts(NamedTuple(st), As...)
-    layers = map(keys(st)) do name
-        broadcast_dims(f, map(nt -> nt[name], nts)...)
+    if any(map(A -> A isa AbstractDimStack, As))
+        st = _firststack(As...)
+        nts = _as_extended_nts(NamedTuple(st), As...)
+        layers = map(keys(st)) do name
+            broadcast_dims(f, map(nt -> nt[name], nts)...)
+        end
+        rebuild_from_arrays(st, layers)
+    else
+        T = Base.Broadcast.combine_eltypes(f, As)
+        broadcast_dims!(f, similar(first(As), T, dims), As...)
     end
-    rebuild_from_arrays(st, layers)
 end
 
 """

From 96f52029415ae38feaa83cbaf267d810e17c45d1 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Wed, 27 Nov 2024 19:18:34 +0100
Subject: [PATCH 05/14] fix DimStack type definition

---
 src/stack/stack.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/stack/stack.jl b/src/stack/stack.jl
index fc9bf4849..98b1a0ff3 100644
--- a/src/stack/stack.jl
+++ b/src/stack/stack.jl
@@ -358,7 +358,7 @@ julia> s[X(At(:a))] isa DimStack
 true
 ```
 """
-struct DimStack{K,T,N,L,D<:Tuple,R<:Tuple,LD,M,LM} <: AbstractDimStack{K,T,N,L,D}
+struct DimStack{K,T<:NamedTuple,N,L,D<:Tuple,R<:Tuple,LD,M,LM} <: AbstractDimStack{K,T,N,L,D}
     data::L
     dims::D
     refdims::R

From e25d5dfff02b15424180c02e3c61592346937cb4 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Wed, 27 Nov 2024 19:24:27 +0100
Subject: [PATCH 06/14] fix more dispatches

---
 src/array/array.jl | 2 +-
 src/utils.jl       | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/array/array.jl b/src/array/array.jl
index b4cc08a60..c57912870 100644
--- a/src/array/array.jl
+++ b/src/array/array.jl
@@ -95,7 +95,7 @@ metadata(A::AbstractDimArray) = A.metadata
 layerdims(A::AbstractDimArray) = basedims(A)
 
 @inline rebuildsliced(A::AbstractBasicDimArray, args...) = rebuildsliced(getindex, A, args...)
-@inline function rebuildsliced(f::Function, A::AbstractBasicDimArray, data::AbstractArray, I::Tuple, name=name(A))
+@inline function rebuildsliced(f::Function, A::AbstractBasicDimArray, data, I, name=name(A))
     I1 = to_indices(A, I)
     rebuild(A, data, slicedims(f, A, I1)..., name)
 end
diff --git a/src/utils.jl b/src/utils.jl
index fe25478c6..3a7375db6 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -85,7 +85,6 @@ modify(CuArray, A)
 This also works for all the data layers in a `DimStack`.
 """
 function modify end
-modify(f, s::AbstractDimStack) = maplayers(a -> modify(f, a), s)
 # Stack optimisation to avoid compilation to build all the `AbstractDimArray` 
 # layers, and instead just modify the parent data directly.
 modify(f, s::AbstractDimStack{<:Any,<:NamedTuple}) = 

From ced12d60c4c3085925915219e6a5e8b22173d37f Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Wed, 27 Nov 2024 19:44:45 +0100
Subject: [PATCH 07/14] add maybeview dispatches

---
 src/array/indexing.jl | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/array/indexing.jl b/src/array/indexing.jl
index 762d5c4cd..df48dd91d 100644
--- a/src/array/indexing.jl
+++ b/src/array/indexing.jl
@@ -265,3 +265,7 @@ Base.@assume_effects :foldable @inline _simplify_dim_indices() = ()
     view(A, args...; kw...)
 @propagate_inbounds Base.maybeview(A::AbstractDimArray, args::Vararg{Union{Number,Base.AbstractCartesianIndex}}; kw...) =
     view(A, args...; kw...)
+@propagate_inbounds Base.maybeview(A::AbstractDimStack, args...; kw...) =
+    view(A, args...; kw...)
+@propagate_inbounds Base.maybeview(A::AbstractDimStack, args::Vararg{Union{Number,Base.AbstractCartesianIndex}}; kw...) =
+    view(A, args...; kw...)

From 987e0404fde8e63e753a81d4b47db8de1326d4a1 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Thu, 28 Nov 2024 10:08:25 +0100
Subject: [PATCH 08/14] fix indexing

---
 src/array/indexing.jl |  6 +-----
 src/stack/indexing.jl | 17 ++++++-----------
 2 files changed, 7 insertions(+), 16 deletions(-)

diff --git a/src/array/indexing.jl b/src/array/indexing.jl
index df48dd91d..9d30017a2 100644
--- a/src/array/indexing.jl
+++ b/src/array/indexing.jl
@@ -264,8 +264,4 @@ Base.@assume_effects :foldable @inline _simplify_dim_indices() = ()
 @propagate_inbounds Base.maybeview(A::AbstractDimArray, args...; kw...) =
     view(A, args...; kw...)
 @propagate_inbounds Base.maybeview(A::AbstractDimArray, args::Vararg{Union{Number,Base.AbstractCartesianIndex}}; kw...) =
-    view(A, args...; kw...)
-@propagate_inbounds Base.maybeview(A::AbstractDimStack, args...; kw...) =
-    view(A, args...; kw...)
-@propagate_inbounds Base.maybeview(A::AbstractDimStack, args::Vararg{Union{Number,Base.AbstractCartesianIndex}}; kw...) =
-    view(A, args...; kw...)
+    view(A, args...; kw...)
\ No newline at end of file
diff --git a/src/stack/indexing.jl b/src/stack/indexing.jl
index 80fc7af1c..acc3f19e5 100644
--- a/src/stack/indexing.jl
+++ b/src/stack/indexing.jl
@@ -49,11 +49,6 @@ for f in (:getindex, :view, :dotview)
             I = to_indices(CartesianIndices(s), Lookups._construct_types(i1, i2, Is...))
             # Check we have the right number of dimensions
             if length(dims(s)) > length(I)
-        @propagate_inbounds function $_dim_f(
-            A::AbstractDimStack, a1::Union{Dimension,DimensionIndsArrays}, args::Union{Dimension,DimensionIndsArrays}...
-        )
-            return merge_and_index(Base.$f, A, (a1, args...))
-        end
                 throw(BoundsError(dims(s), I))
             elseif length(dims(s)) < length(I)
                 # Allow trailing ones
@@ -78,11 +73,6 @@ for f in (:getindex, :view, :dotview)
             $_dim_f(s, _simplify_dim_indices(i)...)
         end
 
-        @propagate_inbounds function $_dim_f(
-            A::AbstractDimStack, a1::Union{Dimension,DimensionIndsArrays}, args::Union{Dimension,DimensionIndsArrays}...
-        )
-            return merge_and_index($_dim_f, A, (a1, args...))
-        end
         # Handle zero-argument getindex, this will error unless all layers are zero dimensional
         @propagate_inbounds function $_dim_f(s::AbstractDimStack)
             map(Base.$f, data(s))
@@ -129,7 +119,7 @@ end
     map((A, x) -> setindex!(A, x, I...; kw...), layers(s), xs)
 end
 
-_map_setindex!(s, xs, i) = map((A, x) -> setindex!(A, x, i...; kw...), layers(s), xs)
+_map_setindex!(s, xs, i; kw...) = map((A, x) -> setindex!(A, x, i...; kw...), layers(s), xs)
 
 _setindex_mixed!(s::AbstractDimStack, x, i::AbstractArray) =
     map(A -> setindex!(A, x, DimIndices(dims(s))[i]), layers(s))
@@ -160,3 +150,8 @@ function merge_and_index(f, s::AbstractDimStack, ds)
     end
     return rebuild_from_arrays(s, newlayers)
 end
+
+@propagate_inbounds Base.maybeview(A::AbstractDimStack, args...; kw...) =
+    view(A, args...; kw...)
+@propagate_inbounds Base.maybeview(A::AbstractDimStack, args::Vararg{Union{Number,Base.AbstractCartesianIndex}}; kw...) =
+    view(A, args...; kw...)

From 41651356825f22b081996872687e5b7f9ad4643c Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Thu, 28 Nov 2024 10:09:06 +0100
Subject: [PATCH 09/14] remove superfluous rand dispatch

---
 src/stack/methods.jl | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/src/stack/methods.jl b/src/stack/methods.jl
index bf7a5fa1d..037706161 100644
--- a/src/stack/methods.jl
+++ b/src/stack/methods.jl
@@ -190,11 +190,4 @@ for fname in (:one, :oneunit, :zero, :copy)
     end
 end
 
-Base.reverse(s::AbstractDimStack; dims=:) = maplayers(A -> reverse(A; dims=dims), s)
-
-# Random
-Random.Sampler(RNG::Type{<:AbstractRNG}, st::AbstractDimStack, n::Random.Repetition) =
-    Random.SamplerSimple(st, Random.Sampler(RNG, DimIndices(st), n))
-
-Random.rand(rng::AbstractRNG, sp::Random.SamplerSimple{<:AbstractDimStack,<:Random.Sampler}) =
-    @inbounds return sp[][rand(rng, sp.data)...]
+Base.reverse(s::AbstractDimStack; dims=:) = maplayers(A -> reverse(A; dims=dims), s)
\ No newline at end of file

From 623e5d49eeee9e99ba899c904e4a7c3140f9a8b5 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Thu, 28 Nov 2024 10:41:18 +0100
Subject: [PATCH 10/14] remove checkbounds again

---
 src/stack/stack.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/stack/stack.jl b/src/stack/stack.jl
index 5f436f017..54c2c4e3c 100644
--- a/src/stack/stack.jl
+++ b/src/stack/stack.jl
@@ -166,8 +166,6 @@ Base.values(s::AbstractDimStack) = _values_gen(s)
 @generated function _values_gen(s::AbstractDimStack{K}) where K
     Expr(:tuple, map(k -> :(s[$(QuoteNode(k))]), K)...)
 end
-Base.checkbounds(s::AbstractDimStack, I...) = checkbounds(CartesianIndices(s), I...)
-Base.checkbounds(T::Type, s::AbstractDimStack, I...) = checkbounds(T, CartesianIndices(s), I...)
 
 @inline Base.keys(s::AbstractDimStack{K}) where K = K
 @inline Base.propertynames(s::AbstractDimStack{K}) where K = K

From ef61ecdf6ede288d055b05f10b0cfb42924ddd6d Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Thu, 28 Nov 2024 17:26:47 +0100
Subject: [PATCH 11/14] change behaviour of broadcast_dims

---
 src/utils.jl  | 20 ++++++++------------
 test/utils.jl |  9 ++++++---
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index cb9b4b127..cc3240668 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -122,17 +122,13 @@ sliced by the dimensions of `B`.
 """
 function broadcast_dims(f, As::AbstractBasicDimArray...)
     dims = combinedims(As...)
-    if any(map(A -> A isa AbstractDimStack, As))
-        st = _firststack(As...)
-        nts = _as_extended_nts(NamedTuple(st), As...)
-        layers = map(keys(st)) do name
-            broadcast_dims(f, map(nt -> nt[name], nts)...)
-        end
-        rebuild_from_arrays(st, layers)
-    else
-        T = Base.Broadcast.combine_eltypes(f, As)
-        broadcast_dims!(f, similar(first(As), T, dims), As...)
-    end
+    T = Base.Broadcast.combine_eltypes(f, As)
+    broadcast_dims!(f, similar(first(As), T, dims), As...)
+end
+function broadcast_dims(f, s::AbstractDimStack, As::AbstractBasicDimArray...)
+    dims = combinedims(s, As...)
+    T = Base.Broadcast.combine_eltypes(f, (s, As...))
+    broadcast_dims!(f, similar(first(layers(s)), T, dims), s, As...)
 end
 
 """
@@ -149,7 +145,7 @@ which they are found.
 - `dest`: `AbstractDimArray` to update.
 - `sources`: `AbstractDimArrays` to broadcast over with `f`.
 """
-function broadcast_dims!(f, dest::AbstractDimArray{<:Any,N}, As::AbstractBasicDimArray...) where {N}
+function broadcast_dims!(f, dest::Union{AbstractDimArray{<:Any,N}, AbstractDimStack{<:Any,<:Any,N}}, As::AbstractBasicDimArray...) where {N}
     As = map(As) do A
         isempty(otherdims(A, dims(dest))) || throw(DimensionMismatch("Cannot broadcast over dimensions not in the dest array"))
         # comparedims(dest, dims(A, dims(dest)))
diff --git a/test/utils.jl b/test/utils.jl
index 91dc30a6d..96d145ad1 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -179,9 +179,12 @@ end
         db1 = DimArray(B1, (Y([:a, :b, :c]),))
         stack1 = DimStack(da3, db1)
         stack2 = DimStack(da3, db1, dc3)
-        @test broadcast_dims(+, stack1, da3, db1) == broadcast_dims(+, da3, db1, stack1)
-        @test broadcast_dims(+, stack1, da3, db1).layer1 == broadcast_dims(+, stack1.layer1, da3, db1)
-        @test broadcast_dims(+, stack1, da3, stack1, db1) == broadcast_dims(+, da3, stack1, db1, stack1)
+        @test broadcast_dims((s, x1, x2) -> s.layer1, stack1, da3, db1) == 
+            broadcast_dims((x1, x2, s) -> s.layer1, da3, db1, stack1)
+        @test broadcast_dims((s, x1, x2) -> s.layer1 + x1 + x2, stack1, da3, db1) == 
+            broadcast_dims(+, stack1.layer1, da3, db1)
+        @test broadcast_dims((s1, x1, s2, x2) -> s1.layer1+x2, stack1, da3, stack1, db1) == 
+            broadcast_dims((x1, s1, x2, s2) -> s1.layer1+x2, da3, stack1, db1, stack1)
         # Cant mix numvers of stack layers
         @test_throws ArgumentError broadcast_dims(+, stack1, da3, db1, stack2)
     end

From 46be6136ce507dc516f8939a600ddc6476e52466 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Fri, 29 Nov 2024 14:49:10 +0100
Subject: [PATCH 12/14] loosen maplayers dispatch

---
 src/stack/stack.jl | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/stack/stack.jl b/src/stack/stack.jl
index 54c2c4e3c..5448a15dd 100644
--- a/src/stack/stack.jl
+++ b/src/stack/stack.jl
@@ -208,14 +208,30 @@ Base.map(f, ::Union{AbstractDimStack,NamedTuple}, xs::Union{AbstractDimStack,Nam
 maplayers(f, s::AbstractDimStack) =
     _maybestack(s, unrolled_map(f, values(s)))
 function maplayers(
-    f, x1::Union{AbstractDimStack,NamedTuple}, xs::Union{AbstractDimStack,NamedTuple}...
+    f, x1::Union{AbstractBasicDimArray,NamedTuple}, xs::Union{AbstractBasicDimArray,NamedTuple}...
 )
-    stacks = (x1, xs...)
-    _check_same_names(stacks...)
-    vals = map(f, map(values, stacks)...)
-    return _maybestack(_firststack(stacks...), vals)
+    xs = (x1, xs...)
+    firststack = _firststack(xs...)
+    if isnothing(firststack) 
+        if all(map(x -> x isa AbstractArray, xs))
+            # all arguments are arrays, but none are stacks, just apply the function
+            f(xs...)
+        else
+            # Arguments are some mix of arrays and NamedTuple
+            throw(ArgumentError("Cannot apply maplayers to NamedTuple and AbstractBasicDimArray"))
+        end
+    else
+        # There is at least one stack, we apply layer-wise
+        _check_same_names(xs...)
+        l = length(values(firststack))
+        vals = map(f, map(s -> _values_or_tuple(s, l), xs)...)
+        return _maybestack(firststack, vals)
+    end
 end
 
+_values_or_tuple(x::Union{AbstractDimStack, NamedTuple}, l) = values(x)
+_values_or_tuple(x::Union{AbstractBasicDimArray}, l) = Tuple(Iterators.repeated(x, l))
+
 # Other interfaces
 
 Extents.extent(A::AbstractDimStack, args...) = Extents.extent(dims(A), args...)
@@ -264,6 +280,9 @@ _check_same_names(::Union{AbstractDimStack{names},NamedTuple{names}},
     ::Union{AbstractDimStack{names},NamedTuple{names}}...) where {names} = nothing
 _check_same_names(::Union{AbstractDimStack,NamedTuple}, ::Union{AbstractDimStack,NamedTuple}...) =
     throw(ArgumentError("Named tuple names do not match."))
+_check_same_names(xs::Union{AbstractDimStack,NamedTuple,AbstractBasicDimArray}...) = 
+    _check_same_names((x for x in xs if x isa Union{AbstractDimStack,NamedTuple})...)
+
 
 _firststack(s::AbstractDimStack, args...) = s
 _firststack(arg1, args...) = _firststack(args...)

From f1e7b2ee58bf7bea2619e9934b75e039fef9cbc9 Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Mon, 20 Jan 2025 17:32:00 +0100
Subject: [PATCH 13/14] add bylayer to broadcast_dims

---
 src/utils.jl  | 39 ++++++++++++++++++++++-----------------
 test/utils.jl | 14 +++++++-------
 2 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/src/utils.jl b/src/utils.jl
index cc3240668..cbdd8a623 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -120,15 +120,16 @@ all passed in arrays in the order in which they are found.
 This is like broadcasting over every slice of `A` if it is
 sliced by the dimensions of `B`.
 """
-function broadcast_dims(f, As::AbstractBasicDimArray...)
-    dims = combinedims(As...)
-    T = Base.Broadcast.combine_eltypes(f, As)
-    broadcast_dims!(f, similar(first(As), T, dims), As...)
-end
-function broadcast_dims(f, s::AbstractDimStack, As::AbstractBasicDimArray...)
-    dims = combinedims(s, As...)
-    T = Base.Broadcast.combine_eltypes(f, (s, As...))
-    broadcast_dims!(f, similar(first(layers(s)), T, dims), s, As...)
+function broadcast_dims(f, As::AbstractBasicDimArray...; bylayer=false)
+    if bylayer
+        maplayers((As...) -> broadcast_dims(f, As...), As...)
+    else
+        A1 = first(As) 
+        _A1 = A1 isa AbstractDimStack ? first(layers(A1)) : A1
+        dims = combinedims(As...)
+        T = Base.Broadcast.combine_eltypes(f, As)
+        broadcast_dims!(f, similar(_A1, T, dims), As...)
+    end
 end
 
 """
@@ -145,15 +146,19 @@ which they are found.
 - `dest`: `AbstractDimArray` to update.
 - `sources`: `AbstractDimArrays` to broadcast over with `f`.
 """
-function broadcast_dims!(f, dest::Union{AbstractDimArray{<:Any,N}, AbstractDimStack{<:Any,<:Any,N}}, As::AbstractBasicDimArray...) where {N}
-    As = map(As) do A
-        isempty(otherdims(A, dims(dest))) || throw(DimensionMismatch("Cannot broadcast over dimensions not in the dest array"))
-        # comparedims(dest, dims(A, dims(dest)))
-        # Lazily permute B dims to match the order in A, if required
-        _maybe_lazy_permute(A, dims(dest))
+function broadcast_dims!(f, dest::Union{AbstractDimArray{<:Any,N}, AbstractDimStack{<:Any,<:Any,N}}, As::AbstractBasicDimArray...; bylayer=false) where {N}
+    if bylayer
+        maplayers((dest,As) -> broadcast_dims!(f, As, dest), As, dest)
+    else
+        As = map(As) do A
+            isempty(otherdims(A, dims(dest))) || throw(DimensionMismatch("Cannot broadcast over dimensions not in the dest array"))
+            # comparedims(dest, dims(A, dims(dest)))
+            # Lazily permute B dims to match the order in A, if required
+            _maybe_lazy_permute(A, dims(dest))
+        end
+        od = map(A -> otherdims(dest, dims(A)), As)
+        return _broadcast_dims_inner!(f, dest, As, od)
     end
-    od = map(A -> otherdims(dest, dims(A)), As)
-    return _broadcast_dims_inner!(f, dest, As, od)
 end
 
 # Function barrier
diff --git a/test/utils.jl b/test/utils.jl
index 96d145ad1..ca926380d 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -179,14 +179,14 @@ end
         db1 = DimArray(B1, (Y([:a, :b, :c]),))
         stack1 = DimStack(da3, db1)
         stack2 = DimStack(da3, db1, dc3)
-        @test broadcast_dims((s, x1, x2) -> s.layer1, stack1, da3, db1) == 
-            broadcast_dims((x1, x2, s) -> s.layer1, da3, db1, stack1)
-        @test broadcast_dims((s, x1, x2) -> s.layer1 + x1 + x2, stack1, da3, db1) == 
-            broadcast_dims(+, stack1.layer1, da3, db1)
-        @test broadcast_dims((s1, x1, s2, x2) -> s1.layer1+x2, stack1, da3, stack1, db1) == 
-            broadcast_dims((x1, s1, x2, s2) -> s1.layer1+x2, da3, stack1, db1, stack1)
+        # currently == returns false because the order of the dimensions is different!
+        @test all(broadcast_dims(+, stack1, da3, stack1, db1; bylayer = true) .== broadcast_dims(+, da3, stack1, db1, stack1; bylayer = true))
+        @test broadcast_dims(+, stack1, da3, db1; bylayer = true).layer1 == broadcast_dims(+, stack1.layer1, da3, db1)
+        @test broadcast_dims(+, stack1, da3, stack1, db1; bylayer = true) == broadcast_dims(+, da3, stack1, db1, stack1; bylayer = true)
         # Cant mix numvers of stack layers
-        @test_throws ArgumentError broadcast_dims(+, stack1, da3, db1, stack2)
+        @test_throws ArgumentError broadcast_dims(+, stack1, da3, db1, stack2; bylayer = true)
+        # If bylayer = false this is performed element-wise and throws a MethodError
+        @test_throws MethodError broadcast_dims(+, stack1, da3, db1, stack2)
     end
 end
 

From 780841d98da0bf0ab582c5022e809217d967346d Mon Sep 17 00:00:00 2001
From: tiemvanderdeure <tiemvanderdeure@gmail.com>
Date: Mon, 20 Jan 2025 17:56:04 +0100
Subject: [PATCH 14/14] fix and test setindex!

---
 src/stack/indexing.jl | 22 ++++++----------------
 test/utils.jl         |  9 ++++++---
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/src/stack/indexing.jl b/src/stack/indexing.jl
index bfa66a142..362c9076f 100644
--- a/src/stack/indexing.jl
+++ b/src/stack/indexing.jl
@@ -112,25 +112,15 @@ end
 #### setindex ####
 @propagate_inbounds Base.setindex!(s::AbstractDimStack, xs, I...; kw...) =
     map((A, x) -> setindex!(A, x, I...; kw...), layers(s), xs)
-@propagate_inbounds Base.setindex!(s::AbstractDimStack, xs::NamedTuple, i::Integer; kw...) =
-    hassamedims(s) ? _map_setindex!(s, xs, i; kw...) : _setindex_mixed!(s, xs, i; kw...)
-@propagate_inbounds Base.setindex!(s::AbstractDimStack, xs::NamedTuple, i::Colon; kw...) =
-    hassamedims(s) ? _map_setindex!(s, xs, i; kw...) : _setindex_mixed!(s, xs, i; kw...)
-@propagate_inbounds Base.setindex!(s::AbstractDimStack, xs::NamedTuple, i::AbstractArray; kw...) =
-    hassamedims(s) ? _map_setindex!(s, xs, i; kw...) : _setindex_mixed!(s, xs, i; kw...)
-
-@propagate_inbounds function Base.setindex!(
-    s::AbstractDimStack, xs::NamedTuple, I...; kw...
-)
-    map((A, x) -> setindex!(A, x, I...; kw...), layers(s), xs)
-end
+@propagate_inbounds Base.setindex!(s::AbstractDimStack, xs::NamedTuple, i...; kw...) =
+    hassamedims(s) ? _map_setindex!(s, xs, i...; kw...) : _setindex_mixed!(s, xs, i...; kw...)
 
 _map_setindex!(s, xs, i; kw...) = map((A, x) -> setindex!(A, x, i...; kw...), layers(s), xs)
 
-_setindex_mixed!(s::AbstractDimStack, x, i::AbstractArray) =
-    map(A -> setindex!(A, x, DimIndices(dims(s))[i]), layers(s))
-_setindex_mixed!(s::AbstractDimStack, i::Integer) =
-    map(A -> setindex!(A, x, DimIndices(dims(s))[i]), layers(s))
+_setindex_mixed!(s::AbstractDimStack, x, i...; kw...) =
+    map(layers(s), x) do A, x
+        A[dims(DimIndices(s)[i...], dims(A))] = x
+    end
 function _setindex_mixed!(s::AbstractDimStack, x, i::Colon)
     map(DimIndices(dims(s))) do D
         map(A -> setindex!(A, D), x, layers(s))
diff --git a/test/utils.jl b/test/utils.jl
index ca926380d..80d90fb6a 100644
--- a/test/utils.jl
+++ b/test/utils.jl
@@ -177,16 +177,19 @@ end
         A3 = cat([1 2 3; 4 5 6], [11 12 13; 14 15 16]; dims=3)
         da3 = DimArray(A3, (X([20, 30]), Y([:a, :b, :c]), Z(10:10:20)))
         db1 = DimArray(B1, (Y([:a, :b, :c]),))
-        stack1 = DimStack(da3, db1)
+        stack1 = DimStack(copy(da3), db1)
         stack2 = DimStack(da3, db1, dc3)
+        @test broadcast_dims(+, stack1, da3, db1; bylayer = true).layer1 == broadcast_dims(+, stack1.layer1, da3, db1)
         # currently == returns false because the order of the dimensions is different!
         @test all(broadcast_dims(+, stack1, da3, stack1, db1; bylayer = true) .== broadcast_dims(+, da3, stack1, db1, stack1; bylayer = true))
-        @test broadcast_dims(+, stack1, da3, db1; bylayer = true).layer1 == broadcast_dims(+, stack1.layer1, da3, db1)
-        @test broadcast_dims(+, stack1, da3, stack1, db1; bylayer = true) == broadcast_dims(+, da3, stack1, db1, stack1; bylayer = true)
         # Cant mix numvers of stack layers
         @test_throws ArgumentError broadcast_dims(+, stack1, da3, db1, stack2; bylayer = true)
         # If bylayer = false this is performed element-wise and throws a MethodError
         @test_throws MethodError broadcast_dims(+, stack1, da3, db1, stack2)
+
+        stack1.layer1 .+= 1
+        broadcast_dims!(nt -> nt[(:layer1, :layer2)], stack1, stack2)
+        @test stack1.layer1 == stack2.layer1
     end
 end