Basic abstract covariance estimation interface (#460)

Author: mateuszbaran

docs/src/cov.md

@@ -5,8 +5,13 @@ This package implements functions for computing scatter matrix, as well as weigh
 ```@docs
 scattermat
 cov
+cov(::CovarianceEstimator, ::AbstractVector)
+cov(::CovarianceEstimator, ::AbstractVector, ::AbstractVector)
+cov(::CovarianceEstimator, ::AbstractMatrix)
 cor
 mean_and_cov
 cov2cor
 cor2cov
+CovarianceEstimator
+SimpleCovariance
 ```

src/StatsBase.jl

@@ -98,6 +98,8 @@ export
     scattermat,     # scatter matrix (i.e. unnormalized covariance)
     cov2cor,        # converts a covariance matrix to a correlation matrix
     cor2cov,        # converts a correlation matrix to a covariance matrix
+    CovarianceEstimator, # abstract type for covariance estimators
+    SimpleCovariance,    # simple covariance estimator
 
     ## counts
     addcounts!,     # add counts to an accumulating array or map

src/cov.jl

@@ -153,3 +153,83 @@ function cor2cov!(C::AbstractMatrix, s::AbstractArray)
     end
     return C
 end
+
+"""
+    CovarianceEstimator
+
+Abstract type for covariance estimators.
+"""
+abstract type CovarianceEstimator end
+
+"""
+    cov(ce::CovarianceEstimator, x::AbstractVector; mean=nothing)
+
+Compute a variance estimate from the observation vector `x` using the  estimator `ce`.
+"""
+cov(ce::CovarianceEstimator, x::AbstractVector; mean=nothing) =
+    error("cov is not defined for $(typeof(ce)) and $(typeof(x))")
+
+"""
+    cov(ce::CovarianceEstimator, x::AbstractVector, y::AbstractVector)
+
+Compute the covariance of the vectors `x` and `y` using estimator `ce`.
+"""
+cov(ce::CovarianceEstimator, x::AbstractVector, y::AbstractVector) =
+    error("cov is not defined for $(typeof(ce)), $(typeof(x)) and $(typeof(y))")
+
+"""
+    cov(ce::CovarianceEstimator, X::AbstractMatrix, [w::AbstractWeights]; mean=nothing, dims::Int=1)
+
+Compute the covariance matrix of the matrix `X` along dimension `dims`
+using estimator `ce`. A weighting vector `w` can be specified.
+The keyword argument `mean` can be:
+
+* `nothing` (default) in which case the mean is estimated and subtracted
+  from the data `X`,
+* a precalculated mean in which case it is subtracted from the data `X`.
+  Assuming `size(X)` is `(N,M)`, `mean` can either be:
+  * when `dims=1`, an `AbstractMatrix` of size `(1,M)`,
+  * when `dims=2`, an `AbstractVector` of length `N` or an `AbstractMatrix`
+    of size `(N,1)`.
+"""
+cov(ce::CovarianceEstimator, X::AbstractMatrix; mean=nothing, dims::Int=1) =
+    error("cov is not defined for $(typeof(ce)) and $(typeof(X))")
+
+cov(ce::CovarianceEstimator, X::AbstractMatrix, w::AbstractWeights; mean=nothing, dims::Int=1) =
+    error("cov is not defined for $(typeof(ce)), $(typeof(X)) and $(typeof(w))")
+
+"""
+    SimpleCovariance(;corrected::Bool=false)
+
+Simple covariance estimator. Estimation calls `cov(x; corrected=corrected)`,
+`cov(x, y; corrected=corrected)` or `cov(X, w, dims; corrected=corrected)`
+where `x`, `y` are vectors, `X` is a matrix and `w` is a weighting vector.
+"""
+struct SimpleCovariance <: CovarianceEstimator
+    corrected::Bool
+    SimpleCovariance(;corrected::Bool=false) = new(corrected)
+end
+
+cov(sc::SimpleCovariance, x::AbstractVector) =
+    cov(x; corrected=sc.corrected)
+
+cov(sc::SimpleCovariance, x::AbstractVector, y::AbstractVector) =
+    cov(x, y; corrected=sc.corrected)
+
+function cov(sc::SimpleCovariance, X::AbstractMatrix; dims::Int=1, mean=nothing)
+    dims ∈ (1, 2) || throw(ArgumentError("Argument dims can only be 1 or 2 (given: $dims)"))
+    if mean === nothing
+        return cov(X; dims=dims, corrected=sc.corrected)
+    else
+        return covm(X, mean, dims, corrected=sc.corrected)
+    end
+end
+
+function cov(sc::SimpleCovariance, X::AbstractMatrix, w::AbstractWeights; dims::Int=1, mean=nothing)
+    dims ∈ (1, 2) || throw(ArgumentError("Argument dims can only be 1 or 2 (given: $dims)"))
+    if mean === nothing
+        return cov(X, w, dims, corrected=sc.corrected)
+    else
+        return covm(X, mean, w, dims, corrected=sc.corrected)
+    end
+end

test/cov.jl

@@ -1,6 +1,8 @@
 using StatsBase
 using LinearAlgebra, Random, Test
 
+struct EmptyCovarianceEstimator <: CovarianceEstimator end
+
 @testset "StatsBase.Covariance" begin
 weight_funcs = (weights, aweights, fweights, pweights)
 
@@ -248,5 +250,50 @@ weight_funcs = (weights, aweights, fweights, pweights)
         @test cor(X, wv1, 1) ≈ expected_cor1
         @test cor(X, wv2, 2) ≈ expected_cor2
     end
+
+    @testset "Abstract covariance estimation" begin
+        Xm1 = mean(X, dims=1)
+        Xm2 = mean(X, dims=2)
+
+        for corrected ∈ (false, true)
+            scc = SimpleCovariance(corrected=corrected)
+            @test_throws ArgumentError cov(scc, X, dims=0)
+            @test_throws ArgumentError cov(scc, X, wv1, dims=0)
+            @test cov(scc, X) ≈ cov(X, corrected=corrected)
+            @test cov(scc, X, mean=Xm1) ≈ StatsBase.covm(X, Xm1, corrected=corrected)
+            @test cov(scc, X, mean=Xm2, dims=2) ≈ StatsBase.covm(X, Xm2, 2, corrected=corrected)
+            if f !== weights || corrected === false
+                @test cov(scc, X, wv1, dims=1) ≈ cov(X, wv1, 1, corrected=corrected)
+                @test cov(scc, X, wv2, dims=2) ≈ cov(X, wv2, 2, corrected=corrected)
+                @test cov(scc, X, wv1, mean=Xm1) ≈ StatsBase.covm(X, Xm1, wv1, corrected=corrected)
+                @test cov(scc, X, wv2, mean=Xm2, dims=2) ≈ StatsBase.covm(X, Xm2, wv2, 2, corrected=corrected)
+            end
+        end
+    end
+end
+
+@testset "Abstract covariance estimation" begin
+    est = EmptyCovarianceEstimator()
+    wv = fweights(rand(2))
+    @test_throws ErrorException cov(est, [1.0 2.0; 3.0 4.0])
+    @test_throws ErrorException cov(est, [1.0 2.0; 3.0 4.0], wv)
+    @test_throws ErrorException cov(est, [1.0 2.0; 3.0 4.0], dims = 2)
+    @test_throws ErrorException cov(est, [1.0 2.0; 3.0 4.0], wv, dims = 2)
+    @test_throws ErrorException cov(est, [1.0 2.0; 3.0 4.0], mean = nothing)
+    @test_throws ErrorException cov(est, [1.0 2.0; 3.0 4.0], wv, mean = nothing)
+    @test_throws ErrorException cov(est, [1.0 2.0; 3.0 4.0], dims = 2, mean = nothing)
+    @test_throws ErrorException cov(est, [1.0 2.0; 3.0 4.0], wv, dims = 2, mean = nothing)
+    @test_throws ErrorException cov(est, [1.0, 2.0], [3.0, 4.0])
+    @test_throws ErrorException cov(est, [1.0, 2.0])
+
+    x = rand(8)
+    y = rand(8)
+
+    for corrected ∈ (false, true)
+        @test_throws MethodError SimpleCovariance(corrected)
+        scc = SimpleCovariance(corrected=corrected)
+        @test cov(scc, x) ≈ cov(x; corrected=corrected)
+        @test cov(scc, x, y) ≈ cov(x, y; corrected=corrected)
+    end
 end
 end # @testset "StatsBase.Covariance"