FR: Make `broadcast_dims` work with `DimArray`s and `Dimension`s

[haakon-e]

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...)
end

I 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

[rafaqz]

Probably we should allow Dimensions, dimensions types and symbols.

And maybe, if the dimensions holds a lookup, we use it as is, if it's a colon we use the dim from the array.

Symbols work the same as types

[haakon-e]

The following modification to the above code should support that:

# setup
julia> x, y = X(1:3), Y(10:13)
↓ X 1:3,
→ Y 10:13

julia> A = ones(x,y)
╭─────────────────────────╮
│ 3×4 DimArray{Float64,2} │
├─────────────────────────┴──────────────────────── dims ┐
  ↓ X Sampled{Int64} 1:3 ForwardOrdered Regular Points,
  → Y Sampled{Int64} 10:13 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────┘
 ↓ →  10    11    12    13
 1     1.0   1.0   1.0   1.0
 2     1.0   1.0   1.0   1.0
 3     1.0   1.0   1.0   1.0

# Method
julia> function broadcast_dims2(f, As::Union{DimensionalData.AbstractBasicDimArray, DimensionalData.Dimensions.Dimension, Type{<:DimensionalData.Dimension}, Symbol}...)
           # 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
                   if parent(A) isa Colon  # If X(:) is passed, look up values from `existing_dims`
                       dim = dims(existing_dims, A)
                       DimArray(parent(dim), dim)
                   else  # otherwise use the values supplied by the dimension
                       DimArray(parent(A), A)
                   end
               elseif A isa Type{<:DimensionalData.Dimension} || A isa Symbol
                   # If e.g. `X` or `:X` is passed, use values supplied by the dimension
                   dim = dims(existing_dims, A)
                   DimArray(parent(dim), dim)
               else
                   A
               end
           end
           broadcast_dims(f, Bs...)
       end

# example
julia> broadcast_dims2(+, A, Y(:))
╭─────────────────────────╮
│ 3×4 DimArray{Float64,2} │
├─────────────────────────┴──────────────────────── dims ┐
  ↓ X Sampled{Int64} 1:3 ForwardOrdered Regular Points,
  → Y Sampled{Int64} 10:13 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────┘
 ↓ →  10    11    12    13
 1    11.0  12.0  13.0  14.0
 2    11.0  12.0  13.0  14.0
 3    11.0  12.0  13.0  14.0

julia> broadcast_dims2(+, A, :X)
╭─────────────────────────╮
│ 3×4 DimArray{Float64,2} │
├─────────────────────────┴──────────────────────── dims ┐
  ↓ X Sampled{Int64} 1:3 ForwardOrdered Regular Points,
  → Y Sampled{Int64} 10:13 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────┘
 ↓ →  10    11    12    13
 1     2.0   2.0   2.0   2.0
 2     3.0   3.0   3.0   3.0
 3     4.0   4.0   4.0   4.0

This also lets you e.g. construct a completely new DimArray from an operation over dimensions

julia> broadcast_dims2(+, x, y)
╭───────────────────────╮
│ 3×4 DimArray{Int64,2} │
├───────────────────────┴────────────────────────── dims ┐
  ↓ X Sampled{Int64} 1:3 ForwardOrdered Regular Points,
  → Y Sampled{Int64} 10:13 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────┘
 ↓ →  10  11  12  13
 1    11  12  13  14
 2    12  13  14  15
 3    13  14  15  16

julia> broadcast_dims2(tuple, x, y)
╭─────────────────────────────────────╮
│ 3×4 DimArray{Tuple{Int64, Int64},2} │
├─────────────────────────────────────┴──────────── dims ┐
  ↓ X Sampled{Int64} 1:3 ForwardOrdered Regular Points,
  → Y Sampled{Int64} 10:13 ForwardOrdered Regular Points
└────────────────────────────────────────────────────────┘
 ↓ →  10         11         12         13
 1      (1, 10)    (1, 11)    (1, 12)    (1, 13)
 2      (2, 10)    (2, 11)    (2, 12)    (2, 13)
 3      (3, 10)    (3, 11)    (3, 12)    (3, 13)

or effortlessly add a new dimension to a DimArray

julia> z = Z(-5:-4)
Z -5:-4

julia> broadcast_dims2(+, A, z)
╭───────────────────────────╮
│ 3×4×2 DimArray{Float64,3} │
├───────────────────────────┴─────────────────────── dims ┐
  ↓ X Sampled{Int64} 1:3 ForwardOrdered Regular Points,
  → Y Sampled{Int64} 10:13 ForwardOrdered Regular Points,
  ↗ Z Sampled{Int64} -5:-4 ForwardOrdered Regular Points
└─────────────────────────────────────────────────────────┘
[:, :, 1]
 ↓ →  10    11    12    13
 1    -4.0  -4.0  -4.0  -4.0
 2    -4.0  -4.0  -4.0  -4.0
 3    -4.0  -4.0  -4.0  -4.0

[rafaqz]

That's awesome. PR?