From 929cd1c4849bc6af7a0602a60b4f968e04776de7 Mon Sep 17 00:00:00 2001
From: "C. Brenhin Keller" <cbkeller@dartmouth.edu>
Date: Mon, 11 Mar 2024 13:52:30 -0400
Subject: [PATCH] Add `movsum`

---
 Project.toml                 |  2 +-
 src/ArrayStats/ArrayStats.jl | 49 +++++++++++++++++++++++++++++++++---
 test/testArrayStats.jl       |  9 ++++++-
 3 files changed, 54 insertions(+), 6 deletions(-)

diff --git a/Project.toml b/Project.toml
index f20f224..9663ca1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "NaNStatistics"
 uuid = "b946abbf-3ea7-4610-9019-9858bfdeaf2d"
 authors = ["C. Brenhin Keller"]
-version = "0.6.34"
+version = "0.6.35"
 
 [deps]
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
diff --git a/src/ArrayStats/ArrayStats.jl b/src/ArrayStats/ArrayStats.jl
index 8a5a546..f5f8a96 100644
--- a/src/ArrayStats/ArrayStats.jl
+++ b/src/ArrayStats/ArrayStats.jl
@@ -508,6 +508,47 @@
     export nanstandardize
 
 
+## --- Moving sum, ignoring NaNs
+
+    """
+    ```julia
+    movsum(x::AbstractVecOrMat, n::Number)
+    ```
+    Simple moving sum of `x` in 1 or 2 dimensions, spanning `n` bins
+    (or n*n in 2D), returning an array of the same size as `x`.
+
+    For the resulting moving sum to be symmetric, `n` must be an odd integer;
+    if `n` is not an odd integer, the first odd integer greater than `n` will be
+    used instead.
+    """
+    function movsum(x::AbstractVector, n::Number)
+        m = similar(x)
+        δ = ceil(Int, (n-1)/2)
+        @inbounds for i ∈ eachindex(x)
+            iₗ = max(i-δ, firstindex(x))
+            iᵤ = min(i+δ, lastindex(x))
+            m[i] = nansum(view(x, iₗ:iᵤ))
+        end
+        return m
+    end
+    function movsum(x::AbstractMatrix, n::Number)
+        m = similar(x)
+        δ = ceil(Int, (n-1)/2)
+        𝐼 = repeat((firstindex(x,1):lastindex(x,1)), 1, size(x,2))
+        𝐽 = repeat((firstindex(x,2):lastindex(x,2))', size(x,1), 1)
+        @inbounds for k ∈ eachindex(𝐼,𝐽)
+            i = 𝐼[k]
+            iₗ = max(i-δ, firstindex(x,1))
+            iᵤ = min(i+δ, lastindex(x,1))
+            j = 𝐽[k]
+            jₗ = max(j-δ, firstindex(x,2))
+            jᵤ = min(j+δ, lastindex(x,2))
+            m[i,j] = nansum(view(x, iₗ:iᵤ, jₗ:jᵤ))
+        end
+        return m
+    end
+    export movsum
+
 ## -- Moving average, ignoring NaNs
 
     """
@@ -521,8 +562,8 @@
     if `n` is not an odd integer, the first odd integer greater than `n` will be
     used instead.
     """
-    function movmean(x::AbstractVector, n::Number)
-        mean_type = Base.promote_op(/, eltype(x), Int64)
+    function movmean(x::AbstractVector{T}, n::Number) where T
+        mean_type = Base.promote_op(/, T, Int64)
         m = Array{mean_type}(undef, size(x))
         δ = ceil(Int, (n-1)/2)
         @inbounds for i ∈ eachindex(x)
@@ -532,8 +573,8 @@
         end
         return m
     end
-    function movmean(x::AbstractMatrix, n::Number)
-        mean_type = Base.promote_op(/, eltype(x), Int64)
+    function movmean(x::AbstractMatrix{T}, n::Number) where T
+        mean_type = Base.promote_op(/, T, Int64)
         m = Array{mean_type}(undef, size(x))
         δ = ceil(Int, (n-1)/2)
         𝐼 = repeat((firstindex(x,1):lastindex(x,1)), 1, size(x,2))
diff --git a/test/testArrayStats.jl b/test/testArrayStats.jl
index d9747f2..4b543fe 100644
--- a/test/testArrayStats.jl
+++ b/test/testArrayStats.jl
@@ -404,7 +404,14 @@
     @test nanstandardize!(collect(1:10.)) ≈ ((1:10) .- mean(1:10)) / std(1:10)
     @test nanstandardize(1:10.) ≈ ((1:10) .- mean(1:10)) / std(1:10)
 
-## --- Moving averages
+## --- Moving sums and averages
+
+    # Moving sum: 1D
+    @test movsum(collect(1:10.),5) == movsum(1:10,5)
+    @test movsum(1:10,5) == [6, 10, 15, 20, 25, 30, 35, 40, 34, 27]
+    # Moving sum: 2D
+    @test movsum(repeat(1:5,1,5),3) == [6 9 9 9 6; 12 18 18 18 12; 18 27 27 27 18; 24 36 36 36 24; 18 27 27 27 18]
+    @test movsum(repeat(1:5,1,5)',3) == [6 12 18 24 18; 9 18 27 36 27; 9 18 27 36 27; 9 18 27 36 27; 6 12 18 24 18]
 
     # Moving average: 1D
     @test movmean(collect(1:10.),5) == movmean(1:10,5)