From 5a058d0c7dd4fcaee46e078c26010390d99d7c42 Mon Sep 17 00:00:00 2001 From: Bagaev Dmitry Date: Mon, 10 Jun 2024 11:11:10 +0200 Subject: [PATCH 1/3] add inplace cov! --- Project.toml | 3 ++- src/statsfuns.jl | 31 +++++++++++++++++++++++++++++++ test/statsfuns_tests.jl | 30 ++++++++++++++++++++++++++++-- 3 files changed, 61 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 12f3446..e152eea 100644 --- a/Project.toml +++ b/Project.toml @@ -35,6 +35,7 @@ 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" @@ -42,4 +43,4 @@ 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"] diff --git a/src/statsfuns.jl b/src/statsfuns.jl index 5ecaea0..80b1e2e 100644 --- a/src/statsfuns.jl +++ b/src/statsfuns.jl @@ -326,3 +326,34 @@ 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. +""" +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' + + mean!(tmp2, Y') + @. tmp4 = Y - tmp2' + + corrected::Bool = true + # mul!(Z, tmp3', tmp4, 1, 0) + BLAS.gemm!('T', 'N', true, tmp3, tmp4, false, Z) + + b = 1//(size(tmp3, 1) - corrected) + @. Z = Z * b + + return Z +end \ No newline at end of file diff --git a/test/statsfuns_tests.jl b/test/statsfuns_tests.jl index 4211b96..7080aca 100644 --- a/test/statsfuns_tests.jl +++ b/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 \ No newline at end of file From abbc00a8fd8ed3e1432171bdb1bf17696c38a910 Mon Sep 17 00:00:00 2001 From: Bagaev Dmitry Date: Mon, 10 Jun 2024 11:12:09 +0200 Subject: [PATCH 2/3] add mcvo to the docs --- docs/src/index.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/index.md b/docs/src/index.md index 7552b0c..d006112 100644 --- a/docs/src/index.md +++ b/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 From 285fe41423409a750984ac579dcddbc640f31db0 Mon Sep 17 00:00:00 2001 From: Bagaev Dmitry Date: Mon, 10 Jun 2024 11:29:56 +0200 Subject: [PATCH 3/3] better performance --- src/BayesBase.jl | 2 +- src/statsfuns.jl | 25 ++++++++++++++++++++----- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/src/BayesBase.jl b/src/BayesBase.jl index 6df188d..c4a3489 100644 --- a/src/BayesBase.jl +++ b/src/BayesBase.jl @@ -1,6 +1,6 @@ module BayesBase -using TinyHugeNumbers +using TinyHugeNumbers, LoopVectorization using StatsAPI, StatsBase, DomainSets, Statistics, Distributions, Random import StatsAPI: params diff --git a/src/statsfuns.jl b/src/statsfuns.jl index 80b1e2e..2e6a363 100644 --- a/src/statsfuns.jl +++ b/src/statsfuns.jl @@ -332,6 +332,7 @@ end 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, @@ -343,17 +344,31 @@ function mcov!( tmp4=similar(Y), ) mean!(tmp1, X') - @. tmp3 = X - tmp1' + # @. 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' + # @. 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 - corrected::Bool = true # mul!(Z, tmp3', tmp4, 1, 0) BLAS.gemm!('T', 'N', true, tmp3, tmp4, false, Z) - b = 1//(size(tmp3, 1) - corrected) - @. Z = Z * b + 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 \ No newline at end of file