Broadcast a vector to a band (#367)

* Broadcast a vector to a band

* copy non-data band to offset vector

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "BandedMatrices"
 uuid = "aae01518-5342-5314-be14-df237901396f"
-version = "0.17.25"
+version = "0.17.26"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/banded/BandedMatrix.jl

@@ -448,18 +448,18 @@ function dataview(V::BandedMatrixBand)
 end
 
 # B[band(i)]
-function copyto!(v::Vector, B::BandedMatrixBand)
+function copyto!(v::AbstractVector, B::BandedMatrixBand)
     A = parent(parent(B))
     if -A.l ≤ band(B) ≤ A.u
         copyto!(v, dataview(B))
     else
         Binds = axes(B,1)
-        v[Binds] .= 0
+        v[firstindex(v)-1 .+ Binds] .= 0
     end
     return v
 end
 
-# B[band(i)] .= x
+# B[band(i)] .= x::Number
 function fill!(Bv::BandedMatrixBand, x)
     b = band(Bv)
     A = parent(parent(Bv))
@@ -472,6 +472,29 @@ function fill!(Bv::BandedMatrixBand, x)
     Bv
 end
 
+# B[band(i)] .= V::AbstractVector
+function _copyto!(Bv::BandedMatrixBand, V::AbstractVector)
+    isempty(V) && return Bv
+    A = parent(parent(Bv))
+    if -A.l ≤ band(Bv) ≤ A.u
+        copyto!(dataview(Bv), V)
+    else
+        # bounds-checking to work around axis offset of V
+        destinds, srcinds = LinearIndices(Bv), LinearIndices(V)
+        idf, isf = first(destinds), first(srcinds)
+        Δi = idf - isf
+        (checkbounds(Bool, destinds, isf+Δi) &
+            checkbounds(Bool, destinds, last(srcinds)+Δi)) ||
+                throw(BoundsError(dest, srcinds))
+
+        all(iszero, V) || throw(BandError(A, band(Bv)))
+    end
+    return Bv
+end
+
+copyto!(Bv::BandedMatrixBand, V::AbstractVector) = _copyto!(Bv, V)
+copyto!(Bv::BandedMatrixBand, V::BandedMatrixBand) = _copyto!(Bv, V)
+
 # ~ indexing along a row
 
 
@@ -530,6 +553,7 @@ end
 
 @propagate_inbounds function setindex!(A::BandedMatrix{T}, v::AbstractVector, ::Colon) where {T}
     A[:, :] = reshape(v,size(A))
+    A
 end
 
 
@@ -542,10 +566,10 @@ end
     row = A.u - b.i + 1
     data, i = A.data, max(b.i + 1, 1)
     for v in V
-        data[row, i] = convert(T, v)::T
+        @inbounds data[row, i] = v
         i += 1
     end
-    V
+    A
 end
 
 
@@ -557,12 +581,14 @@ end
     @boundscheck checkbandmatch(A,V,:,j)
 
     A.data[data_colrange(A,j)] = V[colrange(A,j)]
-    V
+    A
 end
 
 # vector - BandRange - integer -- A[1, BandRange] = 2
-@propagate_inbounds setindex!(A::BandedMatrix, V::AbstractVector, ::BandRangeType, j::Integer) =
-    (A[colrange(A, j), j] = V) # call range method
+@propagate_inbounds function setindex!(A::BandedMatrix, V::AbstractVector, ::BandRangeType, j::Integer)
+    A[colrange(A, j), j] = V # call range method
+    A
+end
 
 # vector - range - integer -- A[1:3, 1] = [1, 2, 3]
 @propagate_inbounds function setindex!(A::BandedMatrix, V::AbstractVector, kr::AbstractRange, j::Integer)
@@ -580,7 +606,7 @@ end
             inbands_setindex!(data, u, v, k, j)
         end
     end
-    V
+    A
 end
 
 
@@ -607,12 +633,14 @@ end
     end
 
     A.data[data_rowrange(A,k)] = V[rowrange(A,k)]
-    V
+    A
 end
 
 # vector - integer - BandRange -- A[1, BandRange] = [1, 2, 3]
-@propagate_inbounds setindex!(A::BandedMatrix, V::AbstractVector, k::Integer, ::BandRangeType) =
-    (A[k, rowstart(A, k):rowstop(A, k)] = V) # call range method
+@propagate_inbounds function setindex!(A::BandedMatrix, V::AbstractVector, k::Integer, ::BandRangeType)
+    A[k, rowstart(A, k):rowstop(A, k)] = V # call range method
+    A
+end
 
 # vector - integer - range -- A[1, 2:3] = [3, 4]
 @propagate_inbounds function setindex!(A::BandedMatrix, V::AbstractVector, k::Integer, jr::AbstractRange)
@@ -630,7 +658,7 @@ end
             inbands_setindex!(data, u, v, k, j)
         end
     end
-    V
+    A
 end
 
 # ~ indexing over a rectangular block
@@ -641,7 +669,7 @@ end
     @boundscheck checkdimensions(kr, jr, V)
     @boundscheck checkbandmatch(A, V, kr, jr)
     copyto!(view(A, kr, jr), V)
-    V
+    A
 end
 
 # ~~ end setindex! ~~

test/test_banded.jl

@@ -471,7 +471,7 @@ Base.similar(::MyMatrix, ::Type{T}, m::Int, n::Int) where T = MyMatrix{T}(undef,
     end
 
     @testset "setindex! with ranges (#348)" begin
-        n = 10; 
+        n = 10;
         X1 = brand(n,n,1,1)
         B = BandedMatrix(Zeros(2n,2n), (3,3))
         B[1:2:end,1:2:end] = X1
@@ -480,20 +480,19 @@ Base.similar(::MyMatrix, ::Type{T}, m::Int, n::Int) where T = MyMatrix{T}(undef,
         @test A == B
     end
 
-    @testset "broadcasting to a band" begin
-        B = brand(Int8, 6, 6, -2, 2)
-        B[band(2)] .= 10
-        @test all(==(10), diag(B, 2))
-        @test all(==(10), B[band(2)])
-
-        B = brand(Int8, 2, 4, 1, 1)
-        B[band(-1)] .= 2
-        B[band(0)] .= 3
-        B[band(1)] .= 4
-        @test B == [3 4 0 0; 2 3 4 0]
-
-        @test_throws BandError B[band(100)] .= 10
-        @test_throws BandError B[band(-100)] .= 10
+    @testset "copy band to offset vector" begin
+        B = BandedMatrix(2=>2:3)
+        # test that a BandedMatrix and a Matrix behave identically
+        p = zeros(3)
+        v = view(p, Base.IdentityUnitRange(2:3))
+        for M in (B, Matrix(B))
+            p .= 0
+            copyto!(v, diag(M, 2))
+            @test p[1] == 0
+            @test @view(p[2:3]) == 2:3
+            copyto!(v, diag(M, 10))
+            @test p[1] == 0
+            @test @view(p[2:3]) == 2:3
+        end
     end
 end
-

test/test_broadcasting.jl

@@ -729,4 +729,30 @@ import BandedMatrices: BandedStyle, BandedRows, BandError
         @test Z .* A ≡ Z
         @test A .* Z ≡ Z
     end
+
+    @testset "broadcasting to a band" begin
+        B = brand(Int8, 6, 6, -2, 2)
+        B[band(2)] .= 10
+        @test all(==(10), diag(B, 2))
+        @test all(==(10), B[band(2)])
+
+        B = brand(Int8, 2, 4, 1, 1)
+        B[band(-1)] .= 2
+        B[band(0)] .= 3
+        B[band(1)] .= 4
+        B[band(2)] .= 0
+        @test B == [3 4 0 0; 2 3 4 0]
+
+        B[band(-1)] .= 3:3
+        @test all(==(3), @view B[band(-1)])
+        B[band(0)] .= 4:5
+        @test (@view B[band(0)]) == 4:5
+        B[band(1)] .= 3:4
+        @test (@view B[band(1)]) == 3:4
+        B[band(2)] .= [0,0]
+        @test all(iszero, @view B[band(2)])
+
+        @test_throws BandError B[band(100)] .= 10
+        @test_throws BandError B[band(-100)] .= 10
+    end
 end