setup recursive broadcast overloads #10

SciML · Nov 26, 2017 · 0a598dc · 0a598dc
1 parent a4cd668
commit 0a598dc
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/src/math.jl b/src/math.jl
@@ -22,6 +22,25 @@ Base.Broadcast.broadcast_indices(::Type{<:AMSA}, A) = indices(A)
     broadcast!(f, similar(A), A, Bs...)
 end
 
+@inline function Base.Broadcast.broadcast_c(f, ::Type{S}, A::AMSA, Bs::AMSA...) where S<:AMSA
+    new_A = similar(A)
+    broadcast!(f,new_A,A,Bs...)
+    new_A
+end
+
+@inline function Base.Broadcast.broadcast_c!(f, ::Type{S}, ::Type, A::AbstractMultiScaleArrayLeaf, Bs::AbstractMultiScaleArrayLeaf...) where S<:AbstractMultiScaleArrayLeaf
+    broadcast!(f, A.values, (B.values for B in Bs)...)
+    A
+end
+
+@inline function Base.Broadcast.broadcast_c!(f, ::Type{S}, ::Type, A::AMSA, Bs::AMSA...) where S<:AMSA
+    for i in eachindex(A.nodes)
+            broadcast!(f, A.nodes[i], (B.nodes[i] for B in Bs)...)
+    end
+    broadcast!(f, A.values, (B.values for B in Bs)...)
+    A
+end
+
 +(m::AbstractMultiScaleArray, y::AbstractMultiScaleArray) = m .+ y
 +(m::AbstractMultiScaleArray, y::Number) = m .+ y
 +(y::Number, m::AbstractMultiScaleArray) = m .+ y

diff --git a/test/indexing_and_creation_tests.jl b/test/indexing_and_creation_tests.jl
@@ -122,8 +122,6 @@ end
 
 ### Test math
 
-#broadcast_getindex(cell1, 1)
-
 g = (x, y) -> x*y
 @test g.(cell1, 2) == [2.; 4; 6]
 # cell1 .= g.(cell1, 2) How to broadcast right???