Add psis! method

Author: Closed-Limelike-Curves

src/ESS.jl

@@ -4,27 +4,29 @@ using Tullio
 export relative_eff, psis_ess, sup_ess
 
 """
-    relative_eff(sample::AbstractArray{<:Real, 3}; [method])
+    relative_eff(
+        sample::AbstractArray{Real, 3}; 
+        method=MCMCDiagnosticTools.FFTESSMethod()
+    )
 
 Calculate the relative efficiency of an MCMC chain, i.e. the effective sample size divided
-by the nominal sample size. If none is provided, use the default method from 
-MCMCDiagnosticTools.
+by the nominal sample size.
 """
 function relative_eff(sample::AbstractArray{<:Real, 3}; maxlag=size(sample, 2), kwargs...)
     dims = size(sample)
     post_sample_size = dims[2] * dims[3]
     ess_sample = inv.(permutedims(sample, [2, 1, 3]))
-    ess, = MCMCDiagnosticTools.ess_rhat(ess_sample; maxlag=maxlag, kwargs...)
+    ess, = MCMCDiagnosticTools.ess_rhat(ess_sample; method=method, maxlag=dims[2])
     r_eff = ess / post_sample_size
     return r_eff
 end
 
 
 """
     function psis_ess(
-        weights::AbstractVector{<:Real},
-        r_eff::AbstractVector{<:Real}
-    ) -> AbstractVector{<:Real}
+        weights::AbstractVector{T<:Real},
+        r_eff::AbstractVector{T}
+    ) -> AbstractVector{T}
 
 Calculate the (approximate) effective sample size of a PSIS sample, using the correction in
 Vehtari et al. 2019. This uses the variance-based definition of ESS, and measures the L2 
@@ -37,6 +39,14 @@ distance of the proposal and target distributions.
 
 See `?relative_eff` to calculate `r_eff`.
 """
+function psis_ess(
+    weights::AbstractVector{T}, r_eff::AbstractVector{T}
+) where {T <: Union{Real, Missing}}
+    @tullio sum_of_squares := weights[x]^2
+    return r_eff ./ sum_of_squares
+end
+
+
 function psis_ess(
     weights::AbstractMatrix{T}, r_eff::AbstractVector{T}
 ) where {T <: Union{Real, Missing}}
@@ -45,7 +55,7 @@ function psis_ess(
 end
 
 
-function psis_ess(weights::AbstractMatrix{<:Real})
+function psis_ess(weights::AbstractMatrix{<:Union{Real, Missing}})
     @warn "PSIS ESS not adjusted based on MCMC ESS. MCSE and ESS estimates " *
           "will be overoptimistic if samples are autocorrelated."
     return psis_ess(weights, ones(size(weights)))
@@ -54,8 +64,8 @@ end
 
 """
     function sup_ess(
-        weights::AbstractVector{<:Real},
-        r_eff::AbstractVector{<:Real}
+        weights::AbstractVector{T},
+        r_eff::AbstractVector{T}
     ) -> AbstractVector
 
 Calculate the supremum-based effective sample size of a PSIS sample, i.e. the inverse of the
@@ -64,8 +74,7 @@ L-∞ norm.
 
 # Arguments
   - `weights`: A set of importance sampling weights derived from PSIS.
-  - `r_eff`: The relative efficiency of the MCMC chains from which PSIS samples were 
-    derived.
+  - `r_eff`: The relative efficiency of the MCMC chains; see also [`relative_eff`]@ref.
 """
 function sup_ess(
     weights::AbstractMatrix{T}, r_eff::V

src/ImportanceSampling.jl

@@ -9,7 +9,7 @@ double check it is correct.
 const MIN_TAIL_LEN = 5  # Minimum size of a tail for PSIS to give sensible answers
 const SAMPLE_SOURCES = ["mcmc", "vi", "other"]
 
-export psis, PsisLoo, PsisLooMethod, Psis
+export psis, psis!, PsisLoo, PsisLooMethod, Psis
 
 
 ###########################
@@ -125,7 +125,7 @@ function psis(
     ξ = similar(r_eff)
     @inbounds Threads.@threads for i in eachindex(tail_length)
         tail_length[i] = _def_tail_length(post_sample_size, r_eff[i])
-        ξ[i] = @views ParetoSmooth._do_psis_i!(weights[i, :], tail_length[i])
+        ξ[i] = @views psis!(weights[i, :], tail_length[i])
     end
 
     @tullio norm_const[i] := weights[i, j]
@@ -159,21 +159,36 @@ function psis(
 end
 
 
+function psis(is_ratios::AbstractVector{<:Real}, args...)
+    new_ratios = copy(is_ratios)
+    ξ = psis!(new_ratios)
+    return new_ratios, ξ
+end
+
+
+
 """
-    _do_psis_i!(is_ratios::AbstractVector{Real}, tail_length::Integer) -> T
+    psis!(is_ratios::AbstractVector{<:Real}, tail_length::Integer) -> Real
+    psis!(is_ratios::AbstractVector{<:Real}, r_eff::Real) -> Real
 
-Do PSIS on a single vector, smoothing its tail values.
+Do PSIS on a single vector, smoothing its tail values *in place* before returning the 
+estimated tail value.
 
 # Arguments
 
   - `is_ratios::AbstractVector{<:Real}`: A vector of importance sampling ratios,
     scaled to have a maximum of 1.
+  - `r_eff::AbstractVector{<:Real}`: A vector of relative effective sample sizes if .
 
 # Returns
 
   - `T<:Real`: ξ, the shape parameter for the GPD; big numbers indicate thick tails.
+
+# Notes
+
+Unlike `psis`, `psis!` performs no checks to make sure the input values are valid.
 """
-function _do_psis_i!(is_ratios::AbstractVector{T}, tail_length::Integer) where {T <: Real}
+function psis!(is_ratios::AbstractVector{<:Real}, tail_length::Integer)
 
     len = length(is_ratios)
     tail_start = len - tail_length + 1  # index of smallest tail value
@@ -198,13 +213,19 @@ function _do_psis_i!(is_ratios::AbstractVector{T}, tail_length::Integer) where {
 end
 
 
+function psis!(is_ratios::AbstractVector{<:Real}, r_eff::Real=1)
+    tail_length = _def_tail_length(length(is_ratios), r_eff)
+    return psis!(is_ratios, tail_length)
+end
+
+
 """
-    _def_tail_length(log_ratios::AbstractVector, r_eff::Real) -> tail_len::Integer
+    _def_tail_length(log_ratios::AbstractVector, r_eff::Real) -> Integer
 
 Define the tail length as in Vehtari et al. (2019).
 """
-function _def_tail_length(length::Int, r_eff::Real)
-    return Int(ceil(min(length / 5, 3 * sqrt(length / r_eff))))
+function _def_tail_length(length::Integer, r_eff::Real=1)
+    return min(cld(length, 5), ceil(3 * sqrt(length / r_eff))) |> Int
 end