Extend Distributions.canonform (#39)

* Extend `Distributions.canonform`

* Add tests

* Add comment

* Bump version

Project.toml

@@ -1,7 +1,7 @@
 name = "KeyedDistributions"
 uuid = "2576fb08-064d-4cab-b15d-8dda7fcb9a6d"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.1.11"
+version = "0.1.12"
 
 [deps]
 AutoHashEquals = "15f4f7f2-30c1-5605-9d31-71845cf9641f"

src/KeyedDistributions.jl

@@ -249,6 +249,11 @@ end
 Distributions.entropy(d::KeyedDistribution) = entropy(distribution(d))
 Distributions.entropy(d::KeyedDistribution, b::Real) = entropy(distribution(d), b)
 
+function Distributions.canonform(d::KeyedDistribution)
+    named_keys = NamedTuple{dimnames(d)}(axiskeys(d)) # Define/exists a shorthand for this?
+    return KeyedDistribution(canonform(distribution(d)); named_keys...)
+end
+
 # Univariate Distributions only
 
 for f in (:logpdf, :quantile, :mgf, :cf)

test/runtests.jl

@@ -212,6 +212,27 @@ using Test
 
             @test isapprox(rand(rng, kd), 0.39349598502717537)
             @test isapprox(rand(rng, kd, 2), [0.519693102856957, 0.6505773044249047])
+
+            @testset "canonform" begin
+                ckd = canonform(kd)
+
+                @test axiskeys(ckd) == axiskeys(kd)
+                @test dimnames(ckd) == dimnames(kd)
+                @test distribution(ckd) == canonform(d)
+            end
+        end
+
+        @testset "Multi-variate" begin
+            d = MvNormal([7.0, 5.0], Matrix([1.0 0.0; 0.0 1.0]))
+
+            @testset "canonform" begin
+                kd = KeyedDistribution(d; dim1=["x1", "x2"])
+                ckd = canonform(kd)
+
+                @test axiskeys(ckd) == axiskeys(kd)
+                @test dimnames(ckd) == dimnames(kd)
+                @test distribution(ckd) == canonform(d)
+            end
         end
 
         @testset "Matrix-variate" begin