add PaddedDiskArray and pad (#219)

* add PaddedDiskArray and pad

* fix pad

* test and bugfix _pad_offset

* bugfixes

* call getindex not readblock

Author: rafaqz

src/DiskArrays.jl

@@ -30,6 +30,7 @@ include("generator.jl")
 include("zip.jl")
 include("show.jl")
 include("cached.jl")
+include("pad.jl")
 
 # The all-in-one macro
 

src/chunks.jl

@@ -61,7 +61,7 @@ struct RegularChunks <: ChunkVector
     arraysize::Int
     function RegularChunks(chunksize::Integer, offset::Integer, arraysize::Integer)
         chunksize > 0 || throw(ArgumentError("Chunk sizes must be strictly positive"))
-        -1 < offset < chunksize || throw(ArgumentError("Offsets must be positive and smaller than the chunk size"))
+        -1 < offset < chunksize || throw(ArgumentError("Offsets must be positive and smaller than the chunk size, got offset: $offset and chunk size: $chunksize"))
         arraysize >= 0 || throw(ArgumentError("Negative axis lengths are not allowed"))
         new(Int(chunksize), Int(offset), Int(arraysize))
     end

src/pad.jl

@@ -0,0 +1,115 @@
+"""
+    PaddedDiskArray <: AbstractDiskArray
+
+    PaddedDiskArray(A, padding; fill=zero(eltype(A)))
+
+An `AbstractDiskArray` that adds padding to the edges of the parent array.
+This can help changing chunk offsets or padding a larger than memory array
+before windowing operations.
+
+# Arguments
+
+- `A`: The parent disk array.
+- `padding`: A tuple of `Int` lower and upper padding tuples, one for each dimension.
+
+# Keywords
+
+- `fill=zero(eltype(A))`: The value to pad the array with.
+"""
+struct PaddedDiskArray{T,N,A<:AbstractArray{T,N},C,F<:T} <: AbstractDiskArray{T,N}
+    parent::A
+    padding::NTuple{N,Tuple{Int,Int}}
+    fill::F
+    chunks::C
+end
+function PaddedDiskArray(A::AbstractArray{T,N}, padding::NTuple{N,Tuple{Int,Int}};
+    fill=zero(eltype(A)),
+) where {T,N}
+    map(padding) do (l, u)
+        (l < 0 || u < 0) && throw(ArgumentError("Padding must be non-negative"))
+    end
+    chunks = GridChunks(map(_pad_offset, eachchunk(A).chunks, padding))
+    PaddedDiskArray(A, padding, fill, chunks)
+end
+
+function _pad_offset(c::RegularChunks, (low, high))
+    chunksize = c.chunksize
+    # Handle lower padding larger than chunksize
+    offset = if low == 0
+        c.offset
+    else
+        c.offset - low + chunksize * (div(low - 1, chunksize) + 1)
+    end
+    size = c.arraysize + low + high
+    return RegularChunks(chunksize, offset, size)
+end
+function _pad_offset(c::IrregularChunks, (low, high))
+    nlowchunks = Int(low > 0) 
+    nhighchunks = Int(high > 0)
+    offsets = Vector{Int}(undef, length(c.offsets) + nlowchunks + nhighchunks)
+    # First offset is always zero
+    offsets[begin] = 0
+    # Increase original offsets by lower padding
+    for (i, o) in enumerate(c.offsets)
+        offsets[i + nlowchunks] = o + low
+    end
+    # Add offset for start of upper padding
+    if nhighchunks > 0
+        offsets[end] = offsets[end-1] + high
+    end
+    return IrregularChunks(offsets)
+end
+
+Base.parent(A::PaddedDiskArray) = A.parent
+function Base.size(A::PaddedDiskArray)
+    map(size(parent(A)), A.padding) do s, (low, high)
+        s + low + high
+    end
+end
+
+haschunks(A::PaddedDiskArray) = haschunks(parent(A))
+eachchunk(A::PaddedDiskArray) = A.chunks
+
+readblock!(A::PaddedDiskArray, data, I::AbstractRange...) =
+    _readblock_padded(A, data, I...)
+writeblock!(A::PaddedDiskArray, data, I...) = 
+    throw(ArgumentError("Cannot write to a PaddedDiskArray"))
+
+function _readblock_padded(A, data, I::AbstractRange...)
+    data .= A.fill
+    Ipadded = map(I, A.padding) do i, (low, high)
+        i .- low
+    end
+    fs = map(axes(parent(A)), Ipadded) do a, ip
+        searchsortedfirst(ip, first(a))
+    end
+    ls = map(axes(parent(A)), Ipadded) do a, ip
+        searchsortedlast(ip, last(a))
+    end
+    return if all(map(<=, fs, ls))
+        Idata = map(:, fs, ls)
+        Iparent = map(getindex, Ipadded, Idata)
+        data[Idata...] .= parent(A)[Iparent...]
+    else
+        # No overlap, don't read
+        data
+    end
+end
+
+"""
+
+    pad(A, padding; fill=zero(eltype(A)))
+
+Pad any `AbstractArray` with fill values, updating chunk patterns.
+
+# Arguments
+
+- `A`: The parent disk array.
+- `padding`: A tuple of `Int` lower and upper padding tuples, one for each dimension.
+
+# Keywords
+
+- `fill=zero(eltype(A))`: The value to pad the array with.
+"""
+pad(A::AbstractArray{<:Any,N}, padding::NTuple{N,Tuple{Int,Int}}; kw...) where N = 
+    PaddedDiskArray(A, padding; kw...)

test/runtests.jl

@@ -973,6 +973,46 @@ end
     end
 end
 
+@testset "Padded disk arrays" begin
+    M = (1:100) * (1:120)'
+    A = cat(M, 2M, 3M, 4M; dims=3)
+    ch = ChunkedDiskArray(A, (128, 128, 2)) 
+    pa = DiskArrays.pad(ch, ((10, 20), (30, 40), (1, 2)); fill=999)
+    @test size(pa) == (130, 190, 7)
+    # All outside
+    @test all(==(999), pa[1:10, 1:30, 1:1])
+    # All inside
+    @test pa[11:20, 31:40, 3:4] == ch[1:10, 1:10, 2:3]
+    @testset "Overlapping lower" begin
+        regions = pa[10:20, 30:40, 1:2]
+        testdata = copy(regions) .= 999
+        testdata[2:11, 2:11, 2:2] .= ch[1:10, 1:10, 1:1]
+        @test regions == testdata
+    end
+    @testset "Overlapping upper" begin
+        regions = pa[101:120, 141:160, 5:7]
+        testdata = copy(regions) .= 999
+        testdata[1:10, 1:10, 1:1] .= ch[91:100, 111:120, 4:4]
+        @test regions == testdata
+    end
+    @testset "offsets larger than chunks" begin
+        pa1 = DiskArrays.pad(ch, ((201, 147), (390, 33), (1, 2)); fill=1)
+        @test sum(pa1) == prod(size(pa1)) + sum(ch) - prod(size(ch))
+    end
+    @testset "_pad_offset" begin
+        c1 = DiskArrays.RegularChunks(10, 0, 100)
+        @test DiskArrays._pad_offset(c1, (5, 2)) == DiskArrays.RegularChunks(10, 5, 107)
+        @test DiskArrays._pad_offset(c1, (30, 2)) == DiskArrays.RegularChunks(10, 0, 132)
+        @test DiskArrays._pad_offset(c1, (10, 10)) == DiskArrays.RegularChunks(10, 0, 120)
+        @test DiskArrays._pad_offset(c1, (0, 0)) == c1
+
+        c2 = DiskArrays.IrregularChunks(chunksizes = [10, 10, 20, 30, 40])
+        #The following test would assume padding ends up in a separate chunk:
+        @test DiskArrays._pad_offset(c2, (5, 5)) == DiskArrays.IrregularChunks(chunksizes = [5, 10, 10, 20, 30, 40, 5])
+        @test DiskArrays._pad_offset(c2, (0, 0)) == c2
+    end
+end
+
 @testset "Range subset identification" begin
     inds = [1, 2, 2, 3, 5, 6, 7, 10, 10]
     readranges, offsets = DiskArrays.find_subranges_sorted(inds, false)