Merge pull request #43 from invenia/mz/var

fix var for distributions backed by KeyedArrays
invenia · Nov 3, 2022 · 479bf73 · 479bf73 · mzgubic · Nov 3, 2022
2 parents 4411448 + cf028f8
commit 479bf73
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "KeyedDistributions"
 uuid = "2576fb08-064d-4cab-b15d-8dda7fcb9a6d"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.10"
+version = "0.1.11"
 
 [deps]
 AutoHashEquals = "15f4f7f2-30c1-5605-9d31-71845cf9641f"

diff --git a/src/KeyedDistributions.jl b/src/KeyedDistributions.jl
@@ -235,7 +235,8 @@ function Distributions.mean(d::KeyedDistribution)
 end
 
 function Distributions.var(d::KeyedDistribution)
-    KeyedArray(NamedDimsArray(var(distribution(d)), dimnames(d)), axiskeys(d))
+    v = AxisKeys.keyless_unname(var(distribution(d)))
+    KeyedArray(NamedDimsArray(v, dimnames(d)), axiskeys(d))
 end
 
 for f in (:cov, :cor)

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -378,4 +378,12 @@ using Test
         @test ks == d
         @test d == ks
     end
+
+    @testset "mean and var when distribution is backed by KeyedArray" begin
+        ka = KeyedArray(rand(3), down=["a", "b", "c"])
+        kd1 = KeyedDistribution(MvNormal(ka), down=["a", "b", "c"])
+        kd2 = KeyedDistribution(MvNormal(AxisKeys.keyless_unname(ka)), down=["a", "b", "c"])
+        @test mean(kd1) == mean(kd2)
+        @test var(kd1) == var(kd2)
+    end
 end