Use DimensionalData.maplayers if available

Author: sethaxen

ext/InferenceObjectsMCMCDiagnosticToolsExt/InferenceObjectsMCMCDiagnosticToolsExt.jl

@@ -13,6 +13,8 @@ else  # using Requires
     using ..Random: Random
 end
 
+maplayers = isdefined(DimensionalData, :maplayers) ? DimensionalData.maplayers : map
+
 include("utils.jl")
 include("bfmi.jl")
 include("ess_rhat.jl")

ext/InferenceObjectsMCMCDiagnosticToolsExt/ess_rhat.jl

@@ -21,7 +21,7 @@ end
 for f in (:ess, :rhat)
     @eval begin
         function MCMCDiagnosticTools.$f(data::InferenceObjects.Dataset; kwargs...)
-            ds = map(data) do var
+            ds = maplayers(data) do var
                 return _as_dimarray(MCMCDiagnosticTools.$f(_params_array(var); kwargs...), var)
             end
             return DimensionalData.rebuild(ds; metadata=DimensionalData.NoMetadata())

ext/InferenceObjectsMCMCDiagnosticToolsExt/mcse.jl

@@ -8,7 +8,7 @@ function MCMCDiagnosticTools.mcse(data::InferenceObjects.InferenceData; kwargs..
     return MCMCDiagnosticTools.mcse(data.posterior; kwargs...)
 end
 function MCMCDiagnosticTools.mcse(data::InferenceObjects.Dataset; kwargs...)
-    ds = map(data) do var
+    ds = maplayers(data) do var
         return _as_dimarray(MCMCDiagnosticTools.mcse(_params_array(var); kwargs...), var)
     end
     return DimensionalData.rebuild(ds; metadata=DimensionalData.NoMetadata())

ext/InferenceObjectsMCMCDiagnosticToolsExt/rstar.jl

@@ -16,7 +16,7 @@ end
 function MCMCDiagnosticTools.rstar(
     rng::Random.AbstractRNG, clf, data::InferenceObjects.Dataset; kwargs...
 )
-    data_array = cat(map(_as_3d_array ∘ _params_array, data)...; dims=3)
+    data_array = cat(maplayers(_as_3d_array ∘ _params_array, data)...; dims=3)
     return MCMCDiagnosticTools.rstar(rng, clf, data_array; kwargs...)
 end
 function MCMCDiagnosticTools.rstar(

src/dataset.jl

@@ -126,6 +126,8 @@ for f in [:data, :dims, :refdims, :metadata, :layerdims, :layermetadata]
     end
 end
 
+DimensionalData.modify(f, s::Dataset) = Dataset(DimensionalData.modify(f, parent(s)))
+
 # Warning: this is not an API function and probably should be implemented abstractly upstream
 DimensionalData.show_after(io, mime, ::Dataset) = nothing
 

test/mcmcdiagnostictools.jl

@@ -7,6 +7,10 @@ using Random
 using Statistics
 using Test
 
+if !isdefined(DimensionalData, :maplayers)
+    maplayers = map
+end
+
 @testset "MCMCDiagnosticTools integration" begin
     nchains, ndraws = 4, 10
     sizes = (x=(), y=(2,), z=(3, 5))
@@ -16,12 +20,12 @@ using Test
     dict1 = Dict(Symbol(k) => randn(ndraws, nchains, sz...) for (k, sz) in pairs(sizes))
     idata1 = from_dict(dict1; dims, coords, sample_stats=Dict(:energy => energy))
     # permute dimensions to test that diagnostics are invariant to dimension order
-    post2 = map(idata1.posterior) do var
+    post2 = maplayers(idata1.posterior) do var
         n = ndims(var)
         permdims = ((3:n)..., 2, 1)
         return permutedims(var, permdims)
     end
-    sample_stats2 = map(permutedims, idata1.sample_stats)
+    sample_stats2 = maplayers(permutedims, idata1.sample_stats)
     idata2 = InferenceData(; posterior=post2, sample_stats=sample_stats2)
 
     @testset for f in (ess, rhat, ess_rhat, mcse)
@@ -35,7 +39,7 @@ using Test
             @test issetequal(keys(metric), keys(idata1.posterior))
             @test metric == f(idata1.posterior; kind)
             @test metric2 == f(idata2.posterior; kind)
-            @test all(map(≈, metric2, metric))
+            @test all(maplayers(≈, metric2, metric))
             for k in keys(sizes)
                 @test all(
                     hasdim(
@@ -81,7 +85,7 @@ using Test
             r4 = rstar(rng, classifier(rng), idata2.posterior; subset)
             rng = Random.seed!(123)
             post_mat = cat(
-                map(var -> reshape(parent(var), ndraws, nchains, :), idata1.posterior)...;
+                maplayers(var -> reshape(parent(var), ndraws, nchains, :), idata1.posterior)...;
                 dims=3,
             )
             r5 = rstar(rng, classifier(rng), post_mat; subset)