Fixes to return types of ess_rhat

[sethaxen]

This PR makes 2 improvements to ess_rhat:

• Don't unnecessarily promote the return type to Float64.
• Remove unused type parameter U
• Return arrays ess and rhat preserve structure of last dimension of input array. This is especially useful when the input array is one with named dimensions, like DimensionalData.DimArray, since the indices in that dimension will be preserved.

[coveralls]

Pull Request Test Coverage Report for Build 3930744323

• 7 of 7 (100.0%) changed or added relevant lines in 1 file are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage increased (+0.4%) to 95.634%

---

Totals
Change from base Build 3906657106:	0.4%
Covered Lines:	679
Relevant Lines:	710

---

💛 - Coveralls

[codecov]

Codecov Report

Base: 95.21% // Head: 95.63% // Increases project coverage by +0.42% 🎉

Coverage data is based on head (7ae645f) compared to base (b6e8de9).
Patch coverage: 100.00% of modified lines in pull request are covered.

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #62      +/-   ##
==========================================
+ Coverage   95.21%   95.63%   +0.42%     
==========================================
  Files          10       10              
  Lines         710      710              
==========================================
+ Hits          676      679       +3     
+ Misses         34       31       -3     

Impacted Files	Coverage Δ
src/ess.jl	100.00% <100.00%> (+1.98%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

[sethaxen]

With DimensionalData v0.24.1 and this PR, this now works:

julia> using DimensionalData, MCMCDiagnosticTools

julia> x = DimArray(randn(Float32, 1000, 4, 3), (:draw, :chain, Dim{:var}([:a, :b, :c])));

julia> S, R = ess_rhat_bulk(x);

julia> S
3-element DimArray{Float32,1} with dimensions: 
  Dim{:var} Categorical{Symbol} Symbol[a, b, c] ForwardOrdered
 :a  3899.19
 :b  3973.59
 :c  3797.71

julia> R
3-element DimArray{Float32,1} with dimensions: 
  Dim{:var} Categorical{Symbol} Symbol[a, b, c] ForwardOrdered
 :a  0.999755
 :b  1.00018
 :c  1.0003

No such luck for any other array types I've tried with named axes/indices, but seems they can opt in if they just implement the documented similar interface.

This PR should be ready for final review.

[devmotion]

Looks good, just one minor question. Possibly something for a different PR though.

[devmotion]

Do we have any tests for this? And any @test_inferred checks - it seems they should fail due to this line?

[sethaxen]

Yes, these lines are tested:

MCMCDiagnosticTools.jl/test/ess.jl

Lines 122 to 134 in b6e8de9

	@testset "ESS and R̂ (single sample)" begin # check that issue #137 is fixed
	x = rand(1, 3, 5)
	for method in (ESSMethod(), FFTESSMethod(), BDAESSMethod())
	# analyze array
	ess_array, rhat_array = ess_rhat(x; method=method)
	@test length(ess_array) == size(x, 3)
	@test all(ismissing, ess_array) # since min(maxlag, niter - 1) = 0
	@test length(rhat_array) == size(x, 3)
	@test all(ismissing, rhat_array)
	end
	end

.

I added more tests that the array type returned here preserves the final axis and some @inferred tests specifying the type union.

[devmotion]

IMO this type union is a bit annoying. I wonder if in a breaking release we should change this to an error, similar to what StatsBase does e.g. in this case:

julia> using StatsBase

julia> autocov(rand(3), [4])
ERROR: lags must be less than the sample length.
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] check_lags
   @ ~/.julia/packages/StatsBase/XgjIN/src/signalcorr.jl:15 [inlined]
 [3] autocov!(r::Vector{Float64}, x::Vector{Float64}, lags::Vector{Int64}; demean::Bool)
   @ StatsBase ~/.julia/packages/StatsBase/XgjIN/src/signalcorr.jl:68
 [4] #autocov#113
   @ ~/.julia/packages/StatsBase/XgjIN/src/signalcorr.jl:115 [inlined]
 [5] autocov(x::Vector{Float64}, lags::Vector{Int64})
   @ StatsBase ~/.julia/packages/StatsBase/XgjIN/src/signalcorr.jl:113
 [6] top-level scope
   @ REPL[3]:1

[sethaxen]

I agree, this is a good idea.