From a7565c5d4efe245283fc3130d717027fb193c613 Mon Sep 17 00:00:00 2001
From: Milan Bouchet-Valat <nalimilan@club.fr>
Date: Sat, 26 Mar 2022 15:17:46 +0100
Subject: [PATCH 1/2] Improve performance of weighted sum

The current code is calling the `AbstractArray` matrix multiplication fallback,
which is slower than BLAS.
---
 src/weights.jl | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/weights.jl b/src/weights.jl
index 712951509..9841c3cc1 100644
--- a/src/weights.jl
+++ b/src/weights.jl
@@ -382,6 +382,14 @@ Compute the weighted sum of an array `v` with weights `w`, optionally over the d
 """
 wsum(v::AbstractArray, w::AbstractVector, dims::Colon=:) = transpose(w) * vec(v)
 
+# Optimized methods (to ensure we use BLAS when possible)
+for W in (AnalyticWeights, FrequencyWeights, ProbabilityWeights, Weights)
+    @eval begin
+        wsum(v::AbstractArray, w::$W, dims::Colon) = transpose(w.values) * vec(v)
+    end
+end
+wsum(v::AbstractArray, w::UnitWeights, dims::Colon) = sum(v)
+
 ## wsum along dimension
 #
 #  Brief explanation of the algorithm:

From a6768f7ba9e7210df974fc635f592857e94f134b Mon Sep 17 00:00:00 2001
From: Milan Bouchet-Valat <nalimilan@club.fr>
Date: Wed, 30 Mar 2022 21:47:41 +0200
Subject: [PATCH 2/2] Reorganize wsum and sum for unit weights, and add test

---
 src/weights.jl  | 12 +++++-------
 test/weights.jl |  2 +-
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/src/weights.jl b/src/weights.jl
index 9841c3cc1..089194b67 100644
--- a/src/weights.jl
+++ b/src/weights.jl
@@ -388,7 +388,11 @@ for W in (AnalyticWeights, FrequencyWeights, ProbabilityWeights, Weights)
         wsum(v::AbstractArray, w::$W, dims::Colon) = transpose(w.values) * vec(v)
     end
 end
-wsum(v::AbstractArray, w::UnitWeights, dims::Colon) = sum(v)
+
+function wsum(A::AbstractArray, w::UnitWeights, dims::Colon)
+    length(A) != length(w) && throw(DimensionMismatch("Inconsistent array dimension."))
+    return sum(A)
+end
 
 ## wsum along dimension
 #
@@ -613,12 +617,6 @@ optionally over the dimension `dims`.
 Base.sum(A::AbstractArray, w::AbstractWeights{<:Real}; dims::Union{Colon,Int}=:) =
     wsum(A, w, dims)
 
-function Base.sum(A::AbstractArray, w::UnitWeights; dims::Union{Colon,Int}=:)
-    a = (dims === :) ? length(A) : size(A, dims)
-    a != length(w) && throw(DimensionMismatch("Inconsistent array dimension."))
-    return sum(A, dims=dims)
-end
-
 ##### Weighted means #####
 
 function wmean(v::AbstractArray{<:Number}, w::AbstractVector)
diff --git a/test/weights.jl b/test/weights.jl
index 7fed9ef97..d2b8b830c 100644
--- a/test/weights.jl
+++ b/test/weights.jl
@@ -476,7 +476,7 @@ end
 @testset "Sum, mean, quantiles and variance for unit weights" begin
     wt = uweights(Float64, 3)
 
-    @test sum([1.0, 2.0, 3.0], wt) ≈ 6.0
+    @test sum([1.0, 2.0, 3.0], wt) ≈ wsum([1.0, 2.0, 3.0], wt) ≈ 6.0
     @test mean([1.0, 2.0, 3.0], wt) ≈ 2.0
 
     @test sum(a, wt, dims=1) ≈ sum(a, dims=1)