Fix ∞-dimensional setindex!(::DiagTrav, ...) (#143)

* Fix ∞-dimensional setindex!(::DiagTrav, ...)

* Update blockkron.jl

Author: dlfivefifty

Project.toml

@@ -1,7 +1,7 @@
 name = "LazyBandedMatrices"
 uuid = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.11.3"
+version = "0.11.4"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/blockkron.jl

@@ -98,7 +98,32 @@ function colsupport(A::DiagTrav{<:Any,2}, _)
 end
 
 
-function getindex(A::DiagTrav, K::Block{1})
+Base.@propagate_inbounds function getindex(A::DiagTrav, Kk::BlockIndex{1})
+    @boundscheck checkbounds(A, Kk)
+    _diagtravgetindex(MemoryLayout(A.array), A.array, Kk)
+end
+
+function _diagtravgetindex(_, A::AbstractMatrix, Kk::BlockIndex{1})
+    K,j = Int(block(Kk)), blockindex(Kk)
+    m,n = size(A)
+    j = max(0,K-m)+j
+    A[K-j+1,j]
+end
+
+Base.@propagate_inbounds function setindex!(A::DiagTrav, v, Kk::BlockIndex{1})
+    @boundscheck checkbounds(A, Kk)
+    _diagtravsetindex!(MemoryLayout(A.array), A.array, v, Kk)
+    A
+end
+
+function _diagtravsetindex!(_, A, v, Kk::BlockIndex{1})
+    K,j = Int(block(Kk)), blockindex(Kk)
+    m,n = size(A)
+    j = max(0,K-m)+j
+    A[K-j+1,j] = v
+end
+
+Base.@propagate_inbounds function getindex(A::DiagTrav, K::Block{1})
     @boundscheck checkbounds(A, K)
     _diagtravgetindex(MemoryLayout(A.array), A.array, K)
 end
@@ -157,6 +182,9 @@ function _diagtravgetindex(::AbstractStridedLayout, A::AbstractArray{T,3}, K::Bl
     ret
 end
 
+# TODO: don't build the block
+_diagtravgetindex(lay, A::AbstractArray{T,3}, Kk::BlockIndex{1}) where T = _diagtravgetindex(lay, A, block(Kk))[blockindex(Kk)]
+
 getindex(A::DiagTrav, k::Int) = A[findblockindex(axes(A,1), k)]
 setindex!(A::DiagTrav, v, k::Int) = A[findblockindex(axes(A,1), k)] = v
 

test/test_lazybandedinf.jl

@@ -76,6 +76,13 @@ const InfKronTravBandedBlockBandedLayout = LazyBandedMatricesInfiniteArraysExt.I
         @test blockcolsupport(A) == Block.(1:6)
         @test A[Block.(1:7)] == [1; 5; 2; 9; 6; 3; 0; 10; 7; 4; 0; 0; 11; 8; 0; 0; 0; 0; 12; zeros(9)]
 
+        C = zeros(∞,∞);
+        A = DiagTrav(C);
+        A[1] = 1
+        @test A[1] == 1
+        A[1:9] = 1:9
+        @test A[1:10] == [1:9; 0]
+
         C = zeros(∞,4);
         C[1:3,1:4] .= [1 2 3 4; 5 6 7 8; 9 10 11 12]
         A = DiagTrav(C)