Skip to content

Commit

Permalink
Clean whitespace
Browse files Browse the repository at this point in the history
  • Loading branch information
brendanjohnharris committed Mar 7, 2024
1 parent 9846bfe commit a93d665
Show file tree
Hide file tree
Showing 2 changed files with 126 additions and 7 deletions.
96 changes: 95 additions & 1 deletion src/array/methods.jl
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ end
ds = DD.dims(A, _astuple(dims))
# Run one slice with dimensions to get the transformed dim
d_inds = map(d -> rebuild(d, 1), otherdims(A, ds))
example_dims = length(d_inds) > 0 ? DD.dims(f(view(A, d_inds...))) : ()
example_dims = DD.dims(f(view(A, d_inds...)))
replacement_dims = if isnothing(example_dims) || length(example_dims) != length(ds)
map(d -> rebuild(d, NoLookup()), ds)
else
Expand Down Expand Up @@ -547,6 +547,100 @@ $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}
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()))
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)
end

"""
Base.stack(A::AbstractVector{<:AbstractDimArray}; dims=Pair(ndims(A[1])+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([da, db], dims=3=>Z(1:2))
╭─────────────────────────╮
│ 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])+1, AnonDim()))
if dims isa Integer
dims = dims => AnonDim()
elseif dims isa Dimension
dims = ndims(A[1])+1 => dims
end

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
B = rebuild(B; dims=format(newdims, B))
end
return B
end

function Base.inv(A::AbstractDimArray{T,2}) where T
newdata = inv(parent(A))
newdims = reverse(dims(A))
Expand Down
37 changes: 31 additions & 6 deletions test/methods.jl
Original file line number Diff line number Diff line change
Expand Up @@ -366,12 +366,6 @@ end
@test size(y) == size(dims(y))
@test dims(y) == dims(x[2:9, :, :])
end

@testset "single dim" begin
x = X(1:10)
A = DimArray(ones(10), x)
@test mapslices(sum, A; dims=X) == fill(10.0, X(1))
end
end

@testset "cumsum" begin
Expand Down Expand Up @@ -661,6 +655,37 @@ end
end
end

@testset "Base.stack" begin
a = [1 2 3; 4 5 6]
da = DimArray(a, (X(4.0:5.0), Y(6.0:8.0)))
b = [7 8 9; 10 11 12]
ca = DimArray(b, (X(4.0:5.0), Y(6.0:8.0)))
db = DimArray(b, (X(6.0:7.0), Y(6.0:8.0)))

@test stack([da, db]; dims=3) == stack([parent(da), parent(db)], dims=3)
@test stack([da, db]; dims=3) == stack([da, db]) # Test default dims
@test_warn "Lookup values for X" stack([da, db]; dims=3)

@test stack([da, ca]; dims=1) == stack([parent(da), parent(ca)], dims=1)
@test_warn "Lookup values for X" stack([da, db]; dims=1)

dc = stack([da, ca], dims=Z(1:2))
@test dims(dc, ndims(da)+1) == Z(1:2)
@test parent(dc) == stack(map(parent, [da, db]))

for d = 1:3
dc = stack([da, ca], dims=d)
@test dims(dc, d) isa AnonDim
@test parent(dc) == stack(map(parent, [da, db]), dims=d)
end

for d = 1:3
dc = stack([da, ca], dims=d=>Z(1:2))
@test dims(dc, d) == Z(1:2)
@test parent(dc) == stack(map(parent, [da, db]), dims=d)
end
end

@testset "unique" begin
a = [1 1 6; 1 1 6]
xs = (X, X(), :X)
Expand Down

0 comments on commit a93d665

Please sign in to comment.