Add special broadcast routines for +

Project.toml

@@ -1,6 +1,6 @@
 name = "FillArrays"
 uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
-version = "0.9.7"
+version = "0.10"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/FillArrays.jl

@@ -158,16 +158,18 @@ convert(::Type{T}, F::T) where T<:Fill = F   # ambiguity fix
 
 getindex(F::Fill{<:Any,0}) = getindex_value(F)
 
-Base._unsafe_getindex(::IndexStyle, F::Fill, kj::Vararg{AbstractVector{II},N}) where {II<:Integer,N} =
-    Fill(getindex_value(F), length.(kj))
+function Base._unsafe_getindex(::IndexStyle, F::AbstractFill, I::Vararg{Union{Real, AbstractArray}, N}) where N
+    shape = Base.index_shape(I...)
+    fillsimilar(F, shape)
+end
 
-function getindex(A::Fill, kr::AbstractVector{Bool})
+function getindex(A::AbstractFill, kr::AbstractVector{Bool})
    length(A) == length(kr) || throw(DimensionMismatch())
-   Fill(getindex_value(A), count(kr))
+   fillsimilar(A, count(kr))
 end
-function getindex(A::Fill, kr::AbstractArray{Bool})
+function getindex(A::AbstractFill, kr::AbstractArray{Bool})
    size(A) == size(kr) || throw(DimensionMismatch())
-   Fill(getindex_value(A), count(kr))
+   fillsimilar(A, count(kr))
 end
 
 sort(a::AbstractFill; kwds...) = a
@@ -202,7 +204,7 @@ end
 function fill_reshape(parent, dims::Integer...)
     n = length(parent)
     prod(dims) == n || _throw_dmrs(n, "size", dims)
-    Fill(getindex_value(parent), dims...)
+    fillsimilar(parent, dims...)
 end
 
 reshape(parent::AbstractFill, dims::Integer...) = reshape(parent, dims)
@@ -267,26 +269,21 @@ for (Typ, funcs, func) in ((:Zeros, :zeros, :zero), (:Ones, :ones, :one))
         copy(F::$Typ) = F
 
         getindex(F::$Typ{T,0}) where T = getindex_value(F)
-        Base._unsafe_getindex(::IndexStyle, F::$Typ{T}, kj::Vararg{AbstractVector{II},N}) where {T,II<:Integer,N} =
-            $Typ{T}(length.(kj))
-        function getindex(A::$Typ{T}, kr::AbstractVector{Bool}) where T
-            length(A) == length(kr) || throw(DimensionMismatch("lengths must match"))
-            $Typ{T}(count(kr))
-        end
-        function getindex(A::$Typ{T}, kr::AbstractArray{Bool}) where T
-            size(A) == size(kr) || throw(DimensionMismatch("sizes must match"))
-            $Typ{T}(count(kr))
-        end
-
-        function fill_reshape(parent::$Typ{T}, dims::Integer...) where T
-            n = length(parent)
-            prod(dims) == n || _throw_dmrs(n, "size", dims)
-            $Typ{T}(dims...)
-        end
     end
 end
 
 
+"""
+    fillsimilar(a::AbstractFill, axes)
+
+creates a fill object that has the same fill value as `a` but
+with the specified axes.
+For example, if `a isa Zeros` then so is the returned object.
+"""
+fillsimilar(a::Ones{T}, axes...) where T = Ones{T}(axes...)
+fillsimilar(a::Zeros{T}, axes...) where T = Zeros{T}(axes...)
+fillsimilar(a::AbstractFill, axes...) = Fill(getindex_value(a), axes...)
+
 
 rank(F::Zeros) = 0
 rank(F::Ones) = 1

src/fillalgebra.jl

@@ -15,10 +15,6 @@ adjoint(a::Zeros{T,2}) where T = Zeros{T}(reverse(a.axes))
 transpose(a::Fill{T,2}) where T = Fill{T}(transpose(a.value), reverse(a.axes))
 adjoint(a::Fill{T,2}) where T = Fill{T}(adjoint(a.value), reverse(a.axes))
 
-fillsimilar(a::Ones{T}, axes) where T = Ones{T}(axes)
-fillsimilar(a::Zeros{T}, axes) where T = Zeros{T}(axes)
-fillsimilar(a::AbstractFill, axes) = Fill(getindex_value(a), axes)
-
 permutedims(a::AbstractFill{<:Any,1}) = fillsimilar(a, (1, length(a)))
 permutedims(a::AbstractFill{<:Any,2}) = fillsimilar(a, reverse(a.axes))
 

src/fillbroadcast.jl

@@ -44,6 +44,14 @@ broadcasted(::DefaultArrayStyle, ::typeof(-), a::Zeros, b::Zeros) = _broadcasted
 broadcasted(::DefaultArrayStyle, ::typeof(-), a::Ones, b::Zeros) = _broadcasted_ones(-, a, b)
 broadcasted(::DefaultArrayStyle, ::typeof(-), a::Ones, b::Ones) = _broadcasted_zeros(-, a, b)
 
+broadcasted(::DefaultArrayStyle{1}, ::typeof(+), a::Zeros{<:Any,1}, b::Zeros{<:Any,1}) = _broadcasted_zeros(+, a, b)
+broadcasted(::DefaultArrayStyle{1}, ::typeof(+), a::Ones{<:Any,1}, b::Zeros{<:Any,1}) = _broadcasted_ones(+, a, b)
+broadcasted(::DefaultArrayStyle{1}, ::typeof(+), a::Zeros{<:Any,1}, b::Ones{<:Any,1}) = _broadcasted_ones(+, a, b)
+
+broadcasted(::DefaultArrayStyle{1}, ::typeof(-), a::Zeros{<:Any,1}, b::Zeros{<:Any,1}) = _broadcasted_zeros(-, a, b)
+broadcasted(::DefaultArrayStyle{1}, ::typeof(-), a::Ones{<:Any,1}, b::Zeros{<:Any,1}) = _broadcasted_ones(-, a, b)
+
+
 broadcasted(::DefaultArrayStyle, ::typeof(*), a::Zeros, b::Zeros) = _broadcasted_zeros(*, a, b)
 
 # In following, need to restrict to <: Number as otherwise we cannot infer zero from type
@@ -109,16 +117,36 @@ _range_convert(::Type{AbstractVector{T}}, a::AbstractRange) where T = convert(T,
 #     end
 # end
 
-function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::Ones{T}, b::AbstractRange{V}) where {T,V}
+function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::Ones{T,1}, b::AbstractRange{V}) where {T,V}
     broadcast_shape(axes(a), axes(b)) == axes(b) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $b to a Vector first."))
     return _range_convert(AbstractVector{promote_type(T,V)}, b)
 end
 
-function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::AbstractRange{V}, b::Ones{T}) where {T,V}
+function broadcasted(::DefaultArrayStyle{1}, ::typeof(*), a::AbstractRange{V}, b::Ones{T,1}) where {T,V}
     broadcast_shape(axes(a), axes(b)) == axes(a) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $b to a Vector first."))
     return _range_convert(AbstractVector{promote_type(T,V)}, a)
 end
 
+for op in (:+, -)
+    @eval begin
+        function broadcasted(::DefaultArrayStyle{1}, ::typeof($op), a::AbstractVector{T}, b::Zeros{V,1}) where {T,V}
+            broadcast_shape(axes(a), axes(b)) == axes(a) || throw(ArgumentError("Cannot broadcast $a and $b. Convert $b to a Vector first."))
+            LinearAlgebra.copy_oftype(a, promote_type(T,V))
+        end
+
+        broadcasted(::DefaultArrayStyle{1}, ::typeof($op), a::AbstractFill{T,1}, b::Zeros{V,1}) where {T,V} = 
+            Base.invoke(broadcasted, Tuple{DefaultArrayStyle, typeof($op), AbstractFill, AbstractFill}, DefaultArrayStyle{1}(), $op, a, b)
+    end
+end
+
+function broadcasted(::DefaultArrayStyle{1}, ::typeof(+), a::Zeros{T,1}, b::AbstractVector{V}) where {T,V}
+    broadcast_shape(axes(a), axes(b))
+    LinearAlgebra.copy_oftype(b, promote_type(T,V))
+end
+
+broadcasted(::DefaultArrayStyle{1}, ::typeof(+), a::Zeros{V,1}, b::AbstractFill{T,1}) where {T,V} = 
+            Base.invoke(broadcasted, Tuple{DefaultArrayStyle, typeof(+), AbstractFill, AbstractFill}, DefaultArrayStyle{1}(), +, a, b)
+
 # Need to prevent array-valued fills from broadcasting over entry
 _broadcast_getindex_value(a::AbstractFill{<:Number}) = getindex_value(a)
 _broadcast_getindex_value(a::AbstractFill) = Ref(getindex_value(a))

test/runtests.jl

@@ -165,6 +165,70 @@ import FillArrays: AbstractFill, RectDiagonal, SquareEye
     end
 end
 
+@testset "indexing" begin
+    A = Fill(3.0,5)
+    @test A[1:3] ≡ Fill(3.0,3)
+    @test A[1:3,1:1] ≡ Fill(3.0,3,1)
+    @test_throws BoundsError A[1:3,2]
+    @test_throws BoundsError A[1:26]
+    @test A[[true, false, true, false, false]] ≡ Fill(3.0, 2)
+    A = Fill(3.0, 2, 2)
+    @test A[[true true; true false]] ≡ Fill(3.0, 3)
+    @test_throws DimensionMismatch A[[true, false]]
+
+    A = Ones{Int}(5,5)
+    @test A[1:3] ≡ Ones{Int}(3)
+    @test A[1:3,1:2] ≡ Ones{Int}(3,2)
+    @test A[1:3,2] ≡ Ones{Int}(3)
+    @test_throws BoundsError A[1:26]
+    A = Ones{Int}(2,2)
+    @test A[[true false; true false]] ≡ Ones{Int}(2)
+    @test A[[true, false, true, false]] ≡ Ones{Int}(2)
+    @test_throws DimensionMismatch A[[true false false; true false false]]
+
+    A = Zeros{Int}(5,5)
+    @test A[1:3] ≡ Zeros{Int}(3)
+    @test A[1:3,1:2] ≡ Zeros{Int}(3,2)
+    @test A[1:3,2] ≡ Zeros{Int}(3)
+    @test_throws BoundsError A[1:26]
+    A = Zeros{Int}(2,2)
+    @test A[[true false; true false]] ≡ Zeros{Int}(2)
+    @test A[[true, false, true, false]] ≡ Zeros{Int}(2)
+    @test_throws DimensionMismatch A[[true false false; true false false]]
+
+    @testset "colon" begin
+        @test Ones(2)[:] ≡ Ones(2)[Base.Slice(Base.OneTo(2))] ≡ Ones(2)
+        @test Zeros(2)[:] ≡ Zeros(2)[Base.Slice(Base.OneTo(2))] ≡ Zeros(2)
+        @test Fill(3.0,2)[:] ≡ Fill(3.0,2)[Base.Slice(Base.OneTo(2))] ≡ Fill(3.0,2)
+
+        @test Ones(2,2)[:,:] ≡ Ones(2,2)[Base.Slice(Base.OneTo(2)),Base.Slice(Base.OneTo(2))] ≡ Ones(2,2)
+        @test Zeros(2,2)[:,:] ≡ Zeros(2)[Base.Slice(Base.OneTo(2)),Base.Slice(Base.OneTo(2))] ≡ Zeros(2,2)
+        @test Fill(3.0,2,2)[:,:] ≡ Fill(3.0,2,2)[Base.Slice(Base.OneTo(2)),Base.Slice(Base.OneTo(2))] ≡ Fill(3.0,2,2)
+    end
+
+    @testset "mixed integer / vector /colon" begin
+        a = Fill(2.0,5)
+        z = Zeros(5)
+        @test a[1:5] ≡ a[:] ≡ a
+        @test z[1:5] ≡ z[:] ≡ z
+
+        A = Fill(2.0,5,6)
+        Z = Zeros(5,6)
+        @test A[:,1] ≡ A[1:5,1] ≡ Fill(2.0,5)
+        @test A[1,:] ≡ A[1,1:6] ≡ Fill(2.0,6)
+        @test A[:,:] ≡ A[1:5,1:6] ≡ A[1:5,:] ≡ A[:,1:6] ≡ A
+        @test Z[:,1] ≡ Z[1:5,1] ≡ Zeros(5)
+        @test Z[1,:] ≡ Z[1,1:6] ≡ Zeros(6)
+        @test Z[:,:] ≡ Z[1:5,1:6] ≡ Z[1:5,:] ≡ Z[:,1:6] ≡ Z
+
+        A = Fill(2.0,5,6,7)
+        Z = Zeros(5,6,7)
+        @test A[:,1,1] ≡ A[1:5,1,1] ≡ Fill(2.0,5)
+        @test A[1,:,1] ≡ A[1,1:6,1] ≡ Fill(2.0,6)
+        @test A[:,:,:] ≡ A[1:5,1:6,1:7] ≡ A[1:5,:,1:7] ≡ A[:,1:6,1:7] ≡ A
+    end
+end
+
 @testset "RectDiagonal" begin
     data = 1:3
     expected_size = (5, 3)
@@ -542,8 +606,8 @@ end
     @testset "range broadcast" begin
         rnge = range(-5.0, step=1.0, length=10)
         @test broadcast(*, Fill(5.0, 10), rnge) == broadcast(*, 5.0, rnge)
-        @test broadcast(*, Zeros(10, 10), rnge) == zeros(10, 10)
-        @test broadcast(*, rnge, Zeros(10, 10)) == zeros(10, 10)
+        @test broadcast(*, Zeros(10, 10), rnge) ≡ Zeros{Float64}(10, 10)
+        @test broadcast(*, rnge, Zeros(10, 10)) ≡ Zeros{Float64}(10, 10)
         @test broadcast(*, Ones{Int}(10), rnge) ≡ rnge
         @test broadcast(*, rnge, Ones{Int}(10)) ≡ rnge
         @test_throws DimensionMismatch broadcast(*, Fill(5.0, 11), rnge)
@@ -554,6 +618,12 @@ end
         deg = 5:5
         @test_throws ArgumentError @inferred(broadcast(*, Fill(5.0, 10), deg)) == broadcast(*, fill(5.0,10), deg)
         @test_throws ArgumentError @inferred(broadcast(*, deg, Fill(5.0, 10))) == broadcast(*, deg, fill(5.0,10))
+
+        @test rnge .+ Zeros(10) ≡ rnge .- Zeros(10) ≡ Zeros(10) .+ rnge ≡ rnge
+
+        @test_throws DimensionMismatch rnge .+ Zeros(5)
+        @test_throws DimensionMismatch rnge .- Zeros(5)
+        @test_throws DimensionMismatch Zeros(5) .+ rnge
     end
 
     @testset "Special Zeros/Ones" begin
@@ -598,6 +668,15 @@ end
             @test Zeros(5) ./ Zeros(5) ≡ Zeros(5) .\ Zeros(5) ≡ Fill(NaN,5)
             @test Zeros{Int}(5,6) ./ Zeros{Int}(5) ≡ Zeros{Int}(5) .\ Zeros{Int}(5,6) ≡ Fill(NaN,5,6)
         end
+
+        @testset "Addition" begin
+            @test Zeros{Int}(5) .+ (1:5) ≡ (1:5) .+ Zeros{Int}(5) ≡ (1:5) .- Zeros{Int}(5) ≡ 1:5
+            @test Zeros{Int}(1) .+ (1:5) ≡ (1:5) .+ Zeros{Int}(1) ≡ (1:5) .- Zeros{Int}(1) ≡ 1:5
+            @test Zeros(5) .+ (1:5) == (1:5) .+ Zeros(5) == (1:5) .- Zeros(5) == 1:5
+            @test Zeros{Int}(5) .+ Fill(1,5) ≡ Fill(1,5) .+ Zeros{Int}(5) ≡ Fill(1,5) .- Zeros{Int}(5) ≡ Fill(1,5)
+            @test_throws DimensionMismatch Zeros{Int}(2) .+ (1:5)
+            @test_throws DimensionMismatch (1:5) .+ Zeros{Int}(2)
+        end
     end
 
     @testset "support Ref" begin
@@ -625,6 +704,11 @@ end
         @test Ones(10) - Ones(10) ≡ Zeros(10)
         @test Fill(1,10) - Zeros(10) ≡ Fill(1.0,10)
 
+        @test Zeros(10) .- Zeros(10) ≡ Zeros(10)
+        @test Ones(10) .- Zeros(10) ≡ Ones(10)
+        @test Ones(10) .- Ones(10) ≡ Zeros(10)
+        @test Fill(1,10) .- Zeros(10) ≡ Fill(1.0,10)
+
         @test Zeros(10) .- Zeros(1,9) ≡ Zeros(10,9)
         @test Ones(10) .- Zeros(1,9) ≡ Ones(10,9)
         @test Ones(10) .- Ones(1,9) ≡ Zeros(10,9)
@@ -652,48 +736,6 @@ end
     @test map(exp,x) === Fill(exp(2),5,3)
 end
 
-@testset "Sub-arrays" begin
-    A = Fill(3.0,5)
-    @test A[1:3] ≡ Fill(3.0,3)
-    @test A[1:3,1:1] ≡ Fill(3.0,3,1)
-    @test_broken A[1:3,2] ≡ Zeros{Int}(3)
-    @test_throws BoundsError A[1:26]
-    @test A[[true, false, true, false, false]] ≡ Fill(3.0, 2)
-    A = Fill(3.0, 2, 2)
-    @test A[[true true; true false]] ≡ Fill(3.0, 3)
-    @test_throws DimensionMismatch A[[true, false]]
-
-    A = Ones{Int}(5,5)
-    @test A[1:3] ≡ Ones{Int}(3)
-    @test A[1:3,1:2] ≡ Ones{Int}(3,2)
-    @test_broken A[1:3,2] ≡ Ones{Int}(3)
-    @test_throws BoundsError A[1:26]
-    A = Ones{Int}(2,2)
-    @test A[[true false; true false]] ≡ Ones{Int}(2)
-    @test A[[true, false, true, false]] ≡ Ones{Int}(2)
-    @test_throws DimensionMismatch A[[true false false; true false false]]
-
-    A = Zeros{Int}(5,5)
-    @test A[1:3] ≡ Zeros{Int}(3)
-    @test A[1:3,1:2] ≡ Zeros{Int}(3,2)
-    @test_broken A[1:3,2] ≡ Zeros{Int}(3)
-    @test_throws BoundsError A[1:26]
-    A = Zeros{Int}(2,2)
-    @test A[[true false; true false]] ≡ Zeros{Int}(2)
-    @test A[[true, false, true, false]] ≡ Zeros{Int}(2)
-    @test_throws DimensionMismatch A[[true false false; true false false]]
-
-    @testset "colon" begin
-        @test Ones(2)[:] ≡ Ones(2)[Base.Slice(Base.OneTo(2))] ≡ Ones(2)
-        @test Zeros(2)[:] ≡ Zeros(2)[Base.Slice(Base.OneTo(2))] ≡ Zeros(2)
-        @test Fill(3.0,2)[:] ≡ Fill(3.0,2)[Base.Slice(Base.OneTo(2))] ≡ Fill(3.0,2)
-
-        @test Ones(2,2)[:,:] ≡ Ones(2,2)[Base.Slice(Base.OneTo(2)),Base.Slice(Base.OneTo(2))] ≡ Ones(2,2)
-        @test Zeros(2,2)[:,:] ≡ Zeros(2)[Base.Slice(Base.OneTo(2)),Base.Slice(Base.OneTo(2))] ≡ Zeros(2,2)
-        @test Fill(3.0,2,2)[:,:] ≡ Fill(3.0,2,2)[Base.Slice(Base.OneTo(2)),Base.Slice(Base.OneTo(2))] ≡ Fill(3.0,2,2)
-    end
-end
-
 @testset "Offset indexing" begin
     A = Fill(3, (Base.Slice(-1:1),))
     @test axes(A)  == (Base.Slice(-1:1),)