Dimensional broadcasts with @d and broadcast_dims

Broadcasting over AbstractDimArray works as usual with Base Julia broadcasts, except that dimensions are checked for compatibility with each other, and that values match. Strict checks can be turned off globally with strict_broadcast!(false). To avoid even dimension name checks, broadcast over parent(dimarray).

The @d macro is a dimension-aware extension to regular dot broadcasting. broadcast_dims is analogous to Base Julia's broadcast.

Because we know the names of the dimensions, there is no ambiguity in which ones we mean to broadcast together. This means we can permute and reshape dims so that broadcasts that would fail with a regular Array just work with a DimArray.

As an added bonus, broadcast_dims even works on DimStacks. Currently, @d does not work on DimStack.

Example: scaling along the time dimension

Define some dimensions:

using DimensionalData
using Dates
using Statistics

julia> x, y, t = X(1:100), Y(1:25), Ti(DateTime(2000):Month(1):DateTime(2000, 12))

(↓ X  1:100,
→ Y  1:25,
↗ Ti Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00"))

A DimArray from 1:12 to scale with:

julia> month_scalars = DimArray(month, t)

╭────────────────────────────────────────╮
│ 12-element DimArray{Int64,1} month(Ti) │
├────────────────────────────────────────┴─────────────────────────────── dims ┐
  ↓ Ti Sampled{Dates.DateTime} Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00") ForwardOrdered Regular Points
└──────────────────────────────────────────────────────────────────────────────┘
 2000-01-01T00:00:00   1
 2000-02-01T00:00:00   2
 2000-03-01T00:00:00   3
 2000-04-01T00:00:00   4
 2000-05-01T00:00:00   5
 2000-06-01T00:00:00   6
 2000-07-01T00:00:00   7
 2000-08-01T00:00:00   8
 2000-09-01T00:00:00   9
 2000-10-01T00:00:00  10
 2000-11-01T00:00:00  11
 2000-12-01T00:00:00  12

And a larger DimArray for example data:

julia> data = rand(x, y, t)

╭───────────────────────────────╮
│ 100×25×12 DimArray{Float64,3} │
├───────────────────────────────┴──────────────────────────────────────── dims ┐
  ↓ X  Sampled{Int64} 1:100 ForwardOrdered Regular Points,
  → Y  Sampled{Int64} 1:25 ForwardOrdered Regular Points,
  ↗ Ti Sampled{Dates.DateTime} Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00") ForwardOrdered Regular Points
└──────────────────────────────────────────────────────────────────────────────┘
[:, :, 1]
   ↓ →  1          2         3          …  23          24         25
   1    0.624685   0.909533  0.43314        0.0541549   0.721101   0.600306
   2    0.546357   0.621601  0.960677       0.129382    0.434374   0.438592
   3    0.407207   0.861652  0.0837074      0.0725917   0.800084   0.749829
   4    0.67539    0.637399  0.416444       0.853879    0.988957   0.44724
   ⋮                                    ⋱                          ⋮
  97    0.474087   0.475786  0.759047       0.704613    0.706165   0.1279
  98    0.844262   0.668594  0.431398       0.746465    0.919534   0.398776
  99    0.0471885  0.788358  0.0145691      0.647328    0.517971   0.910412
 100    0.0442448  0.702787  0.585658   …   0.68326     0.774131   0.403955

A regular broadcast fails:

julia> scaled = data .* month_scalars

ERROR: DimensionMismatch: arrays could not be broadcast to a common size: a has axes DimensionalData.Dimensions.DimUnitRange(Base.OneTo(100), X{Sampled{Int64, UnitRange{Int64}, ForwardOrdered, Regular{Int64}, Points, NoMetadata}}([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100])) and b has axes DimensionalData.Dimensions.DimUnitRange(Base.OneTo(12), Ti{Sampled{Dates.DateTime, StepRange{Dates.DateTime, Dates.Month}, ForwardOrdered, Regular{Dates.Month}, Points, NoMetadata}}([Dates.DateTime("2000-01-01T00:00:00"), Dates.DateTime("2000-02-01T00:00:00"), Dates.DateTime("2000-03-01T00:00:00"), Dates.DateTime("2000-04-01T00:00:00"), Dates.DateTime("2000-05-01T00:00:00"), Dates.DateTime("2000-06-01T00:00:00"), Dates.DateTime("2000-07-01T00:00:00"), Dates.DateTime("2000-08-01T00:00:00"), Dates.DateTime("2000-09-01T00:00:00"), Dates.DateTime("2000-10-01T00:00:00"), Dates.DateTime("2000-11-01T00:00:00"), Dates.DateTime("2000-12-01T00:00:00")]))

But @d knows to broadcast over the Ti dimension:

julia> scaled = @d data .* month_scalars

╭───────────────────────────────╮
│ 100×25×12 DimArray{Float64,3} │
├───────────────────────────────┴──────────────────────────────────────── dims ┐
  ↓ X  Sampled{Int64} 1:100 ForwardOrdered Regular Points,
  → Y  Sampled{Int64} 1:25 ForwardOrdered Regular Points,
  ↗ Ti Sampled{Dates.DateTime} Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00") ForwardOrdered Regular Points
└──────────────────────────────────────────────────────────────────────────────┘
[:, :, 1]
   ↓ →  1          2         3          …  23          24         25
   1    0.624685   0.909533  0.43314        0.0541549   0.721101   0.600306
   2    0.546357   0.621601  0.960677       0.129382    0.434374   0.438592
   3    0.407207   0.861652  0.0837074      0.0725917   0.800084   0.749829
   4    0.67539    0.637399  0.416444       0.853879    0.988957   0.44724
   ⋮                                    ⋱                          ⋮
  97    0.474087   0.475786  0.759047       0.704613    0.706165   0.1279
  98    0.844262   0.668594  0.431398       0.746465    0.919534   0.398776
  99    0.0471885  0.788358  0.0145691      0.647328    0.517971   0.910412
 100    0.0442448  0.702787  0.585658   …   0.68326     0.774131   0.403955

We can see the means of each month are scaled by the broadcast :

julia> mean(eachslice(data; dims=(X, Y)))

╭────────────────────────────────╮
│ 12-element DimArray{Float64,1} │
├────────────────────────────────┴─────────────────────────────────────── dims ┐
  ↓ Ti Sampled{Dates.DateTime} Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00") ForwardOrdered Regular Points
└──────────────────────────────────────────────────────────────────────────────┘
 2000-01-01T00:00:00  0.50854
 2000-02-01T00:00:00  0.499674
 2000-03-01T00:00:00  0.510863
 2000-04-01T00:00:00  0.500818
 2000-05-01T00:00:00  0.501124
 2000-06-01T00:00:00  0.495215
 2000-07-01T00:00:00  0.489497
 2000-08-01T00:00:00  0.500495
 2000-09-01T00:00:00  0.491479
 2000-10-01T00:00:00  0.499282
 2000-11-01T00:00:00  0.491202
 2000-12-01T00:00:00  0.485718

julia> mean(eachslice(scaled; dims=(X, Y)))

╭────────────────────────────────╮
│ 12-element DimArray{Float64,1} │
├────────────────────────────────┴─────────────────────────────────────── dims ┐
  ↓ Ti Sampled{Dates.DateTime} Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00") ForwardOrdered Regular Points
└──────────────────────────────────────────────────────────────────────────────┘
 2000-01-01T00:00:00  0.50854
 2000-02-01T00:00:00  0.999348
 2000-03-01T00:00:00  1.53259
 2000-04-01T00:00:00  2.00327
 2000-05-01T00:00:00  2.50562
 2000-06-01T00:00:00  2.97129
 2000-07-01T00:00:00  3.42648
 2000-08-01T00:00:00  4.00396
 2000-09-01T00:00:00  4.42331
 2000-10-01T00:00:00  4.99282
 2000-11-01T00:00:00  5.40322
 2000-12-01T00:00:00  5.82862

You can also use broadcast_dims the same way:

julia> broadcast_dims(*, data, month_scalars)

╭───────────────────────────────╮
│ 100×25×12 DimArray{Float64,3} │
├───────────────────────────────┴──────────────────────────────────────── dims ┐
  ↓ X  Sampled{Int64} 1:100 ForwardOrdered Regular Points,
  → Y  Sampled{Int64} 1:25 ForwardOrdered Regular Points,
  ↗ Ti Sampled{Dates.DateTime} Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00") ForwardOrdered Regular Points
└──────────────────────────────────────────────────────────────────────────────┘
[:, :, 1]
   ↓ →  1          2         3          …  23          24         25
   1    0.624685   0.909533  0.43314        0.0541549   0.721101   0.600306
   2    0.546357   0.621601  0.960677       0.129382    0.434374   0.438592
   3    0.407207   0.861652  0.0837074      0.0725917   0.800084   0.749829
   4    0.67539    0.637399  0.416444       0.853879    0.988957   0.44724
   ⋮                                    ⋱                          ⋮
  97    0.474087   0.475786  0.759047       0.704613    0.706165   0.1279
  98    0.844262   0.668594  0.431398       0.746465    0.919534   0.398776
  99    0.0471885  0.788358  0.0145691      0.647328    0.517971   0.910412
 100    0.0442448  0.702787  0.585658   …   0.68326     0.774131   0.403955

And with the @d macro you can set the dimension order and other properties of the output array, by passing a single assignment or a NamedTuple argument to @d after the broadcast:

julia> @d data .* month_scalars dims=(Ti, X, Y)

╭───────────────────────────────╮
│ 12×100×25 DimArray{Float64,3} │
├───────────────────────────────┴──────────────────────────────────────── dims ┐
  ↓ Ti Sampled{Dates.DateTime} Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00") ForwardOrdered Regular Points,
  → X  Sampled{Int64} 1:100 ForwardOrdered Regular Points,
  ↗ Y  Sampled{Int64} 1:25 ForwardOrdered Regular Points
└──────────────────────────────────────────────────────────────────────────────┘
[:, :, 1]
 ↓ →                   1         …  98         99          100
  2000-01-01T00:00:00  0.624685      0.844262   0.0471885    0.0442448
  2000-02-01T00:00:00  0.547848      1.79785    0.187495     0.441762
  2000-03-01T00:00:00  2.49232       1.15489    2.82356      0.921936
  2000-04-01T00:00:00  0.439142      2.61341    1.59542      2.24247
 ⋮                               ⋱                           ⋮
  2000-09-01T00:00:00  3.27042       6.61858    3.10246      5.93972
  2000-10-01T00:00:00  0.12404   …   2.53632    7.83904      6.55713
  2000-11-01T00:00:00  1.79388       2.5891     2.14203      5.52389
  2000-12-01T00:00:00  8.84219       6.15004    5.04169      6.10415

Or

julia> @d data .* month_scalars (dims=(Ti, X, Y), name=:scaled)

╭──────────────────────────────────────╮
│ 12×100×25 DimArray{Float64,3} scaled │
├──────────────────────────────────────┴───────────────────────────────── dims ┐
  ↓ Ti Sampled{Dates.DateTime} Dates.DateTime("2000-01-01T00:00:00"):Dates.Month(1):Dates.DateTime("2000-12-01T00:00:00") ForwardOrdered Regular Points,
  → X  Sampled{Int64} 1:100 ForwardOrdered Regular Points,
  ↗ Y  Sampled{Int64} 1:25 ForwardOrdered Regular Points
└──────────────────────────────────────────────────────────────────────────────┘
[:, :, 1]
 ↓ →                   1         …  98         99          100
  2000-01-01T00:00:00  0.624685      0.844262   0.0471885    0.0442448
  2000-02-01T00:00:00  0.547848      1.79785    0.187495     0.441762
  2000-03-01T00:00:00  2.49232       1.15489    2.82356      0.921936
  2000-04-01T00:00:00  0.439142      2.61341    1.59542      2.24247
 ⋮                               ⋱                           ⋮
  2000-09-01T00:00:00  3.27042       6.61858    3.10246      5.93972
  2000-10-01T00:00:00  0.12404   …   2.53632    7.83904      6.55713
  2000-11-01T00:00:00  1.79388       2.5891     2.14203      5.52389
  2000-12-01T00:00:00  8.84219       6.15004    5.04169      6.10415