add dimtrait and use it for plots

ext/DimensionalDataMakie.jl

@@ -470,16 +470,22 @@ end
 
 # Replace the existing dimensions with X/Y/Z so we have a 1:1 
 # relationship with the possible Makie.jl plot axes. 
+# Replacements all rwp the original dimensions so
+# they can be recovered later on.
 function _get_replacement_dims(A::AbstractDimArray{<:Any,N}, replacements::Tuple) where N
     xyz_dims = (X(), Y(), Z())[1:N]
     map(replacements) do d
         # Make sure replacements contain X/Y/Z only
         hasdim(A, d) || throw(ArgumentError("object does not have a dimension $(basetypeof(d))"))
     end
-    # Find and sort remaining dims
-    source_dims_remaining = dims(otherdims(A, replacements), DD.PLOT_DIMENSION_ORDER)
-    xyz_remaining = otherdims(xyz_dims, map(val, replacements))[1:length(source_dims_remaining)]
-    other_replacements = map(rebuild, source_dims_remaining, xyz_remaining)
+    # Wrap remaining dims with their dimtrait
+    not_replaced = otherdims(A, replacements)
+    wrapped = map(not_replaced) do d
+        rebuild(d, dimtrait(d))
+    end
+    sorted_wrapped = dims(wrapped, DD.PLOT_DIMENSION_ORDER)
+    xyz_remaining = otherdims(xyz_dims, map(val, replacements))[1:length(sorted_wrapped)]
+    other_replacements = map(rebuild, sorted_wrapped, xyz_remaining)
     return (replacements..., other_replacements...)
 end
 

src/Dimensions/Dimensions.jl

@@ -33,7 +33,7 @@ using .Lookups: StandardIndices, SelTuple, CategoricalEltypes,
 using Base: tail, OneTo, @propagate_inbounds
 
 export name, label, dimnum, hasdim, hasselection, otherdims, commondims, combinedims,
-    setdims, swapdims, sortdims, lookup, set, format, rebuild, name2dim,
+    setdims, swapdims, sortdims, lookup, set, format, rebuild, name2dim, dimtrait,
     basetypeof, basedims, dims2indices, slicedims, dimsmatch, comparedims, reducedims
 
 export Dimension, IndependentDim, DependentDim, XDim, YDim, ZDim, TimeDim,

src/Dimensions/dimension.jl

@@ -523,3 +523,30 @@ mean(A; dims=Ti)
 @dim Ti TimeDim "Time"
 
 const Time = Ti # For some backwards compat
+
+# dimtrait uses the standard dimensions as traits so that 
+# other dimensions can be matched to axes, mostly in plotting
+dimtrait(::D) where D<:Dimension = dimtrait(D)
+function dimtrait(::Type{D}) where D<:Dimension
+    st = supertype(D)
+    @show st
+    return st == Any ? AnonDim() : dimtrait(st)
+end
+dimtrait(::Type{<:XDim}) = X()
+dimtrait(::Type{<:YDim}) = Y()
+dimtrait(::Type{<:ZDim}) = Z()
+dimtrait(::Type{<:TimeDim}) = Ti()
+Base.@assume_effects :total function dimtrait(::Type{<:Dim{K}}) where K
+    Kl = Symbol(lowercase(string(K)))
+    if Kl in (:t, :time, :ti)
+        Ti()
+    elseif Kl in (:x, :lon, :longitude)
+        X()
+    elseif Kl in (:y, :lat, :latitude)
+        Y()
+    elseif Kl in (:z, :elev, :lev, :elevation)
+        Z()
+    else
+        AnonDim()
+    end
+end

src/plotrecipes.jl

@@ -188,5 +188,14 @@ function refdims_title(lookup::Lookup, refdim::Dimension; kw...)
 end
 
 const PLOT_DIMENSION_ORDER = (TimeDim, XDim, IndependentDim, IndependentDim, YDim, ZDim, DependentDim, DependentDim, Dimension, Dimension, Dimension)
-forward_order_plot_dims(x) = dims(dims(x), PLOT_DIMENSION_ORDER)
+function forward_order_plot_dims(x)
+    # Wrap dimensions as X/Y/Z/Ti/AnonDim
+    ds = map(dims(x)) do d
+        rebuild(dimtrait(d), d)
+    end
+    # Sort
+    sorted = dims(ds, PLOT_DIMENSION_ORDER)
+    # Unwrap
+    map(val, sorted)
+end
 reverse_order_plot_dims(x) = reverse(forward_order_plot_dims(reverse(dims(x))))

test/dimension.jl

@@ -144,3 +144,10 @@ end
     @test typeof(dims(a)[1]) <: X
     @test a.data == cos.(d.val)
 end
+
+@testset "applying function on a dimension" begin
+    @test dimtrait(X) == dimtrait(X()) == dimtrait(Dim{:X}()) == dimtrait(Dim{:x}()) == X()
+    @test dimtrait(Y) == dimtrait(Y()) == dimtrait(Dim{:Y}()) == dimtrait(Dim{:y}()) == Y()
+    @test dimtrait(Z) == dimtrait(Z()) == dimtrait(Dim{:Z}()) == dimtrait(Dim{:z}()) == Z()
+    @test dimtrait(Ti) == dimtrait(Ti()) == dimtrait(Dim{:Time}()) == dimtrait(Dim{:time}()) == Ti()
+end