Return OffDiagonal when slicing Diagonal of LayoutVector (require FillArrays v1.18)

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayLayouts"
 uuid = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
-version = "1.14"
+version = "1.15"
 authors = ["Sheehan Olver <solver@mac.com>"]
 
 [deps]
@@ -16,7 +16,7 @@ ArrayLayoutsSparseArraysExt = "SparseArrays"
 
 [compat]
 Aqua = "0.8"
-FillArrays = "1.2.1"
+FillArrays = "1.18"
 Infinities = "0.1"
 JLArrays = "0.2, 0.3"
 LinearAlgebra = "1"

src/ArrayLayouts.jl

@@ -23,15 +23,15 @@ using LinearAlgebra.BLAS: BlasComplex, BlasFloat, BlasReal
 
 const AdjointQtype{T} = LinearAlgebra.AdjointQ{T}
 
-using FillArrays: AbstractFill, axes_print_matrix_row, getindex_value
+using FillArrays: AbstractFill, OffDiagonal, axes_print_matrix_row, getindex_value
 using StaticArrays
 
 using Base: require_one_based_indexing
 
 export materialize, materialize!, MulAdd, muladd!, Ldiv, Rdiv, Lmul, Rmul, Dot,
         lmul, mul, ldiv, rdiv, mul, MemoryLayout, AbstractStridedLayout,
         DenseColumnMajor, ColumnMajor, ZerosLayout, FillLayout, AbstractColumnMajor, RowMajor, AbstractRowMajor, UnitStride,
-        DiagonalLayout, ScalarLayout, SymTridiagonalLayout, TridiagonalLayout, BidiagonalLayout,
+        DiagonalLayout, OffDiagonalLayout, ScalarLayout, SymTridiagonalLayout, TridiagonalLayout, BidiagonalLayout,
         HermitianLayout, SymmetricLayout, TriangularLayout,
         UnknownLayout, AbstractBandedLayout, ApplyBroadcastStyle, ConjLayout, AbstractFillLayout, DualLayout,
         colsupport, rowsupport, layout_getindex, AbstractQLayout, LayoutArray, LayoutMatrix, LayoutVector,
@@ -285,6 +285,8 @@ copymutable_oftype_layout(_, A, ::Type{S}) where S = copyto!(similar(A, S), A)
 Base.map(::typeof(copy), D::Diagonal{<:LayoutArray}) = Diagonal(map(copy, D.diag))
 Base.permutedims(D::Diagonal{<:Any,<:LayoutVector}) = D
 
+getindex(D::Diagonal{<:Any,<:Union{LayoutVector,SubArray{<:Any,1,<:LayoutVecOrMat}}}, kr::Union{AbstractUnitRange{<:Integer},Colon}, jr::Union{AbstractUnitRange{<:Integer},Colon}) = layout_getindex(D, kr, jr)
+
 
 zero!(A) = zero!(MemoryLayout(A), A)
 zero!(_, A) = fill!(A,zero(eltype(A)))

src/memorylayout.jl

@@ -550,6 +550,14 @@ abstract type AbstractBandedLayout <: MemoryLayout end
 abstract type AbstractTridiagonalLayout <: AbstractBandedLayout end
 
 struct DiagonalLayout{ML} <: AbstractBandedLayout end
+"""
+    OffDiagonalLayout{ML}()
+
+is returned by `MemoryLayout(A)` if a matrix `A` is supported on a single (possibly off-) diagonal,
+as with a `FillArrays.OffDiagonal`, whose entries have layout `ML`. The entries and the index of
+the diagonal are given by `offdiagonaldata(A)` and `offdiagonalindex(A)`.
+"""
+struct OffDiagonalLayout{ML} <: AbstractBandedLayout end
 struct BidiagonalLayout{DV,EV} <: AbstractBandedLayout end
 struct SymTridiagonalLayout{DV,EV} <: AbstractTridiagonalLayout end
 struct TridiagonalLayout{DL,D,DU} <: AbstractTridiagonalLayout end
@@ -562,9 +570,11 @@ bidiagonallayout(d) = bidiagonallayout(d, d)
 tridiagonallayout(d) = tridiagonallayout(d,d,d)
 symtridiagonallayout(d) = symtridiagonallayout(d,d)
 diagonallayout(_) = DiagonalLayout{UnknownLayout}()
+offdiagonallayout(_) = OffDiagonalLayout{UnknownLayout}()
 diagonal(d::AbstractVector) = Diagonal(d) # support non-array diagonal objects like QuasiDiagonal
 
 diagonallayout(::Lay) where Lay<:Union{AbstractStridedLayout, AbstractFillLayout} = DiagonalLayout{Lay}()
+offdiagonallayout(::Lay) where Lay<:Union{AbstractStridedLayout, AbstractFillLayout} = OffDiagonalLayout{Lay}()
 bidiagonallayout(::Lay, ::Lay) where Lay<:Union{AbstractStridedLayout, AbstractFillLayout} = BidiagonalLayout{Lay,Lay}()
 tridiagonallayout(::Lay, ::Lay, ::Lay) where Lay<:Union{AbstractStridedLayout, AbstractFillLayout} = TridiagonalLayout{Lay,Lay,Lay}()
 symtridiagonallayout(::Lay, ::Lay) where Lay<:Union{AbstractStridedLayout, AbstractFillLayout} = SymTridiagonalLayout{Lay,Lay}()
@@ -574,6 +584,7 @@ MemoryLayout(D::Type{Diagonal{T,P}}) where {T,P} = diagonallayout(MemoryLayout(P
 MemoryLayout(::Type{Bidiagonal{T,V}}) where {T,V} = bidiagonallayout(MemoryLayout(V))
 MemoryLayout(::Type{SymTridiagonal{T,P}}) where {T,P} = symtridiagonallayout(MemoryLayout(P))
 MemoryLayout(::Type{Tridiagonal{T,P}}) where {T,P} = tridiagonallayout(MemoryLayout(P))
+MemoryLayout(::Type{<:OffDiagonal{T,V}}) where {T,V} = offdiagonallayout(MemoryLayout(V))
 
 bidiagonaluplo(A::Bidiagonal) = A.uplo
 bidiagonaluplo(A::AdjOrTrans) = bidiagonaluplo(parent(A)) == 'L' ? 'U' : 'L'
@@ -592,7 +603,11 @@ subdiagonaldata(D::SymTridiagonal) = D.ev
 subdiagonaldata(D::Tridiagonal) = D.dl
 supdiagonaldata(D::Tridiagonal) = D.du
 
+offdiagonaldata(A::OffDiagonal) = A.band
+offdiagonalindex(A::OffDiagonal) = A.k
+
 transposelayout(ml::DiagonalLayout) = ml
+transposelayout(ml::OffDiagonalLayout) = ml
 transposelayout(ml::BidiagonalLayout) = ml
 transposelayout(ml::SymTridiagonalLayout) = ml
 transposelayout(ml::TridiagonalLayout) = ml
@@ -614,6 +629,7 @@ subdiagonaldata(U::Union{UnitLowerTriangular,LowerTriangular}) = subdiagonaldata
 adjointlayout(::Type{<:Real}, ml::SymTridiagonalLayout) = ml
 adjointlayout(::Type{<:Real}, ::TridiagonalLayout{DL,D,DU}) where {DL,D,DU} = TridiagonalLayout{DU,D,DL}()
 adjointlayout(::Type{<:Real}, ml::BidiagonalLayout) = ml
+adjointlayout(::Type{<:Real}, ml::OffDiagonalLayout) = ml
 
 symmetriclayout(B::BidiagonalLayout{DV,EV}) where {DV,EV} = SymTridiagonalLayout{DV,EV}()
 hermitianlayout(::Type{<:Real}, B::BidiagonalLayout{DV,EV}) where {DV,EV} = SymTridiagonalLayout{DV,EV}()
@@ -627,6 +643,46 @@ diagonaldata(D::Adjoint{<:Real}) = diagonaldata(parent(D))
 subdiagonaldata(D::Adjoint{<:Real}) = supdiagonaldata(parent(D))
 supdiagonaldata(D::Adjoint{<:Real}) = subdiagonaldata(parent(D))
 
+offdiagonaldata(D::Union{Transpose,Adjoint{<:Real}}) = offdiagonaldata(parent(D))
+offdiagonalindex(D::Union{Transpose,Adjoint{<:Real}}) = -offdiagonalindex(parent(D))
+
+# a matrix with DiagonalLayout is supported on its 0-th diagonal
+_singlebanddata(::DiagonalLayout, A) = diagonaldata(A)
+_singlebanddata(::OffDiagonalLayout, A) = offdiagonaldata(A)
+_singlebandindex(::DiagonalLayout, _) = 0
+_singlebandindex(::OffDiagonalLayout, A) = offdiagonalindex(A)
+
+# slicing by unit ranges keeps a single diagonal.
+# The index type is wider than AbstractUnitRange{Int} to avoid overwriting the sublayout(::DiagonalLayout, ...) in BandedMatrices.jl
+sublayout(::DiagonalLayout{ML}, ::Type{<:Tuple{AbstractUnitRange{<:Integer},AbstractUnitRange{<:Integer}}}) where ML = OffDiagonalLayout{ML}()
+sublayout(::OffDiagonalLayout{ML}, ::Type{<:Tuple{AbstractUnitRange{<:Integer},AbstractUnitRange{<:Integer}}}) where ML = OffDiagonalLayout{ML}()
+
+function offdiagonalindex(V::SubArray{<:Any,2})
+    kr, jr = parentindices(V)
+    P = parent(V)
+    _singlebandindex(MemoryLayout(P), P) + first(kr) - first(jr)
+end
+
+# the indices of the band of the parent that lie in V
+function _offdiagonalband_indices(V::SubArray{<:Any,2})
+    kr, jr = parentindices(V)
+    k = offdiagonalindex(V)
+    len = max(0, k ≥ 0 ? min(length(kr), length(jr) - k) : min(length(kr) + k, length(jr)))
+    # the band of the parent is indexed by min(row, col), so start from the first entry of the k-th diagonal of V
+    st = min(first(kr) + max(0, -k), first(jr) + max(0, k))
+    st:st+len-1
+end
+
+function offdiagonaldata(V::SubArray{<:Any,2})
+    P = parent(V)
+    view(_singlebanddata(MemoryLayout(P), P), _offdiagonalband_indices(V))
+end
+
+function sub_materialize(::OffDiagonalLayout, V::SubArray{<:Any,2}, ax::Tuple{OneTo{Int},OneTo{Int}})
+    P = parent(V)
+    OffDiagonal(_singlebanddata(MemoryLayout(P), P)[_offdiagonalband_indices(V)], offdiagonalindex(V), ax)
+end
+
 diagonaldata(S::HermOrSym) = diagonaldata(parent(S))
 subdiagonaldata(S::HermOrSym) = symmetricuplo(S) == 'L' ? subdiagonaldata(parent(S)) : supdiagonaldata(parent(S))
 supdiagonaldata(S::HermOrSym) = symmetricuplo(S) == 'L' ? subdiagonaldata(parent(S)) : supdiagonaldata(parent(S))
@@ -686,6 +742,17 @@ end
 rowsupport(::DiagonalLayout, _, k) = k
 colsupport(::DiagonalLayout, _, j) = j
 
+function colsupport(::OffDiagonalLayout, A, j)
+    isempty(j) && return 1:0
+    k = offdiagonalindex(A)
+    max(minimum(j)-k,1):min(size(A,1),maximum(j)-k)
+end
+function rowsupport(::OffDiagonalLayout, A, j)
+    isempty(j) && return 1:0
+    k = offdiagonalindex(A)
+    max(minimum(j)+k,1):min(size(A,2),maximum(j)+k)
+end
+
 function colsupport(::BidiagonalLayout, A, j)
     isempty(j) && return 1:0
     bidiagonaluplo(A) == 'L' ? (minimum(j):min(size(A,1),maximum(j)+1)) : (max(minimum(j)-1,1):maximum(j))

test/test_layoutarray.jl

@@ -299,6 +299,20 @@ BroadcastStyle(::Type{<:MyVector2{T}}) where {T} = MyBroadcastStyle{1}()
             @test permutedims(D) ≡ D
         end
 
+        @testset "Diagonal getindex returns OffDiagonal" begin
+            x = randn(5)
+            for D in (Diagonal(MyVector(x)), Diagonal(view(MyVector(x), :)), Diagonal(view(MyMatrix(reshape(x,5,1)), :, 1)))
+                for (kr,jr) in ((2:4,1:5), (1:5,2:4), (2:3,3:5), (:,2:3), (2:3,:), (:,:), (3:2,1:5))
+                    @test D[kr,jr] isa FillArrays.OffDiagonal
+                    @test D[kr,jr] == Diagonal(x)[kr,jr]
+                end
+                @test_throws BoundsError D[0:2,1:3]
+                @test_throws BoundsError D[1:3,2:6]
+            end
+            @test Diagonal(view(MyVector(x), 2:4))[1:2,2:3] isa FillArrays.OffDiagonal
+            @test Diagonal(view(MyVector(x), 2:4))[1:2,2:3] == Diagonal(x[2:4])[1:2,2:3]
+        end
+
         @testset "BroadcastStyle" begin
             A = MyMatrix2(randn(3,6))
             @test BroadcastStyle(typeof(reshape(A, (9, 2)))) == MyBroadcastStyle{2}()

test/test_layouts.jl

@@ -7,7 +7,7 @@ import ArrayLayouts: MemoryLayout, DenseRowMajor, DenseColumnMajor, StridedLayou
                         UnitUpperTriangularLayout, LowerTriangularLayout,
                         UnitLowerTriangularLayout, ScalarLayout, UnknownLayout,
                         hermitiandata, symmetricdata, FillLayout, ZerosLayout, OnesLayout,
-                        DiagonalLayout, TridiagonalLayout, SymTridiagonalLayout, colsupport, rowsupport,
+                        DiagonalLayout, OffDiagonalLayout, TridiagonalLayout, SymTridiagonalLayout, colsupport, rowsupport,
                         diagonaldata, subdiagonaldata, supdiagonaldata, BidiagonalLayout, bidiagonaluplo
 
 struct FooBar end
@@ -266,6 +266,48 @@ struct FooNumber <: Number end
         @test rowsupport(S,3) == rowsupport(H,3) == rowsupport(Sc,3) == rowsupport(Hc,3) == 2:4
     end
 
+    @testset "OffDiagonal" begin
+        A = FillArrays.OffDiagonal(randn(4), 1)
+        @test MemoryLayout(A) == MemoryLayout(A') == MemoryLayout(transpose(A)) == OffDiagonalLayout{DenseColumnMajor}()
+        @test MemoryLayout(FillArrays.OffDiagonal(Fill(1.0, 4), -1)) == OffDiagonalLayout{FillLayout}()
+        @test MemoryLayout(FillArrays.OffDiagonal(1:4, -1)) == OffDiagonalLayout{UnknownLayout}()
+        @test MemoryLayout(FillArrays.OffDiagonal(randn(4) .+ im, 1)') isa UnknownLayout
+
+        @test ArrayLayouts.offdiagonaldata(A) ≡ ArrayLayouts.offdiagonaldata(A') ≡ A.band
+        @test ArrayLayouts.offdiagonalindex(A) == 1
+        @test ArrayLayouts.offdiagonalindex(A') == ArrayLayouts.offdiagonalindex(transpose(A)) == -1
+
+        @test colsupport(A, 1) == 1:0
+        @test colsupport(A, 3) == 2:2
+        @test colsupport(A, 2:4) == 1:3
+        @test rowsupport(A, 1) == 2:2
+        @test rowsupport(A, 5) == 6:5
+        @test colsupport(A', 2) == rowsupport(A, 2) == 3:3
+
+        @testset "views" begin
+            D = Diagonal(randn(5))
+            for P in (A, A', transpose(A), FillArrays.OffDiagonal(Fill(2.0, 3), -2), D, Diagonal(Fill(2.0, 5)))
+                for (kr,jr) in ((1:3,2:4), (2:5,1:2), (:,:), (:,2:3), (2:3,:), (3:2,1:4), (2:3,2:2))
+                    V = view(P, kr, jr)
+                    @test MemoryLayout(V) isa OffDiagonalLayout
+                    @test ArrayLayouts.sub_materialize(V) isa FillArrays.OffDiagonal
+                    @test ArrayLayouts.sub_materialize(V) == Matrix(P)[kr,jr]
+                    @test ArrayLayouts.offdiagonaldata(V) == ArrayLayouts.sub_materialize(V).band
+                    @test ArrayLayouts.offdiagonalindex(V) == ArrayLayouts.sub_materialize(V).k
+                    for j in axes(V,2)
+                        @test all(iszero, Matrix(V)[setdiff(axes(V,1), colsupport(V,j)), j])
+                    end
+                    for k in axes(V,1)
+                        @test all(iszero, Matrix(V)[k, setdiff(axes(V,2), rowsupport(V,k))])
+                    end
+                end
+            end
+            @test MemoryLayout(view(D, 1:3, 2:4)) == OffDiagonalLayout{DenseColumnMajor}()
+            @test MemoryLayout(view(Diagonal(Fill(2.0, 5)), 1:3, 2:4)) == OffDiagonalLayout{FillLayout}()
+            @test MemoryLayout(view(A, [1,2], 2:3)) isa UnknownLayout
+        end
+    end
+
     @testset "triangular MemoryLayout" begin
         A = [1.0 2; 3 4]
         B = [1.0+im 2; 3 4]