axpy! for BandedMatrix (#336)

* axpy for BandedMatrix

* test for matching bandwidth

* Add dimension mismatch tests

Author: jishnub

src/banded/BandedMatrix.jl

@@ -830,3 +830,14 @@ end
 
 bandedbroadcaststyle(_) = BandedStyle()
 BroadcastStyle(::Type{<:BandedMatrix{<:Any,Dat}}) where Dat = bandedbroadcaststyle(BroadcastStyle(Dat))
+
+function banded_axpy!(a::Number, X::BandedMatrix, Y::BandedMatrix)
+    bx = bandwidths(X)
+    by = bandwidths(Y)
+    if bx == by
+        axpy!(a, X.data, Y.data)
+    else
+        banded_generic_axpy!(a, X, Y)
+    end
+    return Y
+end

src/generic/broadcast.jl

@@ -932,10 +932,11 @@ _banded_axpy!(a::Number, X::AbstractMatrix, Y::AbstractMatrix, notbandedX, notba
     banded_dense_axpy!(a, X, Y)
 
 # additions and subtractions
-@propagate_inbounds function banded_generic_axpy!(a::Number, X::AbstractMatrix, Y::AbstractMatrix)
+function banded_generic_axpy!(a::Number, X::AbstractMatrix, Y::AbstractMatrix)
     n,m = size(X)
-    if (n,m) ≠ size(Y)
-        throw(BoundsError())
+    ny,my = size(Y)
+    if (n,m) ≠ (ny,my)
+        throw(DimensionMismatch("X has size $((n,m)) but $Y has size $((ny,my))"))
     end
     Xl, Xu = bandwidths(X)
     Yl, Yu = bandwidths(Y)
@@ -944,17 +945,17 @@ _banded_axpy!(a::Number, X::AbstractMatrix, Y::AbstractMatrix, notbandedX, notba
         return Y
     end
 
-    @boundscheck if Xl > Yl
+    if Xl > Yl
         # test that all entries are zero in extra bands below the diagonal
-        for j=rowsupport(X),k=max(1,j+Yl+1):min(j+Xl,n)
+        @inbounds for j=rowsupport(X),k=max(1,j+Yl+1):min(j+Xl,n)
             if inbands_getindex(X, k, j) ≠ 0
                 throw(BandError(Y, (k,j)))
             end
         end
     end
-    @boundscheck if Xu > Yu
+    if Xu > Yu
         # test that all entries are zero in extra bands above the diagonal
-        for j=rowsupport(X),k=max(1,j-Xu):min(j-Yu-1,n)
+        @inbounds for j=rowsupport(X),k=max(1,j-Xu):min(j-Yu-1,n)
             if inbands_getindex(X, k, j) ≠ 0
                 throw(BandError(Y, (k,j)))
             end
@@ -976,8 +977,10 @@ _banded_axpy!(a::Number, X::AbstractMatrix, Y::AbstractMatrix, notbandedX, notba
 end
 
 function banded_dense_axpy!(a::Number, X::AbstractMatrix, Y::AbstractMatrix)
-    if size(X) != size(Y)
-        throw(DimensionMismatch("+"))
+    n,m = size(X)
+    ny,my = size(Y)
+    if (n,m) ≠ (ny,my)
+        throw(DimensionMismatch("X has size $((n,m)) but $Y has size $((ny,my))"))
     end
     @inbounds for j=rowsupport(X), k=colrange(X,j)
         Y[k,j] += a*inbands_getindex(X,k,j)

test/test_broadcasting.jl

@@ -312,6 +312,9 @@ import BandedMatrices: BandedStyle, BandedRows
         B .= 2.0 .* A .+ B
         @test B == C
 
+        # test with identical bandwidth
+        @test axpy!(3, A, copy(A)) ≈ 4A
+
         @testset "trivial cases" begin
             B = brand(2,4,-1,0) # no bands in B
             B2 = brand(2,4,0,-1) # no bands in B2
@@ -326,6 +329,10 @@ import BandedMatrices: BandedStyle, BandedRows
             axpy!(0.1, C, B) # no bands in dest, but src is zero
             @test B == D
         end
+
+        @test_throws DimensionMismatch axpy!(2, brand(2,2,1,1), brand(3,3,1,1))
+        @test_throws DimensionMismatch axpy!(2, brand(2,2,1,1), brand(3,3,2,2))
+        @test_throws DimensionMismatch axpy!(2, brand(2,2,1,1), zeros(3,3))
     end
 
     @testset "gbmv!" begin