Allow vector-based indices (#181)

* Allow vector-based indices

* Handle discontinous integer indices

* Allow integer indices, now handled by dispatch

* Allow BitVectors

* Add tentative tests

* Define chunks of vector subset through run-length encoding

* Allow AbstractVector views

---------

Co-authored-by: Fabian Gans <[email protected]>

Author: ConnectedSystems

src/chunks.jl

@@ -71,6 +71,36 @@ findchunk(r::RegularChunks, i::Int) = div(i + r.offset - 1, r.cs) + 1
 
 subsetchunks(r, subs) = subsetchunks_fallback(r, subs)
 
+"""
+    chunk_rle(chunks,vec)
+
+Computes the run length of which chunk is accessed how many times consecutively to determine
+
+"""
+function chunk_rle(chunks, vec)
+    out = Int[]
+    currentchunk = -1
+    for i in vec
+        nextchunk = findchunk(chunks, i)
+        if nextchunk == currentchunk
+            out[end] += 1
+        else
+            push!(out, 1)
+            currentchunk = nextchunk
+        end
+    end
+    out
+end
+
+"""
+    subsetchunks(r, subs::BitVector)
+
+Identify chunks from indices which may be discontinuous.
+"""
+function subsetchunks(r, subs::BitVector)
+    return subsetchunks(r, findall(subs))
+end
+
 approx_chunksize(r::RegularChunks) = r.cs
 grid_offset(r::RegularChunks) = r.offset
 max_chunksize(r::RegularChunks) = r.cs
@@ -155,7 +185,8 @@ function subsetchunks_fallback(r, subs)
     elseif issorted(subs; rev=true)
         true
     else
-        throw(ArgumentError("Can only subset chunks for sorted indices"))
+        rle = chunk_rle(r, subs)
+        return chunktype_from_chunksizes(rle)
     end
     cs = zeros(Int, length(r))
     for i in subs
@@ -342,8 +373,6 @@ function estimate_chunksize(s, si)
     return GridChunks(s, cs)
 end
 
-
-
 abstract type ChunkTiledDiskArray{T,N} <: AbstractDiskArray{T,N} end
 Base.size(a::ChunkTiledDiskArray) = arraysize_from_chunksize.(eachchunk(a).chunks)
 function DiskArrays.readblock!(A::ChunkTiledDiskArray{T,N}, data, I...) where {T,N}

src/subarray.jl

@@ -25,8 +25,8 @@ end
 eachchunk(a::SubDiskArray) = eachchunk_view(haschunks(a.v.parent), a.v)
 function eachchunk_view(::Chunked, vv)
     pinds = parentindices(vv)
-    if any(ind->!isa(ind,Union{Int,AbstractRange,Colon}),pinds)
-        throw(ArgumentError("Unable to determine chunksize of non-range views."))
+    if any(ind->!isa(ind,Union{Int,AbstractRange,Colon,AbstractVector{<:Integer}}),pinds)
+        throw(ArgumentError("Unable to determine chunksize for view of type $(typeof.(pinds))."))
     end
     iomit = findints(pinds)
     chunksparent = eachchunk(parent(vv))

test/runtests.jl

@@ -266,7 +266,7 @@ end
     @test subsetchunks(r1, [28, 27, 19, 17, 10, 7]) == [1:3, 4:5, 6:6]
     @test subsetchunks(r1, [1, 2, 3]) == [1:3]
     @test subsetchunks(r1, [1, 2, 3, 10, 11]) == [1:3, 4:5]
-    @test_throws ArgumentError subsetchunks(r1, [3, 4, 5, 1])
+    @test subsetchunks(r1, [3, 4, 11, 13, 1]) == [1:2,3:4,5:5]
 
     r2 = IrregularChunks(; chunksizes=[3, 3, 4, 3, 3, 4])
     @test subsetchunks(r2, 1:20) == r2
@@ -577,6 +577,25 @@ end
     i = Bool[1, 0, 0, 1, 1, 0, 0, 1, 0, 1, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 0, 1, 1, 1, 0, 1, 0, 1, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 0, 1, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 1]
     u = UnchunkedDiskArray(rand(275,305,36))
     @test u[i,1,1][1] == u[findfirst(i),1,1]
+
+    # Test for #180
+    test_arr = rand(100,100,100)
+    a180 = AccessCountDiskArray(test_arr; chunksize=(10,10,20))
+
+    idx = [rand(1:100)]
+    @test all(a180[idx, :, :] .== test_arr[idx, :, :])
+    @test all(a180[:, idx, :] .== test_arr[:, idx, :])
+    @test all(a180[:, :, idx] .== test_arr[:, :, idx])
+
+    sel = findall(rand(100) .<= 0.5)
+    @test all(a180[:, sel, :] .== test_arr[:, sel, :])
+    @test all(a180[sel, :, :] .== test_arr[sel, :, :])
+    @test all(a180[sel, :, sel] .== test_arr[sel, :, sel])
+
+    bit_sel = rand(100) .> 0.5
+    @test all(a180[bit_sel, :, :] .== test_arr[bit_sel, :, :])
+    @test all(a180[:, bit_sel, :] .== test_arr[:, bit_sel, :])
+    @test all(a180[:, bit_sel, bit_sel] .== test_arr[:, bit_sel, bit_sel])
 end
 
 
@@ -712,10 +731,10 @@ end
 
     @test collect(reverse(a_disk)) == reverse(a)
     @test reverse(view(a_disk, :, 1)) == reverse(a[:, 1])
-    @test_broken reverse(view(a_disk, :, 1), 1) == reverse(a[:, 1], 1)
+    @test reverse(view(a_disk, :, 1), 1) == reverse(a[:, 1], 1)
     # ERROR: ArgumentError: Can only subset chunks for sorted indices
-    @test_broken reverse(view(a_disk, :, 1), 5) == reverse(a[:, 1], 5)
-    @test_broken reverse(view(a_disk, :, 1), 5, 10) == reverse(a[:, 1], 5, 10)
+    @test reverse(view(a_disk, :, 1), 5) == reverse(a[:, 1], 5)
+    @test reverse(view(a_disk, :, 1), 5, 10) == reverse(a[:, 1], 5, 10)
     @test collect(reverse(a_disk)) == collect(reverse(a_disk; dims=:)) == 
         collect(reverse(a_disk; dims=(1, 2))) == reverse(a)
     @test collect(reverse(a_disk; dims=2)) == reverse(a; dims=2)