Add inplace cov! computation

Project.toml

@@ -35,11 +35,12 @@ julia = "1.9"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CpuId = "adafc99b-e345-5852-983c-f28acb93d879"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "CpuId", "Test", "ReTestItems", "LinearAlgebra", "StableRNGs", "HCubature"]
+test = ["Aqua", "CpuId", "JET", "Test", "ReTestItems", "LinearAlgebra", "StableRNGs", "HCubature"]

docs/src/index.md

@@ -87,6 +87,7 @@ BayesBase.clamplog
 BayesBase.mvtrigamma
 BayesBase.dtanh
 BayesBase.probvec
+BayesBase.mcov!
 BayesBase.mean_std
 BayesBase.mean_var
 BayesBase.mean_cov

src/BayesBase.jl

@@ -1,6 +1,6 @@
 module BayesBase
 
-using TinyHugeNumbers
+using TinyHugeNumbers, LoopVectorization
 using StatsAPI, StatsBase, DomainSets, Statistics, Distributions, Random
 
 import StatsAPI: params

src/statsfuns.jl

@@ -326,3 +326,49 @@ Base.promote_rule(::Type{CountingReal}, ::Type{T}) where {T<:Real} = CountingRea
 function Base.:(==)(left::CountingReal{T}, right::CountingReal{T}) where {T}
     return (value(left) == value(right)) && (infinities(left) == infinities(right))
 end
+
+"""
+    mcov!(Z, X::AbstractMatrix, Y::AbstractMatrix; tmp1 = zeros(eltype(Z), size(X, 2)), tmp2 = zeros(eltype(Z), size(Y, 2)), tmp3 = similar(X), tmp4 = similar(Y))
+
+Same as `Statistics.cov(X, Y)`, but does not allocate the result. Instead uses a buffer `Z` to store the result in.
+Additionally, it allows for passing temporary buffers `tmp1`, `tmp2`, `tmp3`, `tmp4` to avoid any allocations.
+Always computes `corrected = true` covariance matrix.
+"""
+function mcov!(
+    Z,
+    X::AbstractMatrix,
+    Y::AbstractMatrix;
+    tmp1=zeros(eltype(Z), size(X, 2)),
+    tmp2=zeros(eltype(Z), size(Y, 2)),
+    tmp3=similar(X),
+    tmp4=similar(Y),
+)
+    mean!(tmp1, X')
+    # @. tmp3 = X - tmp1'
+    @turbo warn_check_args = false for j in 1:size(X, 2)
+        for i in 1:size(X, 1)
+            tmp3[i, j] = X[i, j] - tmp1[j]
+        end
+    end
+
+    mean!(tmp2, Y')
+    # @. tmp4 = Y - tmp2'
+    @turbo warn_check_args = false for j in 1:size(X, 2)
+        for i in 1:size(X, 1)
+            tmp4[i, j] = Y[i, j] - tmp2[j]
+        end
+    end
+
+    # mul!(Z, tmp3', tmp4, 1, 0)
+    BLAS.gemm!('T', 'N', true, tmp3, tmp4, false, Z)
+
+    b = 1//(size(tmp3, 1) - 1)
+    # @. Z = Z * b
+    @turbo warn_check_args = false for j in 1:size(Z, 2)
+        for i in 1:size(Z, 1)
+            Z[i, j] *= b
+        end
+    end
+
+    return Z
+end

test/statsfuns_tests.jl

@@ -26,10 +26,10 @@ end
     end
 end
 
-@testitem "dtanh" begin 
+@testitem "dtanh" begin
     for T in (Float32, Float64, BigFloat)
         foreach(rand(T, 10)) do number
-            @test dtanh(number) ≈ 1 - tanh(number) ^ 2
+            @test dtanh(number) ≈ 1 - tanh(number)^2
         end
     end
 end
@@ -87,6 +87,32 @@ end
 
         @test float(convert(CountingReal, r)) ≈ zero(T)
         @test float(convert(CountingReal{Float64}, r)) ≈ zero(Float64)
+    end
+end
+
+@testitem "mcov!" begin
+    using StatsFuns, BayesBase, JET
+    import BayesBase: mcov!
+
+    for n in 2:5:20, j in 3:5:20
+        X = rand(j, n)
+        Y = rand(j, n)
+        Z = rand(n, n)
+
+        @inferred(mcov!(Z, X, Y))
+
+        @test all(Z .≈ cov(X, Y))
+
+        tmp1 = zeros(eltype(Z), size(X, 2))
+        tmp2 = zeros(eltype(Z), size(Y, 2))
+        tmp3 = similar(X)
+        tmp4 = similar(Y)
+
+        @inferred(mcov!(Z, X, Y; tmp1=tmp1, tmp2=tmp2, tmp3=tmp3, tmp4=tmp4))
+
+        @test all(Z .≈ cov(X, Y))
 
+        @report_opt mcov!(Z, X, Y; tmp1=tmp1, tmp2=tmp2, tmp3=tmp3, tmp4=tmp4)
+        @test @allocated(mcov!(Z, X, Y; tmp1=tmp1, tmp2=tmp2, tmp3=tmp3, tmp4=tmp4)) === 0
     end
 end