Faster method to compute the mean of a vector of Measurements

src/MatrixProductBP.jl

@@ -12,7 +12,7 @@ using UnPack: @unpack
 using Random: shuffle!, AbstractRNG, GLOBAL_RNG
 using SparseArrays: rowvals, nonzeros, nzrange, sparse
 using Distributions: Distributions, sample, sample!
-using Measurements: Measurement, ±
+using Measurements: Measurement, ±, value, uncertainty
 using Statistics: mean, std
 using Unzip: unzip
 using StatsBase: weights, proportions
@@ -56,7 +56,8 @@ export
     continuous_sis_sampler, simulate_queue_sis!,
     draw_node_observations!, AtomicVector,
     RecursiveBPFactor, DampedFactor, RecursiveTraceFactor, GenericFactor,
-    RestrictedRecursiveBPFactor
+    RestrictedRecursiveBPFactor,
+    mean_with_uncertainty
 
 
 include("utils.jl")

src/sampling.jl

@@ -158,7 +158,8 @@ function autocorrelations(f, sms::SoftMarginSampler; showprogress::Bool=true,
     dt = showprogress ? 0.1 : Inf
     prog = Progress(N, desc="Autocorrelations from Soft Margin"; dt)
 
-    for (a,i) in pairs(sites)
+    @threads for a in eachindex(sites)
+        i = sites[a]
         for u in axes(r[a], 2), t in max(1,u-maxdist):u-1
             q = nstates(bp, i)
             mtu_avg = zeros(q, q)

src/utils.jl

@@ -21,3 +21,15 @@ end
 sample_noalloc(w) = sample_noalloc(GLOBAL_RNG, w)
 
 
+# Computes the mean of a vector of statistically independent `Measurement`s. Equivalent to `mean(X)`, but faster
+function mean_with_uncertainty(x)
+    v = mean(value(a) for a in x)
+    e = sum(abs2, uncertainty(a) for a in x)
+    s = sqrt(e) / length(x)
+    return v ± s
+end
+function mean_with_uncertainty(x::AbstractArray{<:AbstractArray{<:Measurement}})
+    map(Iterators.product(axes(x[1])...)) do i
+        mean_with_uncertainty(xx[i...] for xx in x)
+    end
+end

test/sampling.jl

@@ -22,6 +22,10 @@
 
     sms = sample(bp, 10; showprogress=false, rng)
     m = marginals(sms)
+    @test all(m) do mᵢ
+        mean.(mᵢ) ≈ mean_with_uncertainty.(mᵢ)
+    end
+    @test mean.(m) ≈ mean_with_uncertainty.(m)
     pb = pair_marginals(sms)
 
     f(x,i) = x-1