Extend AxisKeys.named_axiskeys (#45)

* Add erroring test * Extend `AxisKeys.named_axiskeys` * Bump version * Use `named_axiskeys` in `_rand!` and `canonform` * Bump patch version
invenia · Nov 15, 2022 · 0d8a52d · 0d8a52d · mjp98 · Nov 15, 2022
1 parent 2813aa1
commit 0d8a52d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 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.13"
+version = "0.1.14"
 
 [deps]
 AutoHashEquals = "15f4f7f2-30c1-5605-9d31-71845cf9641f"

diff --git a/src/KeyedDistributions.jl b/src/KeyedDistributions.jl
@@ -148,6 +148,11 @@ AxisKeys.haskeys(d::KeyedDistOrSampleable) = true
 AxisKeys.keyless(d::KeyedDistOrSampleable) = distribution(d)
 AxisKeys.keyless_unname(d::KeyedDistOrSampleable) = distribution(d)
 
+function AxisKeys.named_axiskeys(d::KeyedDistOrSampleable)
+    NT = NamedTuple{dimnames(d)}
+    return NT(axiskeys(d))
+end
+
 # NamedDims functionality
 for T in (:Distribution, :Sampleable)
     KeyedT = Symbol(:Keyed, T)
@@ -175,7 +180,7 @@ function Distributions._rand!(
     x::AbstractVector{<:Real}
 )
     sample = Distributions._rand!(rng, distribution(d), x)
-    named_keys = NamedTuple{dimnames(d)}(axiskeys(d))
+    named_keys = named_axiskeys(d)
     return KeyedArray(sample; named_keys...)
 end
 
@@ -250,7 +255,7 @@ 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?
+    named_keys = named_axiskeys(d)
     return KeyedDistribution(canonform(distribution(d)); named_keys...)
 end
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -64,6 +64,7 @@ using Test
                 @test distribution(kd) == d
                 @test axiskeys(kd) == keys
                 @test dimnames(kd) == names
+                @test named_axiskeys(kd) == named_keys
                 @test length(kd) == length(d) == 3
                 @test isequal(kd, T(d; id=[:a, :b, :c]))
                 @test params(kd) == params(d) == (m, s)