Draft: Add rekey method

Project.toml

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

src/KeyedDistributions.jl

@@ -7,10 +7,7 @@ using IterTools
 using Random: AbstractRNG
 
 export KeyedDistribution, KeyedSampleable
-export axiskeys, distribution
-
-
-# Constructors
+export axiskeys, distribution, rekey
 
 for T in (:Distribution, :Sampleable)
     KeyedT = Symbol(:Keyed, T)
@@ -50,7 +47,15 @@ for T in (:Distribution, :Sampleable)
         The elements of `keys` correspond to the variates of the distribution.
         """
         $KeyedT(d::$T{F, S}, keys::AbstractVector) where {F, S} = $KeyedT(d, (keys, ))
+
+        """
+            rekey(d::$($KeyedT), keys)
+
+        Returns a new [`$($KeyedT)`](@ref) with `keys`.
+        """
+        rekey(d::$KeyedT, keys)::$KeyedT = $KeyedT(distribution(d), keys)
     end
+
 end
 
 _size(d) = (length(d),)
@@ -91,6 +96,7 @@ Return the keys for the variates of the `KeyedDistribution` or `KeyedSampleable`
 """
 AxisKeys.axiskeys(d::KeyedDistOrSampleable) = d.keys
 
+
 # Standard functions to overload for new Distribution and/or Sampleable
 # https://juliastats.org/Distributions.jl/latest/extends/#Create-New-Samplers-and-Distributions
 

test/runtests.jl

@@ -51,6 +51,15 @@ using Test
                 @test ==(kd, T(d, keys))
             end
 
+            @testset "rekey" begin
+                new_keys = [:x, :y, :z]
+                for k in (new_keys, (new_keys,))
+                    new_kd = rekey(kd, k)
+                    @test new_kd isa T
+                    @test axiskeys(new_kd) == (new_keys,)
+                end
+            end
+
             @testset "sampling" begin
                 # Samples from the distribution both wrapped and unwrapped are the same
                 @test rand(StableRNG(1), d) == rand(StableRNG(1), kd)