Skip to content

Commit

Permalink
allow DD keywords in similar (#844)
Browse files Browse the repository at this point in the history
* allow DD keywords in similar

* test similar kw

* bugfix and standardise

* typo

* bugfix
  • Loading branch information
rafaqz authored Nov 5, 2024
1 parent 978f838 commit 175d54e
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 48 deletions.
119 changes: 73 additions & 46 deletions src/array/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ Base.Array{T}(x::UndefInitializer, dims::DimTuple; kw...) where T = Array{T}(x,
function Base.NamedTuple(A1::AbstractDimArray, As::AbstractDimArray...)
arrays = (A1, As...)
keys = map(Symbol name, arrays)
NamedTuple{keys}(arrays)
return NamedTuple{keys}(arrays)
end

# undef constructor for all AbstractDimArray
Expand All @@ -140,27 +140,43 @@ Base.read(A::AbstractDimArray) = A
# Methods that create copies of an AbstractDimArray #######################################

# Need to cover a few type signatures to avoid ambiguity with base
Base.similar(A::AbstractDimArray) =
rebuild(A; data=similar(parent(A)), dims=dims(A), refdims=refdims(A), name=_noname(A), metadata=metadata(A))
Base.similar(A::AbstractDimArray, ::Type{T}) where T =
rebuild(A; data=similar(parent(A), T), dims=dims(A), refdims=refdims(A), name=_noname(A), metadata=metadata(A))
function Base.similar(A::AbstractDimArray;
data=similar(parent(A)),
dims=dims(A), refdims=refdims(A), name=_noname(A), metadata=metadata(A), kw...
)
rebuild(A; dims=format(dims, data), refdims, name, metadata, kw...)
end
function Base.similar(A::AbstractDimArray, ::Type{T};
data=similar(parent(A), T),
dims=dims(A), refdims=refdims(A), name=_noname(A), metadata=metadata(A), kw...
) where T
rebuild(A; data, dims=format(dims, data), refdims, name, metadata, kw...)
end

# We avoid calling `parent` for AbstractBasicDimArray as we don't know what it is/if there is one
Base.similar(A::AbstractBasicDimArray{T,N}) where {T,N} =
rebuild(A; data=similar(Array{T,N}, size(A)), dims=dims(A), refdims=refdims(A), name=_noname(A), metadata=metadata(A))
Base.similar(A::AbstractBasicDimArray{<:Any,N}, ::Type{T}) where {T,N} =
rebuild(A; data=similar(Array{T,N}, size(A)), dims=dims(A), refdims=refdims(A), name=_noname(A), metadata=metadata(A))
function Base.similar(A::AbstractBasicDimArray{T,N};
data=similar(Array{T,N}, size(A)),
dims=dims(A), refdims=refdims(A), name=_noname(A), metadata=NoMetadata(), kw...
) where {T,N}
dimconstructor(dims)(data, dims; refdims, name, metadata, kw...)
end
function Base.similar(A::AbstractBasicDimArray{<:Any,N}, ::Type{T};
data=similar(Array{T,N}, size(A)),
dims=dims(A), refdims=refdims(A), name=_noname(A), metadata=NoMetadata(), kw...
) where {T,N}
dimconstructor(dims)(data, dims; refdims, name, metadata, kw...)
end
# We can't resize the dims or add missing dims, so return the unwraped Array type?
# An alternative would be to fill missing dims with `Anon`, and keep existing
# dims but strip the Lookup? It just seems a little complicated when the methods
# below using DimTuple work better anyway.
Base.similar(A::AbstractDimArray, i::Integer, I::Vararg{Integer}) =
similar(A, eltype(A), (i, I...))
Base.similar(A::AbstractDimArray, I::Tuple{Int,Vararg{Int}}) =
similar(A, eltype(A), I)
Base.similar(A::AbstractDimArray, ::Type{T}, i::Integer, I::Vararg{Integer}) where T =
similar(A, T, (i, I...))
Base.similar(A::AbstractDimArray, ::Type{T}, I::Tuple{Int,Vararg{Int}}) where T =
Base.similar(A::AbstractDimArray, i::Integer, I::Vararg{Integer}; kw...) =
similar(A, eltype(A), (i, I...); kw...)
Base.similar(A::AbstractDimArray, I::Tuple{Int,Vararg{Int}}; kw...) =
similar(A, eltype(A), I; kw...)
Base.similar(A::AbstractDimArray, ::Type{T}, i::Integer, I::Vararg{Integer}; kw...) where T =
similar(A, T, (i, I...); kw...)
Base.similar(A::AbstractDimArray, ::Type{T}, I::Tuple{Int,Vararg{Int}}; kw...) where T =
similar(parent(A), T, I)

const MaybeDimUnitRange = Union{Integer,Base.OneTo,Dimensions.DimUnitRange}
Expand All @@ -172,83 +188,91 @@ const MaybeDimUnitRange = Union{Integer,Base.OneTo,Dimensions.DimUnitRange}
for s1 in (:(Dimensions.DimUnitRange), :MaybeDimUnitRange)
s1 === :MaybeDimUnitRange || @eval begin
function Base.similar(
A::AbstractArray, ::Type{T}, shape::Tuple{$s1,Vararg{MaybeDimUnitRange}}
A::AbstractArray, ::Type{T}, shape::Tuple{$s1,Vararg{MaybeDimUnitRange}}; kw...
) where T
_similar(A, T, shape)
_similar(A, T, shape; kw...)
end
function Base.similar(
::Type{T}, shape::Tuple{$s1,Vararg{MaybeDimUnitRange}}
::Type{T}, shape::Tuple{$s1,Vararg{MaybeDimUnitRange}}; kw...
) where T<:AbstractArray
_similar(T, shape)
_similar(T, shape; kw...)
end
end
for s2 in (:(Dimensions.DimUnitRange), :MaybeDimUnitRange)
all(Base.Fix2(===, :MaybeDimUnitRange), (s1, s2)) || @eval begin
function Base.similar(
A::AbstractArray, T::Type, shape::Tuple{$s1,$s2,Vararg{MaybeDimUnitRange}}
A::AbstractArray, T::Type, shape::Tuple{$s1,$s2,Vararg{MaybeDimUnitRange}}; kw...
)
_similar(A, T, shape)
_similar(A, T, shape; kw...)
end
function Base.similar(
T::Type{<:AbstractArray}, shape::Tuple{$s1,$s2,Vararg{MaybeDimUnitRange}}
T::Type{<:AbstractArray}, shape::Tuple{$s1,$s2,Vararg{MaybeDimUnitRange}}; kw...
)
_similar(T, shape)
_similar(T, shape; kw...)
end
end
for s3 in (:(Dimensions.DimUnitRange), :MaybeDimUnitRange)
all(Base.Fix2(===, :MaybeDimUnitRange), (s1, s2, s3)) || @eval begin
function Base.similar(
A::AbstractArray, T::Type, shape::Tuple{$s1,$s2,$s3,Vararg{MaybeDimUnitRange}}
A::AbstractArray, T::Type, shape::Tuple{$s1,$s2,$s3,Vararg{MaybeDimUnitRange}}; kw...
)
_similar(A, T, shape)
_similar(A, T, shape; kw...)
end
function Base.similar(
T::Type{<:AbstractArray}, shape::Tuple{$s1,$s2,$s3,Vararg{MaybeDimUnitRange}}
T::Type{<:AbstractArray}, shape::Tuple{$s1,$s2,$s3,Vararg{MaybeDimUnitRange}}; kw...
)
_similar(T, shape)
_similar(T, shape; kw...)
end
end
for s4 in (:(Dimensions.DimUnitRange), :MaybeDimUnitRange)
all(Base.Fix2(===, :MaybeDimUnitRange), (s1, s2, s3, s4)) && continue
@eval begin
function Base.similar(
A::AbstractArray, T::Type, shape::Tuple{$s1,$s2,$s3,$s4,Vararg{MaybeDimUnitRange}}
A::AbstractArray, T::Type, shape::Tuple{$s1,$s2,$s3,$s4,Vararg{MaybeDimUnitRange}}; kw...
)
_similar(A, T, shape)
_similar(A, T, shape; kw...)
end
function Base.similar(
T::Type{<:AbstractArray}, shape::Tuple{$s1,$s2,$s3,$s4,Vararg{MaybeDimUnitRange}}
T::Type{<:AbstractArray}, shape::Tuple{$s1,$s2,$s3,$s4,Vararg{MaybeDimUnitRange}}; kw...
)
_similar(T, shape)
_similar(T, shape; kw...)
end
end
end
end
end
end
function _similar(A::AbstractArray, T::Type, shape::Tuple)
function _similar(A::AbstractArray, T::Type, shape::Tuple; kw...)
data = similar(parent(A), T, map(_parent_range, shape))
shape isa Tuple{Vararg{Dimensions.DimUnitRange}} || return data
C = dimconstructor(dims(shape))
return C(data, dims(shape))
return C(data, dims(shape); kw...)
end
function _similar(::Type{T}, shape::Tuple) where {T<:AbstractArray}
function _similar(::Type{T}, shape::Tuple; kw...) where {T<:AbstractArray}
data = similar(T, map(_parent_range, shape))
shape isa Tuple{Vararg{Dimensions.DimUnitRange}} || return data
C = dimconstructor(dims(shape))
return C(data, dims(shape))
return C(data, dims(shape); kw...)
end
_parent_range(r::Dimensions.DimUnitRange) = parent(r)
_parent_range(r) = r

# With Dimensions we can return an `AbstractDimArray`
Base.similar(A::AbstractBasicDimArray, D::DimTuple) = Base.similar(A, eltype(A), D)
Base.similar(A::AbstractBasicDimArray, D::Dimension...) = Base.similar(A, eltype(A), D)
Base.similar(A::AbstractBasicDimArray, ::Type{T}, D::Dimension...) where T =
Base.similar(A, T, D)
Base.similar(A::AbstractDimArray, ::Type{T}, D::DimTuple) where T =
rebuild(A; data=similar(parent(A), T, size(D)), dims=D, refdims=(), metadata=NoMetadata())
Base.similar(A::AbstractDimArray, ::Type{T}, D::Tuple{}) where T =
rebuild(A; data=similar(parent(A), T, ()), dims=(), refdims=(), metadata=NoMetadata())
Base.similar(A::AbstractBasicDimArray, D::DimTuple; kw...) = Base.similar(A, eltype(A), D; kw...)
Base.similar(A::AbstractBasicDimArray, D::Dimension...; kw...) = Base.similar(A, eltype(A), D; kw...)
Base.similar(A::AbstractBasicDimArray, ::Type{T}, D::Dimension...; kw...) where T =
Base.similar(A, T, D; kw...)
function Base.similar(A::AbstractDimArray, ::Type{T}, D::DimTuple;
refdims=(), name=_noname(A), metadata=NoMetadata(), kw...
) where T
data = similar(parent(A), T, _dimlength(D))
dims = _maybestripval(D)
return rebuild(A; data, dims, refdims, metadata, name, kw...)
end
function Base.similar(A::AbstractDimArray, ::Type{T}, D::Tuple{};
refdims=(), name=_noname(A), metadata=NoMetadata(), kw...
) where T
data = similar(parent(A), T, ())
rebuild(A; data, dims=(), refdims, metadata, name, kw...)
end

# Keep the same type in `similar`
_noname(A::AbstractBasicDimArray) = _noname(name(A))
Expand All @@ -257,6 +281,9 @@ _noname(::NoName) = NoName()
_noname(::Symbol) = Symbol("")
_noname(name::Name) = name # Keep the name so the type doesn't change

_parent_range(r::Dimensions.DimUnitRange) = parent(r)
_parent_range(r) = r

for func in (:copy, :one, :oneunit, :zero)
@eval begin
(Base.$func)(A::AbstractDimArray; kw...) = rebuild(A; data=($func)(parent(A)), kw...)
Expand Down
19 changes: 17 additions & 2 deletions test/array.jl
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
using DimensionalData, Test, Unitful, SparseArrays, Dates, Random
using DimensionalData, Test , Unitful, SparseArrays, Dates, Random
using DimensionalData: layerdims, checkdims
using LinearAlgebra

Expand Down Expand Up @@ -131,10 +131,25 @@ end
@test size(da_sim) == size(da)
@test dims(da_sim) === dims(da)
@test refdims(da_sim) === refdims(da)
@test refdims(da_sim) === refdims(da)
@test metadata(da_sim) === metadata(da)
end

@testset "similar with keywords" begin
for x in (da, DimPoints(da))
md = Dict(:new_meta => "b")
rd = (format(Ti(1:1)),)
da_named = similar(x; name=:new_name, metadata=md, refdims=rd)
@test name(da_named) === :new_name
@test metadata(da_named) === md
@test refdims(da_named) === rd
da_float_named = similar(x, Float64; name=:new_name, metadata=md, refdims=rd)
@test eltype(da_float_named) === Float64
@test name(da_float_named) === :new_name
@test metadata(da_float_named) === md
@test refdims(da_float_named) === rd
end
end

@testset "similar with a type" begin
da_float = @inferred similar(da, Float64)
@test eltype(da_float) == Float64
Expand Down

0 comments on commit 175d54e

Please sign in to comment.