Skip to content

add abstract types for views and permuted arrays #255

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
May 31, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
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
1 change: 1 addition & 0 deletions src/DiskArrays.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
module DiskArrays

import ConstructionBase
import Base.PermutedDimsArrays: genperm

using LRUCache: LRUCache, LRU

Expand Down
19 changes: 13 additions & 6 deletions src/permute.jl
Original file line number Diff line number Diff line change
@@ -1,9 +1,16 @@
"""
PermutedDiskArray <: AbstractDiskArray
AbstractPermutedDiskArray <: AbstractDiskArray

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

"""
PermutedDiskArray <: AbstractPermutedDiskArray

A lazily permuted disk array returned by `permutedims(diskarray, permutation)`.
"""
struct PermutedDiskArray{T,N,perm,iperm,A<:AbstractArray{T,N}} <: AbstractDiskArray{T,N}
struct PermutedDiskArray{T,N,perm,iperm,A<:AbstractArray{T,N}} <: AbstractPermutedDiskArray{T,N,perm,iperm,A}
parent::A
end
# We use PermutedDimsArray internals instead of duplicating them,
Expand Down Expand Up @@ -41,24 +48,24 @@ function eachchunk(a::PermutedDiskArray)
# 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!(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!(parent(a), PermutedDimsArray(v, iperm), inew...)
return nothing
end

_getperm(::PermutedDiskArray{<:Any,<:Any,perm}) where {perm} = perm
_getiperm(::PermutedDiskArray{<:Any,<:Any,<:Any,iperm}) where {iperm} = iperm
_getperm(::AbstractPermutedDiskArray{<:Any,<:Any,perm}) where {perm} = perm
_getiperm(::AbstractPermutedDiskArray{<:Any,<:Any,<:Any,iperm}) where {iperm} = iperm

# Implementation macro

Expand Down
40 changes: 23 additions & 17 deletions src/reshape.jl
Original file line number Diff line number Diff line change
@@ -1,39 +1,45 @@
import Base: _throw_dmrs
import Base.PermutedDimsArrays: genperm

"""
ReshapedDiskArray <: AbstractDiskArray
AbstractReshapedDiskArray <: AbstractDiskArray

Abstract supertype for a replacements of `Base.ReshapedArray` for `AbstractDiskArray`s`
"""
abstract type AbstractReshapedDiskArray{T,N,P,M} <: AbstractDiskArray{T,N} end

"""
ReshapedDiskArray <: AbstractReshapedDiskArray

A replacement for `Base.ReshapedArray` for disk arrays,
returned by `reshape`.


Reshaping is really not trivial, because the access pattern would
completely change for reshaped arrays, rectangles would not remain
rectangles in the parent array.

However, we can support the case where only singleton dimensions are added,
later we could allow more special cases like joining two dimensions to one
"""
struct ReshapedDiskArray{T,N,P<:AbstractArray{T},M} <: AbstractDiskArray{T,N}
struct ReshapedDiskArray{T,N,P<:AbstractArray{T},M} <: AbstractReshapedDiskArray{T,N,P,M}
parent::P
keepdim::NTuple{M,Int}
newsize::NTuple{N,Int}
end

# Base methods
Base.size(r::AbstractReshapedDiskArray) = r.newsize
Base.parent(r::AbstractReshapedDiskArray) = r.parent

Base.parent(r::ReshapedDiskArray) = r.parent
Base.size(r::ReshapedDiskArray) = r.newsize
keepdim(r::AbstractReshapedDiskArray) = r.keepdim

# DiskArrays interface

haschunks(a::ReshapedDiskArray) = haschunks(a.parent)
function eachchunk(a::ReshapedDiskArray{<:Any,N}) where {N}
pchunks = eachchunk(a.parent)
haschunks(a::AbstractReshapedDiskArray) = haschunks(parent(a))
function eachchunk(a::AbstractReshapedDiskArray{<:Any,N}) where {N}
pchunks = eachchunk(parent(a))
inow::Int = 0
outchunks = ntuple(N) do idim
if in(idim, a.keepdim)
if in(idim, keepdim(a))
inow += 1
pchunks.chunks[inow]
else
Expand All @@ -42,14 +48,14 @@
end
return GridChunks(outchunks...)
end
function DiskArrays.readblock!(a::ReshapedDiskArray, aout, i::OrdinalRange...)
inew = tuple_tuple_getindex(i, a.keepdim)
DiskArrays.readblock!(a.parent, reshape(aout, map(length, inew)), inew...)
function DiskArrays.readblock!(a::AbstractReshapedDiskArray, aout, i::OrdinalRange...)
inew = tuple_tuple_getindex(i, keepdim(a))
DiskArrays.readblock!(parent(a), reshape(aout, map(length, inew)), inew...)
return nothing
end
function DiskArrays.writeblock!(a::ReshapedDiskArray, v, i::OrdinalRange...)
inew = tuple_tuple_getindex(i, a.keepdim)
DiskArrays.writeblock!(a.parent, reshape(v, map(length, inew)), inew...)
function DiskArrays.writeblock!(a::AbstractReshapedDiskArray, v, i::OrdinalRange...)
inew = tuple_tuple_getindex(i, keepdim(a))
DiskArrays.writeblock!(parent(a), reshape(v, map(length, inew)), inew...)
return nothing
end
function reshape_disk(parent, dims)
Expand Down Expand Up @@ -93,6 +99,6 @@
end

# For ambiguity
function Base._reshape(A::DiskArrays.AbstractDiskArray{<:Any,1}, dims::Tuple{Int64})
function Base._reshape(A::AbstractDiskArray{<:Any,1}, dims::Tuple{Int64})

Check warning on line 102 in src/reshape.jl

View check run for this annotation

Codecov / codecov/patch

src/reshape.jl#L102

Added line #L102 was not covered by tests
return reshape_disk(A, dims)
end
38 changes: 25 additions & 13 deletions src/subarray.jl
Original file line number Diff line number Diff line change
@@ -1,32 +1,44 @@
"""
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
subarray(a::SubDiskArray) = a.v
function Base.view(a::T, i...) where T<:AbstractSubDiskArray
basetype = Base.typename(T).wrapper
basetype(view(subarray(a), i...))
end
Base.view(a::AbstractSubDiskArray, i::CartesianIndices) = view(a, i.indices...)

Check warning on line 23 in src/subarray.jl

View check run for this annotation

Codecov / codecov/patch

src/subarray.jl#L23

Added line #L23 was not covered by tests
Base.size(a::AbstractSubDiskArray) = size(subarray(a))
Base.parent(a::AbstractSubDiskArray) = parent(subarray(a))
Base.parentindices(a::AbstractSubDiskArray) = parentindices(subarray(a))

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

# Diskarrays.jl interface
function readblock!(a::SubDiskArray, aout, i::OrdinalRange...)
pinds = parentindices(view(a.v, i...))
getindex_disk!(aout, parent(a.v), pinds...)
function readblock!(a::AbstractSubDiskArray, aout, i::OrdinalRange...)
pinds = parentindices(view(a, i...))
getindex_disk!(aout, parent(a), pinds...)
end
function writeblock!(a::SubDiskArray, v, i::OrdinalRange...)
pinds = parentindices(view(a.v, i...))
setindex_disk!(parent(a.v), v, pinds...)
function writeblock!(a::AbstractSubDiskArray, v, i::OrdinalRange...)
pinds = parentindices(view(a, i...))
setindex_disk!(parent(a), 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))
eachchunk(a::AbstractSubDiskArray) = eachchunk_view(haschunks(parent(a)), a)

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