Optimized row/col getindex (#373)

Author: jishnub

src/banded/BandedMatrix.jl

@@ -4,7 +4,7 @@
 #   a_21 a_22 a_23
 #   a_31 a_32 a_33 a_34
 #        a_42 a_43 a_44  ]
-# ordering the data like  (cobbndsmns first)
+# ordering the data like  (columns first)
 #       [ *      a_12   a_23    a_34
 #         a_11   a_22   a_33    a_44
 #         a_21   a_32   a_43    *
@@ -379,13 +379,45 @@ diagzero(D::Diagonal{B}, i, j) where B<:BandedMatrix =
 end
 
 
-# scalar - integer - integer
-@inline function getindex(A::BandedMatrix, k::Integer, j::Integer)
+# Int - Int
+@inline function getindex(A::BandedMatrix, k::Int, j::Int)
     @boundscheck checkbounds(A, k, j)
     @inbounds r = banded_getindex(A.data, A.l, A.u, k, j)
     r
 end
 
+# BandRange - Int
+@propagate_inbounds function getindex(A::BandedMatrix, ::BandRangeType, j::Int)
+    @boundscheck checkbounds(A, colrange(A, j), j)
+    A.data[data_colrange(A,j)]
+end
+
+# Colon - Int
+@propagate_inbounds function getindex(A::BandedMatrix, ::Colon, j::Int)
+    @boundscheck checkbounds(A, axes(A,1), j)
+    r = similar(A, axes(A,1))
+    r[firstindex(r):colstart(A,j)-1] .= zero(eltype(r))
+    r[colrange(A,j)] = @view A.data[data_colrange(A,j)]
+    r[colstop(A,j)+1:end] .= zero(eltype(r))
+    return r
+end
+
+# Int - BandRange
+@propagate_inbounds function getindex(A::BandedMatrix, k::Int, j::BandRangeType)
+    @boundscheck checkbounds(A, k, rowrange(A, k))
+    A.data[data_rowrange(A,k)]
+end
+
+# Int - Colon
+@propagate_inbounds function getindex(A::BandedMatrix, k::Int, ::Colon)
+    @boundscheck checkbounds(A, k, axes(A,2))
+    r = similar(A, axes(A,2))
+    r[firstindex(r):rowstart(A,k)-1] .= zero(eltype(r))
+    r[rowrange(A,k)] = @view A.data[data_rowrange(A,k)]
+    r[rowstop(A,k)+1:end] .= zero(eltype(r))
+    return r
+end
+
 # ~ indexing along a band
 # we reduce it to converting a View
 

test/test_indexing.jl

@@ -228,12 +228,21 @@ import BandedMatrices: rowstart, rowstop, colstart, colstop,
                          0  7  10  13  15   0   0;
                          0  0  11  14  16  17   0]
 
-            @test a[BandRange, 1] == [1,   2,  3]
-            @test a[BandRange, 2] == [4,   5,  6,  7]
-            @test a[BandRange, 3] == [8,   9, 10, 11]
-            @test a[BandRange, 4] == [12, 13, 14]
-            @test a[BandRange, 5] == [15, 16]
-            @test a[BandRange, 6] == [17]
+            @test a[BandRange, 1] == @view(a[BandRange, 1]) == [1,   2,  3]
+            @test a[BandRange, 2] == @view(a[BandRange, 2]) == [4,   5,  6,  7]
+            @test a[BandRange, 3] == @view(a[BandRange, 3]) == [8,   9, 10, 11]
+            @test a[BandRange, 4] == @view(a[BandRange, 4]) == [12, 13, 14]
+            @test a[BandRange, 5] == @view(a[BandRange, 5]) == [15, 16]
+            @test a[BandRange, 6] == @view(a[BandRange, 6]) == [17]
+            @test a[BandRange, 7] == @view(a[BandRange, 7]) == Int[]
+
+            @test a[:, 1] == view(a, :, 1) == [1,2,3,0,0]
+            @test a[:, 2] == view(a, :, 2) == [4,5,6,7,0]
+            @test a[:, 3] == view(a, :, 3) == [0,8,9,10,11]
+            @test a[:, 4] == view(a, :, 4) == [0,0,12,13,14]
+            @test a[:, 5] == view(a, :, 5) == [0,0,0,15,16]
+            @test a[:, 6] == view(a, :, 6) == [0,0,0,0,17]
+            @test a[:, 7] == view(a, :, 7) == [0,0,0,0,0]
 
             @test_throws BoundsError a[:, 0] = [1, 2, 3]
             @test_throws DimensionMismatch a[:, 1] = [1, 2, 3]
@@ -310,12 +319,21 @@ import BandedMatrices: rowstart, rowstop, colstart, colstop,
                          0  7  10  13  15   0   0;
                          0  0  11  14  16  17   0]'
 
-            @test a[1, BandRange] == [1,   2,  3]
-            @test a[2, BandRange] == [4,   5,  6,  7]
-            @test a[3, BandRange] == [8,   9, 10, 11]
-            @test a[4, BandRange] == [12, 13, 14]
-            @test a[5, BandRange] == [15, 16]
-            @test a[6, BandRange] == [17]
+            @test a[1, BandRange] == @view(a[1, BandRange]) == [1,   2,  3]
+            @test a[2, BandRange] == @view(a[2, BandRange]) == [4,   5,  6,  7]
+            @test a[3, BandRange] == @view(a[3, BandRange]) == [8,   9, 10, 11]
+            @test a[4, BandRange] == @view(a[4, BandRange]) == [12, 13, 14]
+            @test a[5, BandRange] == @view(a[5, BandRange]) == [15, 16]
+            @test a[6, BandRange] == @view(a[6, BandRange]) == [17]
+            @test a[7, BandRange] == @view(a[7, BandRange]) == Int[]
+
+            @test a[1, :] == @view(a[1, :]) == [1,2,3,0,0]
+            @test a[2, :] == @view(a[2, :]) == [4,5,6,7,0]
+            @test a[3, :] == @view(a[3, :]) == [0,8,9,10,11]
+            @test a[4, :] == @view(a[4, :]) == [0,0,12,13,14]
+            @test a[5, :] == @view(a[5, :]) == [0,0,0,15,16]
+            @test a[6, :] == @view(a[6, :]) == [0,0,0,0,17]
+            @test a[7, :] == @view(a[7, :]) == [0,0,0,0,0]
 
             @test_throws BoundsError a[0, :] = [1, 2, 3]
             @test_throws DimensionMismatch a[1, :] = [1, 2, 3]