DimensionalData + Oceananigans

[tomchor]

I wanted to share here a package that I found called DimensionalData. Maybe people already have been using it and I'm just late to the game, but I found it makes my workflow much easier. I guess the easiest way to explain the package is by saying that it reproduces many of the features of python-xarray natively in Julia.

Basically it all starts by naming and labeling dimensions/axis of arrays to make it easier to select and slice your data. (So it's similar to AxisArrays, but it appears to be more mature.)

Here's a quick example of how it could work with Oceananigans:

julia> using Oceananigans

julia> using DimensionalData

julia> import DimensionalData: DimArray

julia> function DimArray(f::Field; kwargs...)
           X, Y, Z = nodes(f)
           return DimArray(interior(f), (; X, Y, Z), kwargs...)
       end
DimArray

julia> function DimArray(model::Oceananigans.AbstractModel; kwargs...)
           return map(DimArray, fields(model))
       end
DimArray

julia> grid = RectilinearGrid(size=(8,8,8,), extent=(1,1,1))
8×8×8 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
├── Periodic x ∈ [0.0, 1.0)  regularly spaced with Δx=0.125
├── Periodic y ∈ [0.0, 1.0)  regularly spaced with Δy=0.125
└── Bounded  z ∈ [-1.0, 0.0] regularly spaced with Δz=0.125

julia> model = NonhydrostaticModel(grid = grid);

julia> set!(model, u = (x,y,z) -> randn());

julia> dimfields = DimArray(model);

julia> dimfields.u
8×8×8 DimArray{Float64,3} with dimensions: 
  X Sampled Float64[0.0, 0.125, …, 0.75, 0.875] ForwardOrdered Irregular Points,
  Y Sampled Float64[0.0625, 0.1875, …, 0.8125, 0.9375] ForwardOrdered Irregular Points,
  Z Sampled Float64[-0.9375, -0.8125, …, -0.1875, -0.0625] ForwardOrdered Irregular Points
[:, :, 1]
         0.0625      0.1875      0.3125     0.4375      0.5625     0.6875     0.8125      0.9375
 0.0     0.0908888  -1.79494    -0.306636   1.12602     1.33471   -0.254043  -0.39004     0.102122
 0.125  -0.411471    0.0289425  -1.02885    0.539582    0.843286   0.273952   0.529021   -0.0953487
 0.25   -0.652228    0.427576    0.476498   0.136816    0.463953   0.167134   0.646413    1.3374
 0.375  -0.366688    0.99938     0.571352   0.0117112   0.130476  -0.147442   0.281299    0.407026
 0.5    -0.479653    0.47951     0.15681   -0.459028    0.582198  -0.777213   0.665423    0.451173
 0.625  -0.290271   -0.74456    -0.475706   0.454353   -0.385925   1.00917   -0.064887    0.878018
 0.75    1.00748    -1.53969    -0.754316   0.130978    0.608167   0.3162    -0.385334    0.561976
 0.875   1.11118    -2.39276    -0.245676   0.575122    1.25287   -0.413678  -0.0695997  -0.626723
[and 7 more slices...]

julia> dimfields.u[Y=4, Z=Near(-.25)]
8-element DimArray{Float64,1} with dimensions: 
  X Sampled Float64[0.0, 0.125, …, 0.75, 0.875] ForwardOrdered Irregular Points
and reference dimensions: 
  Y Sampled Float64[0.4375] ForwardOrdered Irregular Points,
  Z Sampled Float64[-0.1875] ForwardOrdered Irregular Points
 0.0     0.422572
 0.125   0.533293
 0.25    0.996881
 0.375   0.138838
 0.5     0.167402
 0.625  -0.353312
 0.75    0.300673
 0.875   0.783375

In the last line I selected data value of u in the 4-th y position (index j=4) and in the z index that's the Nearest z=-.25 (which turns out to be z=-.1875), which I think it's pretty handy.

The feature that's helped me the most is the plotting, which it makes super easy. For now it only has recipes for Plots.jl, but they work really well out of the box. For example:

julia> using Plots
WARNING: using Plots.grid in module Main conflicts with an existing identifier.

julia> heatmap(dimfields.u[Z=Near(-.25)], c=:vik, clim=(-.7, .7))

produces this plot:

There are obviously more things (for example the package Rasters.jl understands Oceananigans NetCDF output directly) but I'm curious to hear what people think about it.

[glwagner]

It could be interesting to make the selector syntax u[Y=4, Z=Near(-.25)] work with fields directly. For example, we could have AbstractField subtype AbstractDimArray. We looked into this a while ago, but it was unclear whether it would be a popular enough feature to warrant the maintainance.

This is discussed here: #457