Support Diagonal matrices of Fills (#98)

* Support diagonal matrix of eltype Fill or Zeros

* v0.8.9

Author: dlfivefifty

Project.toml

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

src/FillArrays.jl

@@ -7,7 +7,7 @@ import Base: size, getindex, setindex!, IndexStyle, checkbounds, convert,
     copy, vec, setindex!, count, ==, reshape, _throw_dmrs, map, zero
 
 import LinearAlgebra: rank, svdvals!, tril, triu, tril!, triu!, diag, transpose, adjoint, fill!,
-    norm2, norm1, normInf, normMinusInf, normp, lmul!, rmul!
+    norm2, norm1, normInf, normMinusInf, normp, lmul!, rmul!, diagzero
 
 import Base.Broadcast: broadcasted, DefaultArrayStyle, broadcast_shape
 
@@ -186,7 +186,9 @@ for (Typ, funcs, func) in ((:Zeros, :zeros, :zero), (:Ones, :ones, :one))
     @eval begin
         struct $Typ{T, N, Axes} <: AbstractFill{T, N, Axes}
             axes::Axes
-            @inline $Typ{T, N}(sz::Axes) where Axes<:Tuple{Vararg{AbstractUnitRange,N}} where {T, N} =
+            @inline $Typ{T,N,Axes}(sz::Axes) where Axes<:Tuple{Vararg{AbstractUnitRange,N}} where {T, N} =
+                new{T,N,Axes}(sz)
+            @inline $Typ{T,N}(sz::Axes) where Axes<:Tuple{Vararg{AbstractUnitRange,N}} where {T, N} =
                 new{T,N,Axes}(sz)
             @inline $Typ{T,0,Tuple{}}(sz::Tuple{}) where T = new{T,0,Tuple{}}(sz)
         end

src/fillalgebra.jl

@@ -231,4 +231,10 @@ function rmul!(z::AbstractFill, x::Number)
     # for, e.g., lmul!(NaN, z)
     λ*x == λ || throw(ArgumentError("Cannot scale by $x"))
     z
-end
+end
+
+fillzero(::Type{Fill{T,N,AXIS}}, n, m) where {T,N,AXIS} = Fill{T,N,AXIS}(zero(T), (n, m))
+fillzero(::Type{Zeros{T,N,AXIS}}, n, m) where {T,N,AXIS} = Zeros{T,N,AXIS}((n, m))
+fillzero(::Type{F}, n, m) where F = throw(ArgumentError("Cannot create a zero array of type $F"))
+
+diagzero(D::Diagonal{F}, i, j) where F<:AbstractFill = fillzero(F, axes(D.diag[i], 1), axes(D.diag[j], 2))

test/runtests.jl

@@ -965,4 +965,19 @@ end
     @test rmul!(x,1.0) === x
     @test_throws ArgumentError lmul!(2.0,x)
     @test_throws ArgumentError rmul!(x,2.0)
+end
+
+@testset "Diagonal{<:Fill}" begin
+    D = Diagonal(Fill(Fill(0.5,2,2),10))
+    @test @inferred(D[1,1]) === Fill(0.5,2,2)
+    @test @inferred(D[1,2]) === Fill(0.0,2,2)
+    D = Diagonal(Fill(Zeros(2,2),10))
+    @test @inferred(D[1,1]) === Zeros(2,2)
+    @test @inferred(D[1,2]) === Zeros(2,2)
+
+    D = Diagonal([Zeros(1,1), Zeros(2,2)])
+    @test @inferred(D[1,1]) === Zeros(1,1)
+    @test @inferred(D[1,2]) === Zeros(1,2)
+
+    @test_throws ArgumentError Diagonal(Fill(Ones(2,2),10))[1,2]
 end