Merge pull request #102 from Julia-Tempering/generalize-initialization

add a DistributionLogPotential

Author: miguelbiron

src/Pigeons.jl

@@ -66,6 +66,8 @@ include("includes.jl")
 export pigeons, Inputs, PT, 
     # for running jobs:
     ChildProcess, MPI,
+    # references:
+    DistributionLogPotential,
     # targets:
     TuringLogPotential, StanLogPotential,
     # some examples

src/includes.jl

@@ -26,6 +26,7 @@ include("paths/InterpolatedLogPotential.jl")
 include("paths/InterpolatingPath.jl")
 include("variational/variational.jl")
 include("pt/Inputs.jl")
+include("targets/DistributionLogPotential.jl")
 include("targets/StreamTarget.jl")
 include("pt/Shared.jl")
 include("swap/swap_graphs.jl")

src/replicas/replicas.jl

@@ -87,7 +87,7 @@ end
 function _create_locals(my_global_indices, inputs::Inputs, shared::Shared, ::Nothing)
     master_rng = SplittableRandom(inputs.seed)
     split_rngs = split_slice(my_global_indices, master_rng)
-    states = [initialization(inputs.target, split_rngs[i], my_global_indices[i]) for i in eachindex(split_rngs)]
+    states = [initialization(inputs, split_rngs[i], my_global_indices[i]) for i in eachindex(split_rngs)]
     recorders = [create_recorders(inputs, shared) for i in eachindex(split_rngs)]
     return Replica.(
                 states, 
@@ -96,3 +96,25 @@ function _create_locals(my_global_indices, inputs::Inputs, shared::Shared, ::Not
                 recorders, 
                 my_global_indices)  # <- replica indices
 end
+
+# default method: defer to user-provided method for their target
+initialization(inp::Inputs, args...) = initialization(inp.target, args...)
+
+# generic method for distribution-type references: sample iid for all replicas
+function initialization(
+    inp::Inputs{T, V, E, R},
+    rng::AbstractRNG,
+    ::Int
+    ) where {T, V, E, R <: DistributionLogPotential}
+    rand(rng, inp.reference.dist)
+end
+
+# for univariate references we need to wrap in vector 
+function initialization(
+    inp::Inputs{T, V, E, R},
+    rng::AbstractRNG,
+    ::Int
+    ) where {T, V, E, R <: DistributionLogPotential{<:UnivariateDistribution}}
+    [rand(rng, inp.reference.dist)]
+end
+

src/targets/DistributionLogPotential.jl

@@ -0,0 +1,25 @@
+"""
+Provides a reference type for Pigeons based on an encapsulated Distribution type.
+$FIELDS
+"""
+@auto struct DistributionLogPotential{D<:Distribution}
+    """The encapsulated distribution."""
+    dist::D
+end
+
+# evaluate the log density: general case
+(ref::DistributionLogPotential)(x) = logpdf(ref.dist, x)
+
+# univariate case
+(ref::DistributionLogPotential{<:UnivariateDistribution})(x) = logpdf(ref.dist, first(x))
+
+# iid sampling
+# general case
+function sample_iid!(ref::DistributionLogPotential, replica, shared)
+    rand!(replica.rng, ref.dist, replica.state)
+end
+
+# univariate case
+function sample_iid!(ref::DistributionLogPotential{D}, replica, shared) where {D<:UnivariateDistribution}
+    replica.state[begin] = rand(replica.rng, ref.dist)
+end

src/targets/target.jl

@@ -92,4 +92,4 @@ sample_iid!(reference_log_potential::InterpolatedLogPotential, replica, shared)
         sample_iid!(reference_log_potential.path.ref, replica, shared) 
     else
         error()
-    end
+    end

test/test_DistributionLogPotential.jl

@@ -0,0 +1,27 @@
+using MCMCChains
+
+@testset "DistributionLogPotential" begin
+    @testset "Multivariate" begin
+        function unid_log_potential(x; n_trials=100, n_successes=50)
+            p1, p2 = x
+            ((0 <= p1 <= 1) && (0 <= p2 <= 1)) || return typeof(p1)(-Inf)
+            p = p1 * p2
+            return n_successes*log(p) + (n_trials-n_successes)*log1p(-p)
+        end
+        ref_dist = product_distribution([Uniform(), Uniform()])
+        pt = pigeons(
+            target = unid_log_potential,
+            reference = DistributionLogPotential(ref_dist),
+            record = [traces]
+        )
+        @show Chains(sample_array(pt), variable_names(pt))
+    end
+    @testset "Univariate" begin
+        pt = pigeons(
+            target    = (x -> logpdf(Normal(3,1), x[begin])),
+            reference = DistributionLogPotential(Normal(-3,1)),
+            record    = [traces]
+        )
+        @show Chains(sample_array(pt), variable_names(pt))
+    end
+end