From a5bdd0e50c97a7de0ebcdfa7ba768e00a1f99c67 Mon Sep 17 00:00:00 2001
From: "Thomas M. Huber" <55461592+biothomme@users.noreply.github.com>
Date: Fri, 15 Jul 2022 17:10:25 +0200
Subject: [PATCH] added a weighted mean

---
 src/NaNMath.jl | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/src/NaNMath.jl b/src/NaNMath.jl
index 994c539..ad19a35 100644
--- a/src/NaNMath.jl
+++ b/src/NaNMath.jl
@@ -212,6 +212,27 @@ NaNMath.mean([1., 2., NaN]) # result: 1.5
 function mean(x::AbstractArray{T}) where T<:AbstractFloat
     return mean_count(x)[1]
 end
+          
+"""
+NaNMath.mean(A, W)
+
+##### Args:
+* `A`: An array of floating point numbers
+* `W`: An array of floating point weights corresponding to the values of `A`
+
+##### Returns:
+*    Returns the arithmetic weighted mean of all elements in the array, ignoring NaN's.
+
+##### Examples:
+```julia
+using NaNMath
+NaNMath.mean([1., 2., NaN], [1., .5, 1]) # result: 1.0
+```
+"""
+function mean(x::AbstractArray{T}, w::AbstractArray{T}) where T<:AbstractFloat
+    return mean_count(x, w)[1]
+end
+
 
 """
 Returns a tuple of the arithmetic mean of all elements in the array, ignoring NaN's,
@@ -229,6 +250,22 @@ function mean_count(x::AbstractArray{T}) where T<:AbstractFloat
     return (result, count)
 end
 
+"""
+Returns a tuple of the arithmetic mean of all elements in the array, ignoring NaN's,
+and the sum of weights of non-NaN values and non-NaN weights in the array.
+"""
+function mean_count(x::AbstractArray{T}, w::AbstractArray{T}) where T<:AbstractFloat
+    zx, zw = zero(eltype(x)), zero(eltype(z))
+    sum = zx
+    weightsum = zw
+    @simd for i, j in zip(x, w)
+        weightsum += ifelse(isnan(i)||isnan(j), zw, j)
+        sum += ifelse(isnan(i)||isnan(j), z, i)
+    end
+    result = sum / count
+    return (result, count)
+end
+
 """
 NaNMath.var(A)