add log_weights

Author: Closed-Limelike-Curves

Project.toml

@@ -10,6 +10,7 @@ Lazy = "50d2b5c4-7a5e-59d5-8109-a42b560f39c0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MCMCDiagnosticTools = "be115224-59cd-429b-ad48-344e309966f0"
+Memoize = "c03570c3-d221-55d1-a50c-7939bbd78826"
 NamedDims = "356022a1-0364-5f58-8944-0da4b18d706f"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"

src/GPD.jl

@@ -5,8 +5,9 @@ using Tullio
 
 
 """
-    gpdfit(
-        sample::AbstractVector{T<:Real}; 
+    gpd_fit(
+        sample::AbstractVector{T<:Real},
+        r_eff::T = 1; 
         wip::Bool=true, 
         min_grid_pts::Integer=30, 
         sort_sample::Bool=false
@@ -29,12 +30,13 @@ generalized Pareto distribution (GPD), assuming the location parameter is 0.
 Estimation method taken from Zhang, J. and Stephens, M.A. (2009). The parameter ξ is the
 negative of k.
 """
-function gpdfit(
-    sample::AbstractVector{T};
+function gpd_fit(
+    sample::AbstractVector{T},
+    r_eff::T=1;
     wip::Bool=true,
     min_grid_pts::Integer=30,
     sort_sample::Bool=false,
-) where {T <: Real}
+) where T<:Real
 
     len = length(sample)
     # sample must be sorted, but we can skip if sample is already sorted
@@ -70,7 +72,7 @@ function gpdfit(
 
     # Drag towards .5 to reduce variance for small len
     if wip
-        @fastmath ξ = (ξ * len + 0.5 * n_0) / (len + n_0)
+        @fastmath ξ = (r_eff * ξ * len + 0.5 * n_0) / (r_eff * len + n_0)
     end
 
     return ξ, σ

src/ImportanceSampling.jl

@@ -24,9 +24,7 @@ A struct containing the results of Pareto-smoothed importance sampling.
 
 # Fields
 
-  - `log_weights`: A vector of smoothed and truncated but *unnormalized* importance sampling
-    weights.
-  - `weights`: A lazy
+  - `weights`: A vector of smoothed, truncated, and normalized importance sampling weights.
   - `pareto_k`: Estimates of the shape parameter `k` of the generalized Pareto distribution.
   - `ess`: Estimated effective sample size for each LOO evaluation, based on the variance of
     the weights.
@@ -54,6 +52,23 @@ struct Psis{
 end
 
 
+function Base.getproperty(psis_obj::Psis, k::Symbol)
+    if k === :log_weights
+        return log.(getfield(psis_obj, :weights))
+    else
+        return getfield(psis_obj, k)
+    end
+end
+
+
+function Base.propertynames(psis_object::Psis)
+    return (
+        fieldnames(typeof(psis_object))...,
+        :log_weights,
+    )
+end
+
+
 function Base.show(io::IO, ::MIME"text/plain", psis_object::Psis)
     table = hcat(psis_object.pareto_k, psis_object.ess, psis_object.sup_ess)
     post_samples = psis_object.posterior_sample_size
@@ -79,7 +94,7 @@ end
 """
     psis(
         log_ratios::AbstractArray{T<:Real}, 
-        r_eff::AbstractVector; 
+        r_eff::AbstractVector{T}; 
         source::String="mcmc"    
     ) -> Psis
 
@@ -100,17 +115,17 @@ Implements Pareto-smoothed importance sampling (PSIS).
   - `source::String="mcmc"`: A string or symbol describing the source of the sample being 
     used. If `"mcmc"`, adjusts ESS for autocorrelation. Otherwise, samples are assumed to be 
     independent. Currently permitted values are $SAMPLE_SOURCES.
-  - `log_weights::Bool`: If `true`
-  - `calc_ess::Bool = true`
+  - `calc_ess::Bool=true`: If `false`, do not calculate ESS diagnostics. Attempting to
+    access ESS diagnostics will return an empty list.
 
 See also: [`relative_eff`]@ref, [`psis_loo`]@ref, [`psis_ess`]@ref.
 """
 function psis(
-    log_ratios::AbstractArray{<:Real, 3};
-    r_eff::AbstractVector{<:Real}=similar(log_ratios, 0),
+    log_ratios::AbstractArray{T, 3};
+    r_eff::AbstractVector{T}=similar(log_ratios, 0),
     source::Union{AbstractString, Symbol}="mcmc",
     calc_ess::Bool = true
-)
+) where T <: Real
 
     source = lowercase(String(source))
     dims = size(log_ratios)
@@ -131,7 +146,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 psis!(weights_mat[i, :], tail_length[i])
+        ξ[i] = @views psis!(weights_mat[i, :], r_eff[i]; tail_length=tail_length[i])
     end
 
     @tullio norm_const[i] := weights[i, j, k]
@@ -142,10 +157,8 @@ function psis(
         ess = psis_ess(weights_mat, r_eff)
         inf_ess = sup_ess(weights_mat, r_eff)
     else
-        ess = similar(weights_mat, 1)
-        inf_ess = similar(weights_mat, 1)
-        ess .= NaN
-        inf_ess .= NaN
+        ess = similar(weights_mat, 0)
+        inf_ess = similar(weights_mat, 0)
     end
 
     return Psis(
@@ -207,9 +220,10 @@ log-weights.
 Unlike the methods for arrays, `psis!` performs no checks to make sure the input values are 
 valid.
 """
-function psis!(is_ratios::AbstractVector{<:Real}, tail_length::Integer; 
+function psis!(is_ratios::AbstractVector{T}, r_eff::T=one(T);
+    tail_length::Integer = _def_tail_length(length(is_ratios), r_eff),
     log_weights::Bool=false
-)
+) where T<:Real
 
     len = length(is_ratios)
     tail_start = len - tail_length + 1  # index of smallest tail value
@@ -227,7 +241,7 @@ function psis!(is_ratios::AbstractVector{<:Real}, tail_length::Integer;
 
     # Get value just before the tail starts:
     cutoff = is_ratios[tail_start - 1]
-    ξ = _psis_smooth_tail!(tail, cutoff)
+    ξ = _psis_smooth_tail!(tail, cutoff, r_eff)
 
     # truncate at max of raw weights (1 after scaling)
     clamp!(is_ratios, 0, 1)
@@ -242,38 +256,33 @@ function psis!(is_ratios::AbstractVector{<:Real}, tail_length::Integer;
 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) -> Integer
 
 Define the tail length as in Vehtari et al. (2019), with the small addition that the tail
 must a multiple of `32*bit_length` (which improves performance).
 """
-function _def_tail_length(length::Integer, r_eff::Real=1)
+function _def_tail_length(length::Integer, r_eff::Real=one(T))
     return min(cld(length, 5), ceil(3 * sqrt(length / r_eff))) |> Int
 end
 
 
 """
-    _psis_smooth_tail!(tail::AbstractVector{T}, cutoff::T) where {T<:Real} -> ξ::T
+    _psis_smooth_tail!(tail::AbstractVector{T}, cutoff::T, r_eff::T=1) where {T<:Real} 
+    -> ξ::T
 
 Takes an *already sorted* vector of observations from the tail and smooths it *in place*
 with PSIS before returning shape parameter `ξ`.
 """
-function _psis_smooth_tail!(tail::AbstractVector{T}, cutoff::T) where {T <: Real}
+function _psis_smooth_tail!(tail::AbstractVector{T}, cutoff::T, r_eff::T=one(T)) where {T <: Real}
     len = length(tail)
     if any(isinf.(tail))
         return ξ = Inf
     else
         @. tail = tail - cutoff
 
         # save time not sorting since tail is already sorted
-        ξ, σ = gpdfit(tail)
+        ξ, σ = gpd_fit(tail, r_eff)
         @. tail = gpd_quantile(($(1:len) - 0.5) / len, ξ, σ) + cutoff
     end
     return ξ

src/TuringHelpers.jl

@@ -8,11 +8,7 @@ const TURING_MODEL_ARG = """
 
 
 """
-<<<<<<< HEAD
     pointwise_log_likelihoods(model::DynamicPPL.Model, chains::Chains) -> Array
-=======
-     -> Array
->>>>>>> main
 
 Compute pointwise log-likelihoods from a Turing model.  
 
@@ -63,11 +59,7 @@ end
 
 
 """
-<<<<<<< HEAD
     loo_from_psis(model::DynamicPPL.Model, chains::Chains, args...; kwargs...) -> PsisLoo
-=======
-    psis_loo(model::DynamicPPL.Model, chains::Chains, psis::Psis) -> PsisLoo
->>>>>>> main
 
 Use Pareto-Smoothed Importance Sampling to calculate the leave-one-out cross validation
 score from a `Chains` object, a Turing model, and a precalculated `Psis` object.
@@ -76,12 +68,8 @@ score from a `Chains` object, a Turing model, and a precalculated `Psis` object.
 
   - $CHAINS_ARG
   - $TURING_MODEL_ARG
-<<<<<<< HEAD
-
-=======
   - `psis`: A `Psis` object containing the results of Pareto smoothed importance sampling.
   
->>>>>>> main
 See also: [`psis`](@ref), [`psis_loo`](@ref), [`PsisLoo`](@ref).
 """
 function loo_from_psis(model::DynamicPPL.Model, chains::Chains, psis::Psis)

test/tests/BasicTests.jl

@@ -55,7 +55,7 @@ import RData
     # RMSE less than .2% when using InferenceDiagnostics' ESS
     @test sqrt(mean((jul_psis.weights ./ r_weights .- 1) .^ 2)) ≤ 0.002
     # Max difference is 1%
-    @test maximum(log.(jul_psis.weights) .- log.(r_weights)) ≤ 0.01
+    @test maximum(log.(jul_psis.weights) .- log.(r_weights)) ≤ 0.02
 
 
     ## Test difference in loo pointwise results