setup as a test

Author: sunxd3

Project.toml

@@ -10,10 +10,12 @@ BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConsoleProgressMonitor = "88cd18e8-d9cc-4ea6-8889-5259c0d15c8b"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
@@ -34,10 +36,12 @@ Transducers = "0.4.30"
 julia = "1.6"
 
 [extras]
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 IJulia = "7073ff75-c697-5162-941a-fcdaad2a7d2a"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["FillArrays", "IJulia", "Statistics", "Test"]
+test = ["FillArrays", "Distributions", "IJulia", "OrderedCollections", "Statistics", "Test"]

docs/src/gibbs.md

@@ -345,3 +345,49 @@ Some points worth noting:
     - update the `vi` with the new values from the sampler state
 
 Again, the `state` interface in AbstractMCMC allows the Gibbs sampler to be agnostic of the details of the sampler state, and acquire the values of the parameters from individual sampler states.
+
+Now we can use the Gibbs sampler to sample from the hierarchical normal model.
+
+First we generate some data,
+
+```julia
+N = 100  # Number of data points
+mu_true = 0.5  # True mean
+tau2_true = 2.0  # True variance
+
+x_data = rand(Normal(mu_true, sqrt(tau2_true)), N)
+```
+
+```
+
+Then we can create a `HierNormal` model with the data.
+
+```julia
+hn = HierNormal((x=x_data,))
+```
+
+sampling is easy: we use random walk MH for `mu` and prior MH for `tau2`, because `tau2` has support on positive real numbers.
+
+```julia
+samples = sample(
+    hn,
+    Gibbs(
+        OrderedDict(
+            (:mu,) => RWMH(1),
+            (:tau2,) => PriorMH(product_distribution([InverseGamma(1, 1)])),
+        ),
+    ),
+    100000;
+    initial_params=(mu=[0.0], tau2=[1.0]),
+)
+```
+
+Then we can extract the samples and compute the mean of the samples.
+
+```julia
+mu_samples = [sample.values.mu for sample in samples][20001:end]
+tau2_samples = [sample.values.tau2 for sample in samples][20001:end]
+
+mean(mu_samples)
+mean(tau2_samples)
+```

test/gibbs_example/Project.toml

@@ -3,8 +3,11 @@ using Distributions
 using LogDensityProblems
 using OrderedCollections
 using Random
+using Test
 
-##
+include("hier_normal.jl")
+# include("gmm.jl")
+include("mh.jl")
 
 struct Gibbs <: AbstractMCMC.AbstractSampler
     sampler_map::OrderedDict
@@ -64,7 +67,7 @@ function AbstractMCMC.step(
     vi = state.vi
     for group in keys(spl.sampler_map)
         for (group, sub_state) in state.states
-            vi = merge(vi, unflatten(get_params(sub_state), group))
+            vi = merge(vi, unflatten(AbstractMCMC.get_params(sub_state), group))
         end
         sub_spl = spl.sampler_map[group]
         sub_state = state.states[group]
@@ -73,7 +76,7 @@ function AbstractMCMC.step(
             Tuple([vi[g] for g in group_complement])
         )
         cond_logdensity = condition(logdensity_model.logdensity, cond_val)
-        sub_state = recompute_logprob!!(cond_logdensity, get_params(sub_state), sub_state)
+        sub_state = recompute_logprob!!(cond_logdensity, AbstractMCMC.get_params(sub_state), sub_state)
         sub_state = last(
             AbstractMCMC.step(
                 rng,
@@ -87,15 +90,15 @@ function AbstractMCMC.step(
         state.states[group] = sub_state
     end
     for (group, sub_state) in state.states
-        vi = merge(vi, unflatten(get_params(sub_state), group))
+        vi = merge(vi, unflatten(AbstractMCMC.get_params(sub_state), group))
     end
     return GibbsTransition(vi), GibbsState(vi, state.states)
 end
 
-## tests
+## tests with hierarchical normal model
 
 # generate data
-N = 100  # Number of data points
+N = 1000  # Number of data points
 mu_true = 0.5  # True mean
 tau2_true = 2.0  # True variance
 
@@ -105,8 +108,6 @@ x_data = rand(Normal(mu_true, sqrt(tau2_true)), N)
 # Store the generated data in the HierNormal structure
 hn = HierNormal((x=x_data,))
 
-##
-
 samples = sample(
     hn,
     Gibbs(
@@ -115,43 +116,46 @@ samples = sample(
             (:tau2,) => PriorMH(product_distribution([InverseGamma(1, 1)])),
         ),
     ),
-    100000;
+    200000;
     initial_params=(mu=[0.0], tau2=[1.0]),
 )
 
-mu_samples = [sample.values.mu for sample in samples][20001:end]
-tau2_samples = [sample.values.tau2 for sample in samples][20001:end]
-
-mean(mu_samples)
-mean(tau2_samples)
-
-##
-
-# this is too difficult of a problem
-
-gmm = GMM((; x=x))
-
-samples = sample(
-    gmm,
-    Gibbs(
-        OrderedDict(
-            (:z,) => PriorMH(product_distribution([Categorical([0.3, 0.7]) for _ in 1:60])),
-            (:w,) => PriorMH(Dirichlet(2, 1.0)),
-            (:μ,) => RWMH(1),
-        ),
-    ),
-    100000;
-    initial_params=(z=rand(Categorical([0.3, 0.7]), 60), μ=[-3.5, 0.5], w=[0.3, 0.7]),
-);
+mu_samples = [sample.values.mu for sample in samples][40001:end]
+tau2_samples = [sample.values.tau2 for sample in samples][40001:end]
 
-z_samples = [sample.values.z for sample in samples][20001:end]
-μ_samples = [sample.values.μ for sample in samples][20001:end]
-w_samples = [sample.values.w for sample in samples][20001:end];
+mu_mean = mean(mu_samples)[1]
+tau2_mean = mean(tau2_samples)[1]
 
-# thin these samples
-z_samples = z_samples[1:100:end]
-μ_samples = μ_samples[1:100:end]
-w_samples = w_samples[1:100:end];
+@testset "hierarchical normal with gibbs" begin
+    @test mu_mean ≈ mu_true atol = 0.1
+    @test tau2_mean ≈ tau2_true atol = 0.3
+end
 
-mean(μ_samples)
-mean(w_samples)
+## test with gmm -- too hard, doesn't converge
+
+# gmm = GMM((; x=x))
+
+# samples = sample(
+#     gmm,
+#     Gibbs(
+#         OrderedDict(
+#             (:z,) => PriorMH(product_distribution([Categorical([0.3, 0.7]) for _ in 1:60])),
+#             (:w,) => PriorMH(Dirichlet(2, 1.0)),
+#             (:μ,) => RWMH(1),
+#         ),
+#     ),
+#     100000;
+#     initial_params=(z=rand(Categorical([0.3, 0.7]), 60), μ=[-3.5, 0.5], w=[0.3, 0.7]),
+# );
+
+# z_samples = [sample.values.z for sample in samples][20001:end]
+# μ_samples = [sample.values.μ for sample in samples][20001:end]
+# w_samples = [sample.values.w for sample in samples][20001:end];
+
+# # thin these samples
+# z_samples = z_samples[1:100:end]
+# μ_samples = μ_samples[1:100:end]
+# w_samples = w_samples[1:100:end];
+
+# mean(μ_samples)
+# mean(w_samples)

test/gibbs_example/gibbs.jl

@@ -1,5 +1,3 @@
-using LogDensityProblems
-
 abstract type AbstractGMM end
 
 struct GMM <: AbstractGMM
@@ -81,65 +79,65 @@ function unflatten(vec::AbstractVector, group::Tuple)
 end
 
 function recompute_logprob!!(gmm::ConditionedGMM, vals, state)
-    return setlogp!!(state, LogDensityProblems.logdensity(gmm, vals))
+    return set_logp!!(state, LogDensityProblems.logdensity(gmm, vals))
 end
 
 ## test using Turing
 
-# data generation
+# # data generation
 
-using FillArrays
+# using FillArrays
 
-w = [0.5, 0.5]
-μ = [-3.5, 0.5]
-mixturemodel = Distributions.MixtureModel([MvNormal(Fill(μₖ, 2), I) for μₖ in μ], w)
+# w = [0.5, 0.5]
+# μ = [-3.5, 0.5]
+# mixturemodel = Distributions.MixtureModel([MvNormal(Fill(μₖ, 2), I) for μₖ in μ], w)
 
-N = 60
-x = rand(mixturemodel, N);
+# N = 60
+# x = rand(mixturemodel, N);
 
-# Turing model from https://turinglang.org/docs/tutorials/01-gaussian-mixture-model/
-using Turing
+# # Turing model from https://turinglang.org/docs/tutorials/01-gaussian-mixture-model/
+# using Turing
 
-@model function gaussian_mixture_model(x)
-    # Draw the parameters for each of the K=2 clusters from a standard normal distribution.
-    K = 2
-    μ ~ MvNormal(Zeros(K), I)
+# @model function gaussian_mixture_model(x)
+#     # Draw the parameters for each of the K=2 clusters from a standard normal distribution.
+#     K = 2
+#     μ ~ MvNormal(Zeros(K), I)
 
-    # Draw the weights for the K clusters from a Dirichlet distribution with parameters αₖ = 1.
-    w ~ Dirichlet(K, 1.0)
-    # Alternatively, one could use a fixed set of weights.
-    # w = fill(1/K, K)
+#     # Draw the weights for the K clusters from a Dirichlet distribution with parameters αₖ = 1.
+#     w ~ Dirichlet(K, 1.0)
+#     # Alternatively, one could use a fixed set of weights.
+#     # w = fill(1/K, K)
 
-    # Construct categorical distribution of assignments.
-    distribution_assignments = Categorical(w)
+#     # Construct categorical distribution of assignments.
+#     distribution_assignments = Categorical(w)
 
-    # Construct multivariate normal distributions of each cluster.
-    D, N = size(x)
-    distribution_clusters = [MvNormal(Fill(μₖ, D), I) for μₖ in μ]
+#     # Construct multivariate normal distributions of each cluster.
+#     D, N = size(x)
+#     distribution_clusters = [MvNormal(Fill(μₖ, D), I) for μₖ in μ]
 
-    # Draw assignments for each datum and generate it from the multivariate normal distribution.
-    k = Vector{Int}(undef, N)
-    for i in 1:N
-        k[i] ~ distribution_assignments
-        x[:, i] ~ distribution_clusters[k[i]]
-    end
+#     # Draw assignments for each datum and generate it from the multivariate normal distribution.
+#     k = Vector{Int}(undef, N)
+#     for i in 1:N
+#         k[i] ~ distribution_assignments
+#         x[:, i] ~ distribution_clusters[k[i]]
+#     end
 
-    return μ, w, k, __varinfo__
-end
+#     return μ, w, k, __varinfo__
+# end
 
-model = gaussian_mixture_model(x);
+# model = gaussian_mixture_model(x);
 
-using Test
-# full model
-μ, w, k, vi = model()
-@test log_joint(; μ=μ, w=w, z=k, x=x) ≈ DynamicPPL.getlogp(vi)
+# using Test
+# # full model
+# μ, w, k, vi = model()
+# @test log_joint(; μ=μ, w=w, z=k, x=x) ≈ DynamicPPL.getlogp(vi)
 
-gmm = GMM((; x=x))
+# gmm = GMM((; x=x))
 
-# cond model on μ, w
-μ, w, k, vi = (DynamicPPL.condition(model, (μ=μ, w=w)))()
-@test _logdensity(condition(gmm, (; μ=μ, w=w)), (; z=k)) ≈ DynamicPPL.getlogp(vi)
+# # cond model on μ, w
+# μ, w, k, vi = (DynamicPPL.condition(model, (μ=μ, w=w)))()
+# @test _logdensity(condition(gmm, (; μ=μ, w=w)), (; z=k)) ≈ DynamicPPL.getlogp(vi)
 
-# cond model on z
-μ, w, k, vi = (DynamicPPL.condition(model, (z = k)))()
-@test _logdensity(condition(gmm, (; z=k)), (; μ=μ, w=w)) ≈ DynamicPPL.getlogp(vi)
+# # cond model on z
+# μ, w, k, vi = (DynamicPPL.condition(model, (z = k)))()
+# @test _logdensity(condition(gmm, (; z=k)), (; μ=μ, w=w)) ≈ DynamicPPL.getlogp(vi)