Extending Base.stack for DimArrays #645
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 == ()) | ||||||
| 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 == ()) | ||||||
| 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...) | ||||||
| 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)...; | ||||||
| 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. | ||||||
|
|
||||||
| 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. 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 | ||||||
|
|
||||||
| 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...) | ||||||
| dims isa Integer && (dims = dims => AnonDim()) | ||||||
| dims isa Dimension && (dims = ndims(A[1])+1 => dims) | ||||||
rafaqz marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||
|
|
||||||
| 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.
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)) | ||||||
|
|
||||||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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