Update/unify hdi/ci implementations

Author: sethaxen

ext/InferenceObjectsPosteriorStatsExt/InferenceObjectsPosteriorStatsExt.jl

@@ -6,13 +6,15 @@ using InferenceObjects: InferenceObjects
 using PosteriorStats: PosteriorStats
 using StatsBase: StatsBase
 
-import PosteriorStats: hdi, loo, loo_pit, r2_score, summarize, waic
+import PosteriorStats: eti, hdi, loo, loo_pit, r2_score, summarize, waic
 import StatsBase: summarystats
 
-export hdi, loo, loo_pit, r2_score, summarize, waic, summarystats
+export eti, hdi, loo, loo_pit, r2_score, summarize, waic, summarystats
+
+maplayers = isdefined(DimensionalData, :maplayers) ? DimensionalData.maplayers : map
 
 include("utils.jl")
-include("hdi.jl")
+include("ci.jl")
 include("loo.jl")
 include("waic.jl")
 include("loo_pit.jl")

ext/InferenceObjectsPosteriorStatsExt/ci.jl

@@ -0,0 +1,27 @@
+
+for (ci_fun, ci_desc) in
+    (:eti => "equal-tailed interval (ETI)", :hdi => "highest density interval (HDI)")
+    @eval begin
+        # this pattern ensures that the type is completely specified at compile time
+        @doc """
+            $($ci_fun)(data::InferenceData; kwargs...) -> Dataset
+            $($ci_fun)(data::Dataset; kwargs...) -> Dataset
+
+        Calculate the $($ci_desc) for each parameter in the data.
+
+        For more details and a description of the `kwargs`, see
+        [`PosteriorStats.$($ci_fun)`](@extref).
+        """
+        function PosteriorStats.$(ci_fun)(data::InferenceObjects.InferenceData; kwargs...)
+            return PosteriorStats.$(ci_fun)(data.posterior; kwargs...)
+        end
+        function PosteriorStats.$(ci_fun)(data::InferenceObjects.Dataset; kwargs...)
+            ds = maplayers(data) do var
+                return _as_dimarray(
+                    PosteriorStats.$(ci_fun)(_params_array(var); kwargs...), var
+                )
+            end
+            return DimensionalData.rebuild(ds; metadata=DimensionalData.NoMetadata())
+        end
+    end
+end

ext/InferenceObjectsPosteriorStatsExt/hdi.jl

@@ -159,8 +159,17 @@ function observations_and_predictions(
     end
 end
 
+# reshape to (ndraws, nchains, nparams...) and drop the dimensions
+function _params_array(data)
+    sample_dims = Dimensions.dims(data, InferenceObjects.DEFAULT_SAMPLE_DIMS)
+    param_dims = Dimensions.otherdims(data, sample_dims)
+    dims_combined = Dimensions.combinedims((sample_dims..., param_dims...))
+    Dimensions.dimsmatch(Dimensions.dims(data), dims_combined) && return data
+    return PermutedDimsArray(data, dims_combined)
+end
+
 _as_dimarray(x::DimensionalData.AbstractDimArray, ::DimensionalData.AbstractDimArray) = x
-function _as_dimarray(x::Union{Real,Missing}, arr::DimensionalData.AbstractDimArray)
+function _as_dimarray(x, arr::DimensionalData.AbstractDimArray)
     return Dimensions.rebuild(arr, fill(x), ())
 end
 

ext/InferenceObjectsPosteriorStatsExt/utils.jl

@@ -6,8 +6,11 @@ using Statistics
 using StatsBase
 using Test
 
+_as_array(x) = fill(x)
+_as_array(x::AbstractArray) = x
+
 @testset "PosteriorStats integration" begin
-    @testset "hdi" begin
+    @testset for ci_fun in (eti, hdi)
         nt = (x=randn(1000, 3), y=randn(1000, 3, 4), z=randn(1000, 3, 4, 2))
         posterior = convert_to_dataset(nt)
         posterior_perm = convert_to_dataset((
@@ -17,23 +20,17 @@ using Test
         ))
         idata = InferenceData(; posterior)
         @testset for prob in (0.76, 0.93)
-            if VERSION ≥ v"1.9"
-                @test_broken @inferred hdi(posterior; prob)
-            end
-            r1 = hdi(posterior; prob)
-            r1_perm = hdi(posterior_perm; prob)
+            @test_broken @inferred ci_fun(posterior; prob)
+            r1 = ci_fun(posterior; prob)
+            r1_perm = ci_fun(posterior_perm; prob)
             for k in (:x, :y, :z)
-                rk = hdi(posterior[k]; prob)
-                @test r1[k][hdi_bound=At(:lower)] == rk.lower
-                @test r1[k][hdi_bound=At(:upper)] == rk.upper
+                rk = ci_fun(posterior[k]; prob)
+                @test r1[k] == _as_array(rk)
                 # equality check is safe because these are always data values
-                @test r1_perm[k][hdi_bound=At(:lower)] == rk.lower
-                @test r1_perm[k][hdi_bound=At(:upper)] == rk.upper
-            end
-            if VERSION ≥ v"1.9"
-                @test_broken @inferred hdi(idata; prob)
+                @test r1_perm[k] == _as_array(rk)
             end
-            r2 = hdi(idata; prob)
+            @test_broken @inferred ci_fun(idata; prob)
+            r2 = ci_fun(idata; prob)
             @test r1 == r2
         end
     end