Remove LoopVectorization.jl (#47)

* Remove LoopVectorization.jl

* Remove FFTW dep

* Accidentally added PSIS.jl

* Remove FFTW

* Fix method overwrite

* fix method overwwrite

* Remove StatsFuns dependency

Author: Carlos Parada

Project.toml

@@ -5,37 +5,26 @@ version = "0.6.6"
 
 [deps]
 AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
-Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
-FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
+LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MCMCDiagnosticTools = "be115224-59cd-429b-ad48-344e309966f0"
 NamedDims = "356022a1-0364-5f58-8944-0da4b18d706f"
-Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
-StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [compat]
 AxisKeys = "0.1.18"
-Distributions = "0.25.10"
-DocStringExtensions = "0.8"
-FFTW = "1.4.3"
-LoopVectorization = "0.12.37"
 MCMCDiagnosticTools = "0.1.0"
 NamedDims = "0.2.35"
-Polyester = "0.3.4, 0.4, 0.5"
 PrettyTables = "1.1.0"
 Requires = "1.1.3"
 StatsBase = "0.33.10"
-StatsFuns = "0.9.9"
 Tullio = "0.3.0"
 julia = "1.6"
 

src/ESS.jl

@@ -1,37 +1,30 @@
-using FFTW
-using LoopVectorization
 using MCMCDiagnosticTools
-
 using Tullio
 
 export relative_eff, psis_ess, sup_ess
 
 """
-    relative_eff(
-        sample::AbstractArray{Real, 3}; 
-        method=MCMCDiagnosticTools.FFTESSMethod()
-    )
+    relative_eff(sample::AbstractArray{<:Real, 3}; [method])
 
 Calculate the relative efficiency of an MCMC chain, i.e. the effective sample size divided
-by the nominal sample size.
+by the nominal sample size. If none is provided, use the default method from 
+MCMCDiagnosticTools.
 """
-function relative_eff(
-    sample::AbstractArray{T, 3}; method=MCMCDiagnosticTools.FFTESSMethod()
-) where {T <: Union{Real, Missing}}
+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; method=method, maxlag=dims[2])
+    ess, = MCMCDiagnosticTools.ess_rhat(ess_sample; maxlag=maxlag, kwargs...)
     r_eff = ess / post_sample_size
     return r_eff
 end
 
 
 """
     function psis_ess(
-        weights::AbstractVector{T},
-        r_eff::AbstractVector{T}
-    ) -> AbstractVector{T}
+        weights::AbstractVector{<:Real},
+        r_eff::AbstractVector{<:Real}
+    ) -> AbstractVector{<:Real}
 
 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 
@@ -44,23 +37,15 @@ 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 @turbo r_eff ./ sum_of_squares
-end
-
-
 function psis_ess(
     weights::AbstractMatrix{T}, r_eff::AbstractVector{T}
 ) where {T <: Union{Real, Missing}}
     @tullio sum_of_squares[x] := weights[x, y]^2
-    return @turbo r_eff ./ sum_of_squares
+    return r_eff ./ sum_of_squares
 end
 
 
-function psis_ess(weights::AbstractMatrix{<:Union{Real, Missing}})
+function psis_ess(weights::AbstractMatrix{<:Real})
     @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)))
@@ -69,8 +54,8 @@ end
 
 """
     function sup_ess(
-        weights::AbstractVector{T},
-        r_eff::AbstractVector{T}
+        weights::AbstractVector{<:Real},
+        r_eff::AbstractVector{<:Real}
     ) -> AbstractVector
 
 Calculate the supremum-based effective sample size of a PSIS sample, i.e. the inverse of the
@@ -79,10 +64,11 @@ 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 from which PSIS samples were 
+    derived.
 """
 function sup_ess(
     weights::AbstractMatrix{T}, r_eff::V
- ) where {T<:Union{Real, Missing}, V<:AbstractVector{T}}
-    return @turbo inv.(dropdims(maximum(weights; dims=2); dims=2)) .* r_eff
+) where {T<:Real, V<:AbstractVector{T}}
+    return inv.(dropdims(maximum(weights; dims=2); dims=2)) .* r_eff
 end

src/GPD.jl

@@ -1,5 +1,5 @@
 using LinearAlgebra
-using LoopVectorization
+using LogExpFunctions
 using Statistics
 using Tullio
 
@@ -45,22 +45,23 @@ function gpdfit(
 
     grid_size = min_grid_pts + isqrt(len)  # isqrt = floor sqrt
     n_0 = 10  # determines how strongly to nudge ξ towards .5
-    x_star::T = inv(3 * sample[(len + 2) ÷ 4])  # magic number. ¯\_(ツ)_/¯
-
+    x_star = inv(3 * sample[(len + 2) ÷ 4])  # magic number. ¯\_(ツ)_/¯
+    invmax = inv(sample[len])
 
     # build pointwise estimates of ξ and θ at each grid point
     θ_hats = similar(sample, grid_size)
-    ξ_hats = similar(sample, grid_size)
-    invmax = inv(sample[len])
-    @tullio threads=false θ_hats[i] = invmax + (1 - sqrt((grid_size + 1) / i)) * x_star
-    @tullio threads=false ξ_hats[i] = log1p(-θ_hats[i] * sample[j]) |> _ / len
-    log_like = similar(ξ_hats)
+    @fastmath @. θ_hats = invmax + (1 - sqrt((grid_size + 1) / $(1:grid_size))) * x_star
+    @tullio threads=false ξ_hats[i] := log1p(-θ_hats[i] * sample[j])
+    ξ_hats /= len
+
+    log_like = ξ_hats  # Reuse preallocated array
     # Calculate profile log-likelihood at each estimate:
-    @tullio threads=false log_like[i] =
+    @tullio threads=false ξ_hats[i] =
         len * (log(-θ_hats[i] / ξ_hats[i]) - ξ_hats[i] - 1)
     # Calculate weights from log-likelihood:
-    weights = ξ_hats  # Reuse preallocated array
-    @tullio threads=false weights[y] = exp(log_like[x] - log_like[y]) |> inv
+    weights = log_like  # Reuse preallocated array
+    log_norm = logsumexp(log_like)
+    @tullio threads=false log_like[x] = exp(log_like[x] - log_norm)
     # Take weighted mean:
     @tullio threads=false θ_hat := weights[x] * θ_hats[x]
     @tullio threads=false ξ := log1p(-θ_hat * sample[i])
@@ -72,7 +73,7 @@ function gpdfit(
         @fastmath ξ = (ξ * len + 0.5 * n_0) / (len + n_0)
     end
 
-    return ξ::T, σ::T
+    return ξ, σ
 
 end
 

src/ImportanceSampling.jl

@@ -1,5 +1,3 @@
-using LoopVectorization
-using StatsBase
 using Tullio
 
 const LIKELY_ERROR_CAUSES = """
@@ -118,8 +116,9 @@ function psis(
     r_eff = _generate_r_eff(log_ratios, dims, r_eff, source)
     weights = similar(log_ratios)
     # Shift ratios by maximum to prevent overflow
-    @tturbo @. weights = exp(log_ratios - $maximum(log_ratios; dims=2))
+    @. weights = exp(log_ratios - $maximum(log_ratios; dims=2))
 
+    r_eff = _generate_r_eff(weights, dims, r_eff, source)
     _check_input_validity_psis(reshape(log_ratios, dims), r_eff)
 
     tail_length = Vector{Int}(undef, data_size)
@@ -130,7 +129,7 @@ function psis(
     end
 
     @tullio norm_const[i] := weights[i, j]
-    @tturbo weights .= weights ./ norm_const
+    @. weights = weights / norm_const
     ess = psis_ess(weights, r_eff)
     inf_ess = sup_ess(weights, r_eff)
 
@@ -151,9 +150,10 @@ end
 
 function psis(
     log_ratios::AbstractMatrix{<:Real};
-    chain_index::AbstractVector{<:Integer}=_assume_one_chain(log_ratios),
+    chain_index::AbstractVector=_assume_one_chain(log_ratios),
     kwargs...,
 )
+    chain_index = Vector(Int.(chain_index))
     new_log_ratios = _convert_to_array(log_ratios, chain_index)
     return psis(new_log_ratios; kwargs...)
 end
@@ -166,7 +166,7 @@ Do PSIS on a single vector, smoothing its tail values.
 
 # Arguments
 
-  - `is_ratios::AbstractVector{Real}`: A vector of importance sampling ratios,
+  - `is_ratios::AbstractVector{<:Real}`: A vector of importance sampling ratios,
     scaled to have a maximum of 1.
 
 # Returns
@@ -219,11 +219,11 @@ function _psis_smooth_tail!(tail::AbstractVector{T}, cutoff::T) where {T <: Real
     if any(isinf.(tail))
         return ξ = Inf
     else
-        @turbo @. tail = tail - cutoff
+        @. tail = tail - cutoff
 
         # save time not sorting since tail is already sorted
         ξ, σ = gpdfit(tail)
-        @turbo @. tail = gpd_quantile(($(1:len) - 0.5) / len, ξ, σ) + cutoff
+        @. tail = gpd_quantile(($(1:len) - 0.5) / len, ξ, σ) + cutoff
     end
     return ξ
 end

src/InternalHelpers.jl

@@ -1,5 +1,5 @@
 const CHAIN_INDEX_DOC = """
-`chain_index::Vector`: An optional vector of integers specifying which chain each step
+`chain_index::Vector{Int}`: An optional vector of integers specifying which chain each step
 belongs to. For instance, `chain_index[step]` should return `2` if `log_likelihood[:, step]`
 belongs to the second chain.
 """
@@ -14,9 +14,10 @@ of that point. This function must take the form `f(θ[1], ..., θ[n], data)`, wh
 parameter vector. See also the `splat` keyword argument.
 """
 
-const LIKELIHOOD_ARRAY_ARG = """
+const LOG_LIK_ARR = """
 `log_likelihood::Array`: A matrix or 3d array of log-likelihood values indexed as
-`[data, step, chain]`. See the `chain_index` argument if leaving the `chain` index off.
+`[data, step, chain]`. The chain argument can be left off if `chain_index` is provided
+or if all posterior samples were drawn from a single chain.
 """
 
 const R_EFF_DOC = """

src/LeaveOneOut.jl

@@ -1,7 +1,5 @@
 using AxisKeys
-using Distributions
 using InteractiveUtils
-using LoopVectorization
 using NamedDims
 using Statistics
 using Printf
@@ -90,16 +88,16 @@ end
 
 """
     function psis_loo(
-        log_likelihood::Array{Real} [, args...];
-        [, chain_index::Vector{Integer}, kwargs...]
+        log_likelihood::AbstractArray{<:Real} [, args...];
+        [, chain_index::Vector{Int}, kwargs...]
     ) -> PsisLoo
 
 Use Pareto-Smoothed Importance Sampling to calculate the leave-one-out cross validation
 score.
 
 # Arguments
 
-  - $LIKELIHOOD_ARRAY_ARG
+  - $LOG_LIK_ARR
   - $ARGS [`psis`](@ref).
   - $CHAIN_INDEX_DOC
   - $KWARGS [`psis`](@ref).
@@ -110,30 +108,31 @@ function psis_loo(log_likelihood::AbstractArray{<:Real, 3}, args...; kwargs...)
     psis_object = psis(-log_likelihood, args...; kwargs...)
     return loo_from_psis(log_likelihood, psis_object)
 end
-psis_loo
+
 
 function psis_loo(
     log_likelihood::AbstractMatrix{<:Real},
     args...;
     chain_index::AbstractVector=_assume_one_chain(log_likelihood),
     kwargs...,
 )
+    chain_index = Int.(chain_index)
     new_log_ratios = _convert_to_array(log_likelihood, chain_index)
     return psis_loo(new_log_ratios, args...; kwargs...)
 end
 
 
 """
     loo_from_psis(
-        log_likelihood::AbstractArray, psis_object::Psis; 
-        chain_index::AbstractVector{Integer}
+        log_likelihood::AbstractArray{<:Real}, psis_object::Psis; 
+        chain_index::Vector{<:Integer}
     )
 
 Use a precalculated `Psis` object to estimate the leave-one-out cross validation score.
 
 # Arguments
 
-  - $LIKELIHOOD_ARRAY_ARG
+  - $LOG_LIK_ARR
   - `psis_object`: A precomputed `Psis` object used to estimate the LOO-CV score.
   - $CHAIN_INDEX_DOC
 
@@ -171,9 +170,8 @@ function loo_from_psis(log_likelihood::AbstractArray{<:Real, 3}, psis_object::Ps
     table = _generate_loo_table(pointwise)
 
     gmpd = exp.(table(column=:mean, statistic=:cv_elpd))
-    @tullio mcse := pointwise_mcse[i]^2 
-    mcse = sqrt(mcse)
 
+    mcse = sum(abs2, pointwise_mcse) |> sqrt
     return PsisLoo(table, pointwise, psis_object, gmpd, mcse)
 end
 
@@ -182,6 +180,7 @@ function loo_from_psis(
     log_likelihood::AbstractMatrix{<:Real}, psis_object::Psis, args...;
     chain_index::AbstractVector=_assume_one_chain(log_likelihood), kwargs...
 )
+    chain_index = Int.(chain_index)
     new_log_ratios = _convert_to_array(log_likelihood, chain_index)
     return loo_from_psis(new_log_ratios, psis_object, args...; kwargs...)
 end
@@ -227,9 +226,9 @@ function _calc_mcse(weights, log_likelihood, pointwise_loo, r_eff)
     pointwise_gmpd = exp.(pointwise_loo)
     @tullio pointwise_var[i] := 
         (weights[i, j, k] * (exp(log_likelihood[i, j, k]) - pointwise_gmpd[i]))^2
-    # If MCMC draws follow a log-normal distribution, we can use method of moments to est
-    # the standard deviation of their log:
-    @turbo @. pointwise_var = log1p(pointwise_var / pointwise_gmpd^2)
+    # If MCMC draws follow a log-normal distribution, then their log has this std. error:
+    @. pointwise_var = log1p(pointwise_var / pointwise_gmpd^2)
+    # (google "log-normal method of moments" for a proof)
     # apply MCMC correlation correction:
-    return @turbo @. sqrt(pointwise_var / r_eff)
+    return @. sqrt(pointwise_var / r_eff)
 end

src/ModelComparison.jl

@@ -1,5 +1,3 @@
-using StatsFuns
-using LoopVectorization
 import Base.show
 
 export loo_compare, ModelComparison
@@ -47,7 +45,7 @@ end
 
 """
     function loo_compare(
-        cv_results::PsisLoo...;
+        cv_results...;
         sort_models::Bool=true,
         best_to_worst::Bool=true,
         [, model_names::Tuple{Symbol}]
@@ -111,13 +109,13 @@ function loo_compare(
     se_total = NamedTuple{name_tuple}(se_total)
 
     log_norm = logsumexp(cv_elpd)
-    weights = @turbo warn_check_args=false @. exp(cv_elpd - log_norm)
+    weights = @. exp(cv_elpd - log_norm)
 
-    gmpd = @turbo @. exp(cv_elpd / data_size)
+    gmpd = @. exp(cv_elpd / data_size)
     gmpd = NamedTuple{name_tuple}(gmpd)
 
-    @turbo warn_check_args=false @. cv_elpd = cv_elpd - cv_elpd[1]
-    @turbo warn_check_args=false avg_elpd = cv_elpd ./ data_size
+    @. cv_elpd = cv_elpd - cv_elpd[1]
+    avg_elpd = cv_elpd ./ data_size
     total_diffs = KeyedArray(
         hcat(cv_elpd, avg_elpd, weights);
         model=model_names,