Add OffDiagonal, returned when slicing Diagonal fills and RectDiagonal

Project.toml

@@ -1,6 +1,6 @@
 name = "FillArrays"
 uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
-version = "1.17.1"
+version = "1.18.0"
 
 
 [deps]

src/FillArrays.jl

@@ -715,6 +715,7 @@ end
 # include
 #########
 
+include("offdiagonal.jl")
 include("fillalgebra.jl")
 include("fillbroadcast.jl")
 include("trues.jl")
@@ -736,6 +737,7 @@ Base.print_matrix_row(io::IO,
                  AbstractFillMatrix,
                  Diagonal{<:Any,<:AbstractFillVector},
                  RectDiagonal,
+                 OffDiagonal,
                  UpperOrLowerTriangular{<:Any,<:AbstractFillMatrix}
                  }, A::Vector,
         i::Integer, cols::AbstractVector, sep::AbstractString, idxlast::Integer=last(axes(X, 2))) =

src/fillhash.jl

@@ -60,6 +60,7 @@ _hasinfaxes(A::AbstractArray) = any(ax -> isinf(length(ax)), axes(A))
 
 const HashableFill = Union{AbstractFill,
                            RectDiagonal,
+                           OffDiagonal,
                            OneElement,
                            Diagonal{<:Any,<:AbstractFillVector},
                            Bidiagonal{<:Any,<:AbstractFillVector},

src/offdiagonal.jl

@@ -0,0 +1,128 @@
+"""
+    OffDiagonal(band::AbstractVector, k::Integer, axes)
+    OffDiagonal(band::AbstractVector, k::Integer, m::Integer, n::Integer)
+    OffDiagonal(band::AbstractVector, k::Integer)
+
+A matrix whose only non-zero entries lie on its `k`-th diagonal, where they are given by `band`.
+As for `diag`, `k > 0` refers to a superdiagonal and `k < 0` to a subdiagonal, and
+`length(band)` must equal the length of that diagonal. If the size is omitted, the matrix is
+square with `length(band) + abs(k)` rows.
+
+Indexing a `Diagonal` of an `AbstractFill`, a `RectDiagonal` or an `OffDiagonal` by unit ranges
+returns an `OffDiagonal`.
+
+# Examples
+
+```jldoctest
+julia> FillArrays.OffDiagonal([1, 2, 3], 1)
+4×4 FillArrays.OffDiagonal{Int64, Vector{Int64}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}:
+ ⋅  1  ⋅  ⋅
+ ⋅  ⋅  2  ⋅
+ ⋅  ⋅  ⋅  3
+ ⋅  ⋅  ⋅  ⋅
+
+julia> A = Eye(5)[3:5, 1:4];
+
+julia> A.k
+2
+
+julia> A.band
+2-element Ones{Float64}
+```
+"""
+struct OffDiagonal{T,V<:AbstractVector{T},Axes<:Tuple{Vararg{AbstractUnitRange,2}}} <: AbstractMatrix{T}
+    band::V
+    k::Int
+    axes::Axes
+
+    @inline function OffDiagonal{T,V,Axes}(band::V, k::Integer, axes::Axes) where {T,V<:AbstractVector{T},Axes<:Tuple{Vararg{AbstractUnitRange,2}}}
+        Base.require_one_based_indexing(band)
+        len = _bandlength(map(length, axes)..., k)
+        length(band) == len || throw(DimensionMismatch(LazyString("band has length ", length(band), " but diagonal ", k, " of a ", Base.dims2string(map(length, axes)), " matrix has length ", len)))
+        A = new{T,V,Axes}(band, k, axes)
+        Base.require_one_based_indexing(A)
+        return A
+    end
+end
+
+# number of entries on the k-th diagonal of an m × n matrix, which is zero when the diagonal lies outside the matrix
+_bandlength(m, n, k) = max(0, k ≥ 0 ? min(m, n - k) : min(m + k, n))
+
+@inline OffDiagonal{T,V}(band::V, k::Integer, axes::Axes) where {T,V,Axes<:Tuple{Vararg{AbstractUnitRange,2}}} = OffDiagonal{T,V,Axes}(band, k, axes)
+@inline OffDiagonal{T,V}(band::V, k::Integer, sz::Tuple{Vararg{Integer,2}}) where {T,V} = OffDiagonal{T,V}(band, k, oneto.(sz))
+@inline OffDiagonal{T,V}(band::V, k::Integer, m::Integer, n::Integer) where {T,V} = OffDiagonal{T,V}(band, k, (m, n))
+@inline function OffDiagonal{T,V}(band::V, k::Integer) where {T,V}
+    n = length(band) + abs(k)
+    OffDiagonal{T,V}(band, k, (n, n))
+end
+@inline function OffDiagonal{T}(band::AbstractVector, args...) where T
+    b = convert(AbstractVector{T}, band)
+    OffDiagonal{T,typeof(b)}(b, args...)
+end
+@inline OffDiagonal(band::V, args...) where {V<:AbstractVector} = OffDiagonal{eltype(V),V}(band, args...)
+
+axes(A::OffDiagonal) = A.axes
+size(A::OffDiagonal) = map(length, A.axes)
+
+@inline function getindex(A::OffDiagonal{T}, i::Integer, j::Integer) where T
+    @boundscheck checkbounds(A, i, j)
+    if j - i == A.k
+        @inbounds r = A.band[min(i, j)]
+    else
+        r = zero(T)
+    end
+    return r
+end
+
+function setindex!(A::OffDiagonal, v, i::Integer, j::Integer)
+    @boundscheck checkbounds(A, i, j)
+    if j - i == A.k
+        @inbounds A.band[min(i, j)] = v
+    elseif !iszero(v)
+        throw(ArgumentError(LazyString("cannot set entry (", i, ", ", j, ") off diagonal ", A.k, " to a nonzero value (", v, ")")))
+    end
+    return v
+end
+
+function diag(A::OffDiagonal{T}, k::Integer=0) where T
+    k == A.k && return copy(A.band)
+    fill!(similar(A.band, length(diagind(A, k))), zero(T))
+end
+
+Base.replace_in_print_matrix(A::OffDiagonal, i::Integer, j::Integer, s::AbstractString) =
+    j - i == A.k ? s : Base.replace_with_centered_mark(s)
+
+##
+# Slicing a matrix supported on a single diagonal by unit ranges keeps a single diagonal
+##
+
+const SingleBandMatrix = Union{Diagonal{<:Number,<:AbstractFillVector}, RectDiagonal, OffDiagonal}
+
+_band(D::Union{Diagonal,RectDiagonal}) = D.diag
+_band(A::OffDiagonal) = A.band
+_bandindex(::Union{Diagonal,RectDiagonal}) = 0
+_bandindex(A::OffDiagonal) = A.k
+
+Base.@propagate_inbounds function getindex(A::SingleBandMatrix, kr::Union{AbstractUnitRange{<:Integer},Colon}, jr::Union{AbstractUnitRange{<:Integer},Colon})
+    I = (_onebased_index(A, 1, kr), _onebased_index(A, 2, jr))
+    _singleband_getindex(Base.index_shape(I...), A, I...)
+end
+
+# all SingleBandMatrix are one-based, so a Colon may be replaced by a OneTo even if the axes are not OneTo
+_onebased_index(A, d, ::Colon) = oneto(size(A, d))
+_onebased_index(_, _, r) = r
+
+# indices with offset or infinite axes use the generic fallback
+Base.@propagate_inbounds _singleband_getindex(_, A, kr, jr) = invoke(getindex, Tuple{AbstractArray, Vararg{Any}}, A, kr, jr)
+
+@inline function _singleband_getindex(shape::Tuple{Base.OneTo,Base.OneTo}, A, kr, jr)
+    @boundscheck checkbounds(A, kr, jr)
+    # entry (i, j) of the result is A[kr[i], jr[j]], which lies on the band when jr[j] - kr[i] == _bandindex(A)
+    k = _bandindex(A) + first(kr) - first(jr)
+    m, n = length(kr), length(jr)
+    len = _bandlength(m, n, k)
+    # the band of A is indexed by min(row, col), so start from the first entry of the k-th diagonal of the result
+    st = min(first(kr) + max(0, -k), first(jr) + max(0, k))
+    @inbounds band = _band(A)[st:st+len-1]
+    OffDiagonal(band, k, shape)
+end

test/runtests.jl

@@ -1,6 +1,6 @@
 using FillArrays, LinearAlgebra, PDMats, SparseArrays, StaticArrays, ReverseDiff, Random, Test, Statistics, Quaternions
 
-import FillArrays: AbstractFill, RectDiagonal, SquareEye
+import FillArrays: AbstractFill, RectDiagonal, OffDiagonal, SquareEye
 
 using Documenter
 DocMeta.setdocmeta!(FillArrays, :DocTestSetup, :(using FillArrays))
@@ -386,6 +386,91 @@ end
     @test stringmime("text/plain", D) == "3×2 RectDiagonal{Float64, Vector{Float64}, Tuple{Base.OneTo{$Int}, Base.OneTo{$Int}}}:\n 1.0   ⋅ \n  ⋅   2.0\n  ⋅    ⋅ "
 end
 
+@testset "OffDiagonal" begin
+    @testset "constructors" begin
+        A = OffDiagonal([1, 2, 3], 1)
+        @test A isa OffDiagonal{Int, Vector{Int}}
+        @test size(A) == (4, 4)
+        @test axes(A) == (Base.OneTo(4), Base.OneTo(4))
+        @test A == diagm(1 => [1, 2, 3])
+        @test OffDiagonal([1, 2, 3], -1) == diagm(-1 => [1, 2, 3])
+        @test OffDiagonal(1:3, 2, 3, 5) == OffDiagonal(1:3, 2, (3, 5)) ==
+            OffDiagonal(1:3, 2, (Base.OneTo(3), Base.OneTo(5))) == [0 0 1 0 0; 0 0 0 2 0; 0 0 0 0 3]
+        @test OffDiagonal(1:2, -1, 3, 2) == [0 0; 1 0; 0 2]
+        @test OffDiagonal{Float64}(1:3, 1) isa OffDiagonal{Float64}
+        @test OffDiagonal{Float64}(1:3, 1) == A
+        # the band lies outside the matrix
+        @test OffDiagonal(Int[], 4, 2, 3) == zeros(Int, 2, 3)
+        @test OffDiagonal(Int[], -3, 2, 3) == zeros(Int, 2, 3)
+        @test_throws DimensionMismatch OffDiagonal([1, 2], 1, 4, 4)
+        @test_throws DimensionMismatch OffDiagonal([1], 4, 2, 3)
+    end
+
+    @testset "indexing" begin
+        A = OffDiagonal([1, 2, 3], -2, 5, 4)
+        @test A == [0 0 0 0; 0 0 0 0; 1 0 0 0; 0 2 0 0; 0 0 3 0]
+        @test A[4, 2] == 2
+        @test A[2, 4] == 0
+        @test_throws BoundsError A[6, 1]
+        @test diag(A, -2) == [1, 2, 3]
+        @test diag(A, -2) !== A.band
+        @test diag(A) == zeros(Int, 4)
+        @test diag(A, 3) == [0]
+        @test diag(Eye(5)[2:5, 1:4], 1) ≡ Ones(3)
+
+        A[4, 2] = 5
+        @test A.band == [1, 5, 3]
+        A[1, 1] = 0
+        @test_throws ArgumentError A[1, 1] = 1
+        @test_throws BoundsError A[6, 1] = 0
+
+        @test stringmime("text/plain", OffDiagonal([1, 2], 1)) ==
+            "3×3 $OffDiagonal{$Int, Vector{$Int}, Tuple{Base.OneTo{$Int}, Base.OneTo{$Int}}}:\n ⋅  1  ⋅\n ⋅  ⋅  2\n ⋅  ⋅  ⋅"
+        @test hash(OffDiagonal([1, 2], 1)) == hash([0 1 0; 0 0 2; 0 0 0])
+    end
+
+    @testset "slicing returns OffDiagonal" begin
+        for M in (Eye(5), Eye{Int}(4), Diagonal(Fill(2, 6)), Diagonal(Zeros(3)), Eye(4, 6), Eye(6, 4),
+                  RectDiagonal([1, 2, 3]), RectDiagonal([1, 2, 3], (1:3, 1:5)), RectDiagonal([1, 2, 3], (1:5, 1:3)),
+                  RectDiagonal(Fill(2, 3), 3, 4), OffDiagonal([1, 2, 3], 1), OffDiagonal([1, 2, 3], -2),
+                  OffDiagonal(Fill(2, 3), 2, 3, 5), OffDiagonal(Int[], 4, 3, 4))
+            Md = Matrix(M)
+            for a in axes(M, 1), b in a-1:size(M, 1), c in axes(M, 2), d in c-1:size(M, 2)
+                S = M[a:b, c:d]
+                @test S isa OffDiagonal{eltype(M)}
+                @test S == Md[a:b, c:d]
+            end
+            @test M[:, :] isa OffDiagonal
+            @test M[:, :] == Md
+            @test M[Base.OneTo(2), :] == Md[1:2, :]
+            @test M[:, 2:3] == Md[:, 2:3]
+            @test_throws BoundsError M[0:1, 1:1]
+            @test_throws BoundsError M[1:1, 1:size(M, 2)+1]
+        end
+
+        # lazy bands stay lazy
+        @test Eye(5)[2:4, 1:3] ≡ OffDiagonal(Ones(2), 1, 3, 3)
+        @test Diagonal(Fill(2, 6))[2:4, 3:6] ≡ OffDiagonal(Fill(2, 2), -1, 3, 4)
+        @test Eye(4, 6)[:, 3:6] ≡ OffDiagonal(Ones(2), -2, 4, 4)
+        @test Eye(5)[1:2, 4:5] ≡ OffDiagonal(Ones(0), -3, 2, 2)
+        @test OffDiagonal(Fill(2, 4), 1)[2:4, 3:5] ≡ OffDiagonal(Fill(2, 3), 0, 3, 3)
+        @test OffDiagonal(1:4, -1)[2:5, 1:3] == OffDiagonal(1:3, 0, 4, 3)
+
+        # Diagonals not backed by a fill are unaffected
+        @test Diagonal([1, 2, 3])[1:2, 1:2] isa Matrix{Int}
+
+        # indices whose axes are not Base.OneTo use the generic getindex
+        for M in (Eye(5), Eye(4, 6), OffDiagonal([1, 2, 3], 1))
+            S = M[SOneTo(3), SOneTo(2)]
+            @test S isa MMatrix{3, 2, eltype(M)}
+            @test S == Matrix(M)[1:3, 1:2]
+        end
+
+        @test @inferred(Eye(5)[2:4, :]) isa OffDiagonal
+        @test @inferred(OffDiagonal(1:4, 1)[:, 2:3]) isa OffDiagonal
+    end
+end
+
 # Check that all pair-wise combinations of + / - elements of As and Bs yield the correct
 # type, and produce numerically correct results.
 as_array(x::AbstractArray) = Array(x)