-
Notifications
You must be signed in to change notification settings - Fork 41
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
Extending Base.stack for DimArrays #645
base: main
Are you sure you want to change the base?
Changes from 7 commits
4ea5e40
ec3b2cf
0b762ce
6affd9b
ac3affc
83ace60
e6cfb1b
ce16068
bdf151f
7a4cf26
3314852
0ac2eb0
884ab68
c369d06
2c7b8b9
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
@@ -131,15 +131,15 @@ | |||||||||||||||
""" | ||||||||||||||||
function Base.eachslice(A::AbstractDimArray; dims) | ||||||||||||||||
dimtuple = _astuple(dims) | ||||||||||||||||
if !(dimtuple == ()) | ||||||||||||||||
if !(dimtuple == ()) | ||||||||||||||||
all(hasdim(A, dimtuple...)) || throw(DimensionMismatch("A doesn't have all dimensions $dims")) | ||||||||||||||||
end | ||||||||||||||||
_eachslice(A, dimtuple) | ||||||||||||||||
end | ||||||||||||||||
else | ||||||||||||||||
@inline function Base.eachslice(A::AbstractDimArray; dims, drop=true) | ||||||||||||||||
dimtuple = _astuple(dims) | ||||||||||||||||
if !(dimtuple == ()) | ||||||||||||||||
if !(dimtuple == ()) | ||||||||||||||||
all(hasdim(A, dimtuple...)) || throw(DimensionMismatch("A doesn't have all dimensions $dims")) | ||||||||||||||||
end | ||||||||||||||||
_eachslice(A, dimtuple, drop) | ||||||||||||||||
|
@@ -195,7 +195,7 @@ | |||||||||||||||
return rebuild(A, newdata, newdims) | ||||||||||||||||
end | ||||||||||||||||
|
||||||||||||||||
|
||||||||||||||||
Base.cumsum(A::AbstractDimVector) = rebuild(A, Base.cumsum(parent(A))) | ||||||||||||||||
Base.cumsum(A::AbstractDimArray; dims) = rebuild(A, cumsum(parent(A); dims=dimnum(A, dims))) | ||||||||||||||||
Base.cumsum!(B::AbstractArray, A::AbstractDimVector) = cumsum!(B, parent(A)) | ||||||||||||||||
|
@@ -293,7 +293,7 @@ | |||||||||||||||
else | ||||||||||||||||
# vcat the index for the catdim in each of Xin | ||||||||||||||||
joindims = map(A -> dims(A, catdim), Xin) | ||||||||||||||||
if !check_cat_lookups(joindims...) | ||||||||||||||||
if !check_cat_lookups(joindims...) | ||||||||||||||||
return rebuild(catdim, NoLookup()) | ||||||||||||||||
end | ||||||||||||||||
_vcat_dims(joindims...) | ||||||||||||||||
|
@@ -367,7 +367,7 @@ | |||||||||||||||
(catdim,) | ||||||||||||||||
else | ||||||||||||||||
# Make sure this is the same dimension for all arrays | ||||||||||||||||
if !comparedims(Bool, map(x -> dims(x, 2), As)...; | ||||||||||||||||
if !comparedims(Bool, map(x -> dims(x, 2), As)...; | ||||||||||||||||
val=true, warn = " Can't `vcat` AbstractDimArray, applying to `parent` object." | ||||||||||||||||
) | ||||||||||||||||
return Base.vcat(map(parent, As)...) | ||||||||||||||||
|
@@ -542,11 +542,104 @@ | |||||||||||||||
_span_string(D, S, span) = _cat_warn_string(D, "not all lookups have `$S` spans. Found $(basetypeof(span))") | ||||||||||||||||
_cat_warn_string(D, message) = """ | ||||||||||||||||
`cat` cannot concatenate `Dimension`s, falling back to `parent` type: | ||||||||||||||||
$message on dimension $D. | ||||||||||||||||
$message on dimension $D. | ||||||||||||||||
|
||||||||||||||||
To fix for `AbstractDimArray`, pass new lookup values as `cat(As...; dims=$D(newlookupvals))` keyword or `dims=$D()` for empty `NoLookup`. | ||||||||||||||||
""" | ||||||||||||||||
|
||||||||||||||||
function Base._typed_stack(::Colon, ::Type{T}, ::Type{S}, A, Aax=_iterator_axes(A)) where {T,S<:AbstractDimArray} | ||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. How stable do you think these methods are... could we add a method to I know we use internals elsewhere, but we should stop: |
||||||||||||||||
origdims = map(dims, A) | ||||||||||||||||
_A = parent.(A) | ||||||||||||||||
t = eltype(_A) | ||||||||||||||||
_A = Base._typed_stack(:, T, t, A) | ||||||||||||||||
|
||||||||||||||||
if !comparedims(Bool, origdims...; | ||||||||||||||||
order=true, val=true, warn=" Can't `stack` AbstractDimArray, applying to `parent` object." | ||||||||||||||||
) | ||||||||||||||||
return _A | ||||||||||||||||
else | ||||||||||||||||
DimArray(_A, (first(origdims)..., AnonDim())) | ||||||||||||||||
rafaqz marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||||||
end | ||||||||||||||||
end | ||||||||||||||||
|
||||||||||||||||
function Base._dim_stack(newdim::Integer, ::Type{T}, ::Type{S}, A) where {T,S<:AbstractDimArray} | ||||||||||||||||
origdims = dims.(A) | ||||||||||||||||
_A = parent.(A) | ||||||||||||||||
t = eltype(_A) | ||||||||||||||||
_A = Base._dim_stack(newdim, T, t, A) | ||||||||||||||||
|
||||||||||||||||
if !comparedims(Bool, origdims...; | ||||||||||||||||
order=true, val=true, warn=" Can't `stack` AbstractDimArray, applying to `parent` object." | ||||||||||||||||
) | ||||||||||||||||
return _A | ||||||||||||||||
end | ||||||||||||||||
|
||||||||||||||||
newdims = first(origdims) | ||||||||||||||||
newdims = ntuple(length(newdims) + 1) do d | ||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is this type-stable? |
||||||||||||||||
if d == newdim | ||||||||||||||||
AnonDim() | ||||||||||||||||
else # Return the old dimension, shifted across once if it comes after the new dim | ||||||||||||||||
newdims[d-(d>newdim)] | ||||||||||||||||
end | ||||||||||||||||
end | ||||||||||||||||
DimArray(_A, newdims) | ||||||||||||||||
rafaqz marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||||||
end | ||||||||||||||||
|
||||||||||||||||
""" | ||||||||||||||||
Base.stack(A::AbstractVector{<:AbstractDimArray}; dims=Pair(ndims(A[1]), AnonDim())) | ||||||||||||||||
|
||||||||||||||||
Stack arrays along a new axis while preserving the dimensional information of other axes. | ||||||||||||||||
|
||||||||||||||||
The optional keyword argument `dims` has the following behavior: | ||||||||||||||||
- `dims isa Integer`: The dimension of the new axis is an `AnonDim` at position `dims` | ||||||||||||||||
- `dims isa Dimension`: The new axis is at `ndims(A[1])+1` and has a dimension of `dims`. | ||||||||||||||||
- `dims isa Pair{Integer, Dimension}`: The new axis is at `first(dims)` and has a dimension | ||||||||||||||||
of `last(dims)`. | ||||||||||||||||
|
||||||||||||||||
If `dims` contains a `Dimension`, that `Dimension` must have the same length as A. | ||||||||||||||||
|
||||||||||||||||
# Examples | ||||||||||||||||
```julia-repl | ||||||||||||||||
julia> da = DimArray([1 2 3; 4 5 6], (X(10:10:20), Y(300:-100:100))); | ||||||||||||||||
julia> db = DimArray([6 5 4; 3 2 1], (X(10:10:20), Y(300:-100:100))); | ||||||||||||||||
|
||||||||||||||||
# Stack along a new dimension `Z` | ||||||||||||||||
julia> dc = stack(Z(1:2), [da, db], dims=3) | ||||||||||||||||
╭─────────────────────────╮ | ||||||||||||||||
│ 2×3×2 DimArray{Int64,3} │ | ||||||||||||||||
├─────────────────────────┴──────────────────────────────── dims ┐ | ||||||||||||||||
↓ X Sampled{Int64} 10:10:20 ForwardOrdered Regular Points, | ||||||||||||||||
→ Y Sampled{Int64} 300:-100:100 ReverseOrdered Regular Points, | ||||||||||||||||
↗ Z 1:2 | ||||||||||||||||
└────────────────────────────────────────────────────────────────┘ | ||||||||||||||||
|
||||||||||||||||
julia> dims(dc, 3) == Z(1:2) | ||||||||||||||||
true | ||||||||||||||||
julia> parent(dc) == stack(map(parent, [da, db]), dims=3) | ||||||||||||||||
true | ||||||||||||||||
``` | ||||||||||||||||
""" | ||||||||||||||||
function Base.stack(A::AbstractVector{<:AbstractDimArray}; dims=Pair(ndims(A[1]), AnonDim()), kwargs...) | ||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
Isn't the default |
||||||||||||||||
dims isa Integer && (dims = dims => AnonDim()) | ||||||||||||||||
dims isa Dimension && (dims = ndims(A[1])+1 => dims) | ||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
Assigning after |
||||||||||||||||
|
||||||||||||||||
B = Base._stack(first(dims), A) | ||||||||||||||||
|
||||||||||||||||
if B isa AbstractDimArray | ||||||||||||||||
newdims = ntuple(ndims(B)) do d | ||||||||||||||||
if d == first(dims) # Use the new provided dimension | ||||||||||||||||
last(dims) | ||||||||||||||||
else | ||||||||||||||||
DimensionalData.dims(B, d) | ||||||||||||||||
end | ||||||||||||||||
end | ||||||||||||||||
newdims = format(newdims, B) | ||||||||||||||||
|
||||||||||||||||
B = rebuild(B; dims=newdims) | ||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
User specified dims are usually incomplete and possibly incorrect |
||||||||||||||||
end | ||||||||||||||||
return B | ||||||||||||||||
end | ||||||||||||||||
|
||||||||||||||||
function Base.inv(A::AbstractDimArray{T,2}) where T | ||||||||||||||||
newdata = inv(parent(A)) | ||||||||||||||||
newdims = reverse(dims(A)) | ||||||||||||||||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Would be nice to remove all these whitespace changes so we can see the real changes