FR: Make broadcast_dims work with DimArrays and Dimensions #748
Description
Recording this idea from slack discussions in github.
It would be nice if I could do something like
broadcast_dims(some_function, A, X, Ti)where A is a DimArray with dimensions (a subset of) (X, Ti), and some_function(,xi,ti) is a function that takes (here) three arguments, where a is an elements from A, and xi, ti are elements from the lookup of each dimension.
Currently, you can do something like:
lon, ts = X(1:3), Ti(1:5)
A = ones(lon, ts)
broadcast_dims(+, A, DimArray(parent(lon), lon))but it would be neat to do something like
broadcast_dims(+, A, X)The most naive implementation I could come up with looks like this:
function broadcast_dims2(f, As::Union{DimensionalData.AbstractBasicDimArray, DimensionalData.Dimensions.Dimension, Type{<:DimensionalData.Dimension}}...)
# have to look up dims for any actual DimArrays first if support for `X`, `Ti`, etc, as input should work (because we need the lookup array)
existing_dims = DimensionalData.combinedims(filter(A -> A isa DimensionalData.AbstractBasicDimArray, As)...)
Bs = map(As) do A
if A isa DimensionalData.Dimension
DimArray(parent(A), A)
elseif A isa Type{<:DimensionalData.Dimension}
dim = dims(existing_dims, A)
DimArray(parent(dim), dim)
else
A
end
end
broadcast_dims(f, Bs...)
endI had to do it this way because e.g. combinedims(A, lon) gives a StackOverflowError for me.
With this, I can use both the Dimension type (e.g. X) or a reference to a concrete dimension (e.g. lon)
broadcast_dims2(+, A, X, Ti)
╭─────────────────────────╮
│ 3×5 DimArray{Float64,2} │
├─────────────────────────┴──────────────────────── dims ┐
↓ X Sampled{Int64} 1:3 ForwardOrdered Regular Points,
→ Ti Sampled{Int64} 1:5 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────┘
↓ → 1 2 3 4 5
1 3.0 4.0 5.0 6.0 7.0
2 4.0 5.0 6.0 7.0 8.0
3 5.0 6.0 7.0 8.0 9.0