Add missing shape checks for the means argument to var[m] and std[m] (#32)

We use `@inbounds`, but the shape of the `means` argument was never checked
for the general `AbstractArray` method.
With an incorrect shape, invalid results or crashes would happen.
To avoid breaking existing code which was working, allow trailing singleton
dimensions. Sync the `SparseMatrixCSV` method, which makes it less strict.

Author: nalimilan

src/Statistics.jl

@@ -247,6 +247,11 @@ centralize_sumabs2(A::AbstractArray, m, ifirst::Int, ilast::Int) =
 function centralize_sumabs2!(R::AbstractArray{S}, A::AbstractArray, means::AbstractArray) where S
     # following the implementation of _mapreducedim! at base/reducedim.jl
     lsiz = Base.check_reducedims(R,A)
+    for i in 1:max(ndims(R), ndims(means))
+        if axes(means, i) != axes(R, i)
+            throw(DimensionMismatch("dimension $i of `mean` should have indices $(axes(R, i)), but got $(axes(means, i))"))
+        end
+    end
     isempty(R) || fill!(R, zero(S))
     isempty(A) && return R
 
@@ -295,9 +300,9 @@ function varm!(R::AbstractArray{S}, A::AbstractArray, m::AbstractArray; correcte
 end
 
 """
-    varm(itr, m; dims, corrected::Bool=true)
+    varm(itr, mean; dims, corrected::Bool=true)
 
-Compute the sample variance of collection `itr`, with known mean(s) `m`.
+Compute the sample variance of collection `itr`, with known mean(s) `mean`.
 
 The algorithm returns an estimator of the generative distribution's variance
 under the assumption that each entry of `itr` is an IID drawn from that generative
@@ -308,7 +313,8 @@ whereas the sum is scaled with `n` if `corrected` is
 `false` with `n` the number of elements in `itr`.
 
 If `itr` is an `AbstractArray`, `dims` can be provided to compute the variance
-over dimensions, and `m` may contain means for each dimension of `itr`.
+over dimensions. In that case, `mean` must be an array with the same shape as
+`mean(itr, dims=dims)` (additional trailing singleton dimensions are allowed).
 
 !!! note
     If array contains `NaN` or [`missing`](@ref) values, the result is also
@@ -331,7 +337,7 @@ end
 
 
 """
-    var(itr; dims, corrected::Bool=true, mean=nothing)
+    var(itr; corrected::Bool=true, mean=nothing[, dims])
 
 Compute the sample variance of collection `itr`.
 
@@ -343,10 +349,12 @@ If `corrected` is `true`, then the sum is scaled with `n-1`,
 whereas the sum is scaled with `n` if `corrected` is
 `false` with `n` the number of elements in `itr`.
 
-A pre-computed `mean` may be provided.
-
 If `itr` is an `AbstractArray`, `dims` can be provided to compute the variance
-over dimensions, and `mean` may contain means for each dimension of `itr`.
+over dimensions.
+
+A pre-computed `mean` may be provided. When `dims` is specified, `mean` must be
+an array with the same shape as `mean(itr, dims=dims)` (additional trailing
+singleton dimensions are allowed).
 
 !!! note
     If array contains `NaN` or [`missing`](@ref) values, the result is also
@@ -416,11 +424,13 @@ If `corrected` is `true`, then the sum is scaled with `n-1`,
 whereas the sum is scaled with `n` if `corrected` is
 `false` with `n` the number of elements in `itr`.
 
-A pre-computed `mean` may be provided.
-
 If `itr` is an `AbstractArray`, `dims` can be provided to compute the standard deviation
 over dimensions, and `means` may contain means for each dimension of `itr`.
 
+A pre-computed `mean` may be provided. When `dims` is specified, `mean` must be
+an array with the same shape as `mean(itr, dims=dims)` (additional trailing
+singleton dimensions are allowed).
+
 !!! note
     If array contains `NaN` or [`missing`](@ref) values, the result is also
     `NaN` or `missing` (`missing` takes precedence if array contains both).
@@ -445,9 +455,9 @@ std(iterable; corrected::Bool=true, mean=nothing) =
     sqrt(var(iterable, corrected=corrected, mean=mean))
 
 """
-    stdm(itr, m; corrected::Bool=true)
+    stdm(itr, mean; corrected::Bool=true)
 
-Compute the sample standard deviation of collection `itr`, with known mean(s) `m`.
+Compute the sample standard deviation of collection `itr`, with known mean(s) `mean`.
 
 The algorithm returns an estimator of the generative distribution's standard
 deviation under the assumption that each entry of `itr` is an IID drawn from that generative
@@ -457,10 +467,9 @@ If `corrected` is `true`, then the sum is scaled with `n-1`,
 whereas the sum is scaled with `n` if `corrected` is
 `false` with `n` the number of elements in `itr`.
 
-A pre-computed `mean` may be provided.
-
 If `itr` is an `AbstractArray`, `dims` can be provided to compute the standard deviation
-over dimensions, and `m` may contain means for each dimension of `itr`.
+over dimensions. In that case, `mean` must be an array with the same shape as
+`mean(itr, dims=dims)` (additional trailing singleton dimensions are allowed).
 
 !!! note
     If array contains `NaN` or [`missing`](@ref) values, the result is also
@@ -1065,7 +1074,11 @@ end
 function centralize_sumabs2!(R::AbstractArray{S}, A::SparseMatrixCSC{Tv,Ti}, means::AbstractArray) where {S,Tv,Ti}
     require_one_based_indexing(R, A, means)
     lsiz = Base.check_reducedims(R,A)
-    size(means) == size(R) || error("size of means must match size of R")
+    for i in 1:max(ndims(R), ndims(means))
+        if axes(means, i) != axes(R, i)
+            throw(DimensionMismatch("dimension $i of `mean` should have indices $(axes(R, i)), but got $(axes(means, i))"))
+        end
+    end
     isempty(R) || fill!(R, zero(S))
     isempty(A) && return R
 

test/runtests.jl

@@ -311,6 +311,22 @@ end
     @test var(Int[]) isa Float64
     @test isequal(var(skipmissing(Int[])), NaN)
     @test var(skipmissing(Int[])) isa Float64
+
+    # over dimensions with provided means
+    for x in ([1 2 3; 4 5 6], sparse([1 2 3; 4 5 6]))
+        @test var(x, dims=1, mean=mean(x, dims=1)) == var(x, dims=1)
+        @test var(x, dims=1, mean=reshape(mean(x, dims=1), 1, :, 1)) == var(x, dims=1)
+        @test var(x, dims=2, mean=mean(x, dims=2)) == var(x, dims=2)
+        @test var(x, dims=2, mean=reshape(mean(x, dims=2), :)) == var(x, dims=2)
+        @test var(x, dims=2, mean=reshape(mean(x, dims=2), :, 1, 1)) == var(x, dims=2)
+        @test_throws DimensionMismatch var(x, dims=1, mean=ones(size(x, 1)))
+        @test_throws DimensionMismatch var(x, dims=1, mean=ones(size(x, 1), 1))
+        @test_throws DimensionMismatch var(x, dims=2, mean=ones(1, size(x, 2)))
+        @test_throws DimensionMismatch var(x, dims=1, mean=ones(1, 1, size(x, 2)))
+        @test_throws DimensionMismatch var(x, dims=2, mean=ones(1, size(x, 2), 1))
+        @test_throws DimensionMismatch var(x, dims=2, mean=ones(size(x, 1), 1, 5))
+        @test_throws DimensionMismatch var(x, dims=1, mean=ones(1, size(x, 2), 5))
+    end
 end
 
 function safe_cov(x, y, zm::Bool, cr::Bool)