almost have it

Author: ChrisRackauckas

src/math.jl

@@ -6,15 +6,20 @@ Base.map!(f::F, m::AMSA, A0::AbstractArray, As::AbstractArray...) where {F} =
 Base.map!(f::F, m::AMSA, A0, As...) where {F} =
         broadcast!(f, m, A0, As...)
 
-struct AMSAStyle{Style <: Broadcast.BroadcastStyle} <: Broadcast.AbstractArrayStyle{Any} end
+Base.BroadcastStyle(::Type{<:AMSA}) = Broadcast.ArrayStyle{AMSA}()
+Base.BroadcastStyle(::Type{<:AbstractMultiScaleArrayLeaf}) = Broadcast.ArrayStyle{AbstractMultiScaleArrayLeaf}()
+
+#=
 AMSAStyle(::S) where {S} = AMSAStyle{S}()
 AMSAStyle(::S, ::Val{N}) where {S,N} = AMSAStyle(S(Val(N)))
 AMSAStyle(::Val{N}) where N = AMSAStyle{Broadcast.DefaultArrayStyle{N}}()
 
+
 # promotion rules
 function Broadcast.BroadcastStyle(::AMSAStyle{AStyle}, ::AMSAStyle{BStyle}) where {AStyle, BStyle}
     AMSAStyle(Broadcast.BroadcastStyle(AStyle(), BStyle()))
 end
+=#
 
 #=
 combine_styles(args::Tuple{})         = Broadcast.DefaultArrayStyle{0}()
@@ -28,13 +33,23 @@ function Broadcast.BroadcastStyle(::Type{AMSA{T}}) where {T}
 end
 =#
 
-@inline function Base.copy(bc::Broadcast.Broadcasted{AMSAStyle{Style}}) where Style
-    nnodes(bc)
+@inline function Base.copy(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AMSA}})
+    N = nnodes(bc)
     @inline function f(i)
         copy(unpack(bc, i))
     end
-    CommonType = get_common_type(bc)
-    construct(CommonType, map(f,N), f(nothing))
+    first_amsa = find_amsa(bc)
+    construct(first_amsa, map(f,N), f(nothing))
+end
+
+@inline function Base.copy(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AbstractMultiScaleArrayLeaf}})
+    @show bc
+    @inline function f(i)
+        copy(unpack(bc, i))
+    end
+    first_amsa = find_amsa(bc)
+    @show first_amsa
+    construct(first_amsa, f(nothing))
 end
 
 @inline function Base.copyto!(dest::AMSA, bc::Broadcast.Broadcasted{Nothing})
@@ -51,7 +66,7 @@ end
 
 # drop axes because it is easier to recompute
 @inline unpack(bc::Broadcast.Broadcasted{Style}, i) where Style = Broadcast.Broadcasted{Style}(bc.f, unpack_args(i, bc.args))
-@inline unpack(bc::Broadcast.Broadcasted{AMSAStyle{Style}}, i) where Style = Broadcast.Broadcasted{Style}(bc.f, unpack_args(i, bc.args))
+@inline unpack(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AMSA}}, i) = Broadcast.Broadcasted{Broadcast.ArrayStyle{AMSA}}(bc.f, unpack_args(i, bc.args))
 unpack(x,::Any) = x
 unpack(x::AMSA, i) = x.nodes[i]
 unpack(x::AMSA, ::Nothing) = x.values
@@ -69,23 +84,16 @@ nnodes(A, Bs...) = common_number(nnodes(A), _nnodes(Bs))
 _nnodes(args::Tuple{Any}) = nnodes(args[1])
 _nnodes(args::Tuple{}) = 0
 
-get_common_type(A) = Nothing
-get_common_type(A::AMSA) = typeof(A)
-get_common_type(bc::Broadcast.Broadcasted) = _nnodes(bc.args)
-get_common_type(A, Bs...) = common_type(get_common_type(A), _get_common_type(Bs))
-
-@inline _get_common_type(args::Tuple) = get_common_type(get_common_type(args[1]), _get_common_type(Base.tail(args)))
-_get_common_type(args::Tuple{Any}) = get_common_type(args[1])
-_get_common_type(args::Tuple{}) = Nothing
+"`A = find_amsa(As)` returns the first AMSA among the arguments."
+find_amsa(bc::Base.Broadcast.Broadcasted) = find_amsa(bc.args)
+find_amsa(args::Tuple) = find_amsa(find_amsa(args[1]), Base.tail(args))
+find_amsa(x) = x
+find_amsa(a::AMSA, rest) = a
+find_amsa(::Any, rest) = find_amsa(rest)
 
 ## utils
 common_number(a, b) =
     a == 0 ? b :
     (b == 0 ? a :
      (a == b ? a :
       throw(DimensionMismatch("number of nodes must be equal"))))
-
-common_type(a::T, b::T) where T = T
-common_type(a::T, b::Nothing) where T = T
-common_type(a::Nothing, b::T) where T = T
-common_type(a::Nothing, b::Nothing) where T = Nothing

test/indexing_and_creation_tests.jl

@@ -28,6 +28,8 @@ end
 cell1 = Cell([1.0; 2.0; 3.0])
 cell2 = Cell([4.0; 5])
 
+cell1 .+ cell1
+
 sim_cell = similar(cell1)
 
 @test length(cell1) == 3