Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 21 additions & 15 deletions src/permute.jl
Original file line number Diff line number Diff line change
@@ -1,46 +1,52 @@
"""
PermutedDiskArray <: AbstractDiskArray
AbstractPermutedDiskArray <: AbstractDiskArray

Abstract supertype for diskarray with permuted dimensions.
"""
abstract type AbstractPermutedDiskArray{T,N,P<:PermutedDimsArray{T,N}} <: AbstractDiskArray{T,N} end

"""
PermutedDiskArray <: AbstractPermutedDiskArray

A lazily permuted disk array returned by `permutedims(diskarray, permutation)`.
"""
struct PermutedDiskArray{T,N,P<:PermutedDimsArray{T,N}} <: AbstractDiskArray{T,N}
struct PermutedDiskArray{T,N,P<:PermutedDimsArray{T,N}} <: AbstractPermutedDiskArray{T,N,P}
a::P
end

# Base methods

Base.size(a::PermutedDiskArray) = size(a.a)

Base.size(a::AbstractPermutedDiskArray) = size(a.a)
Base.parent(a::AbstractPermutedDiskArray) = a.a.parent
# DiskArrays interface

haschunks(a::PermutedDiskArray) = haschunks(a.a.parent)
function eachchunk(a::PermutedDiskArray)
haschunks(a::AbstractPermutedDiskArray) = haschunks(parent(a))
function eachchunk(a::AbstractPermutedDiskArray)
# Get the parent chunks
gridchunks = eachchunk(a.a.parent)
perm = _getperm(a.a)
gridchunks = eachchunk(parent(a))
perm = _getperm(a)
# Return permuted GridChunks
return GridChunks(genperm(gridchunks.chunks, perm)...)
end
function DiskArrays.readblock!(a::PermutedDiskArray, aout, i::OrdinalRange...)
function DiskArrays.readblock!(a::AbstractPermutedDiskArray, aout, i::OrdinalRange...)
iperm = _getiperm(a)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is now an internal method in an Abstract typed function... Should we remove the underscore?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove which underscore? You mean in _getiperm and _getperm

# Permute the indices
inew = genperm(i, iperm)
# Permute the dest block and read from the true parent
DiskArrays.readblock!(a.a.parent, PermutedDimsArray(aout, iperm), inew...)
DiskArrays.readblock!(parent(a), PermutedDimsArray(aout, iperm), inew...)
return nothing
end
function DiskArrays.writeblock!(a::PermutedDiskArray, v, i::OrdinalRange...)
function DiskArrays.writeblock!(a::AbstractPermutedDiskArray, v, i::OrdinalRange...)
iperm = _getiperm(a)
inew = genperm(i, iperm)
# Permute the dest block and write from the true parent
DiskArrays.writeblock!(a.a.parent, PermutedDimsArray(v, iperm), inew...)
DiskArrays.writeblock!(parent(a), PermutedDimsArray(v, iperm), inew...)
return nothing
end

_getperm(a::PermutedDiskArray) = _getperm(a.a)
_getperm(a::AbstractPermutedDiskArray) = _getperm(a.a)
_getperm(::PermutedDimsArray{<:Any,<:Any,perm}) where {perm} = perm

_getiperm(a::PermutedDiskArray) = _getiperm(a.a)
_getiperm(a::AbstractPermutedDiskArray) = _getiperm(a.a)
_getiperm(::PermutedDimsArray{<:Any,<:Any,<:Any,iperm}) where {iperm} = iperm

# Implementaion macros
Expand Down
25 changes: 16 additions & 9 deletions src/subarray.jl
Original file line number Diff line number Diff line change
@@ -1,32 +1,39 @@
"""
SubDiskArray <: AbstractDiskArray

Abstract supertype for a view of an AbstractDiskArray
"""
abstract type AbstractSubDiskArray{T,N,P,I,L} <: AbstractDiskArray{T,N} end

"""
SubDiskArray <: AbstractDiskArray

A replacement for `Base.SubArray` for disk arrays, returned by `view`.
"""
struct SubDiskArray{T,N,P,I,L} <: AbstractDiskArray{T,N}
struct SubDiskArray{T,N,P,I,L} <: AbstractSubDiskArray{T,N,P,I,L}
v::SubArray{T,N,P,I,L}
end

# Base methods
Base.view(a::SubDiskArray, i...) = SubDiskArray(view(a.v, i...))
Base.view(a::SubDiskArray, i::CartesianIndices) = view(a, i.indices...)
Base.size(a::SubDiskArray) = size(a.v)
Base.parent(a::SubDiskArray) = a.v.parent
Base.view(a::AbstractSubDiskArray, i...) = SubDiskArray(view(a.v, i...))
Base.view(a::AbstractSubDiskArray, i::CartesianIndices) = view(a, i.indices...)

Check warning on line 19 in src/subarray.jl

View check run for this annotation

Codecov / codecov/patch

src/subarray.jl#L19

Added line #L19 was not covered by tests
Base.size(a::AbstractSubDiskArray) = size(a.v)
Base.parent(a::AbstractSubDiskArray) = a.v.parent

Check warning on line 21 in src/subarray.jl

View check run for this annotation

Codecov / codecov/patch

src/subarray.jl#L21

Added line #L21 was not covered by tests

_replace_colon(s, ::Colon) = Base.OneTo(s)
_replace_colon(s, r) = r

# Diskarrays.jl interface
function readblock!(a::SubDiskArray, aout, i::OrdinalRange...)
function readblock!(a::AbstractSubDiskArray, aout, i::OrdinalRange...)
pinds = parentindices(view(a.v, i...))
getindex_disk!(aout, parent(a.v), pinds...)
end
function writeblock!(a::SubDiskArray, v, i::OrdinalRange...)
function writeblock!(a::AbstractSubDiskArray, v, i::OrdinalRange...)
pinds = parentindices(view(a.v, i...))
setindex_disk!(parent(a.v), v, pinds...)
end
haschunks(a::SubDiskArray) = haschunks(parent(a.v))
eachchunk(a::SubDiskArray) = eachchunk_view(haschunks(a.v.parent), a.v)
haschunks(a::AbstractSubDiskArray) = haschunks(parent(a.v))
eachchunk(a::AbstractSubDiskArray) = eachchunk_view(haschunks(a.v.parent), a.v)

function eachchunk_view(::Chunked, vv)
pinds = parentindices(vv)
Expand Down
Loading