add hierarchical normal problem

Author: sunxd3

gibbs_example/gibbs.jl

@@ -92,6 +92,41 @@ end
 
 ## tests
 
+# generate data
+N = 100  # Number of data points
+mu_true = 0.5  # True mean
+tau2_true = 2.0  # True variance
+
+# Generate data based on true parameters
+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(
+        OrderedDict(
+            (:mu,) => RWMH(1),
+            (:tau2,) => PriorMH(product_distribution([InverseGamma(1, 1)])),
+        ),
+    ),
+    100000;
+    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(
@@ -104,12 +139,17 @@ samples = sample(
         ),
     ),
     100000;
-    initial_params=(z=rand(Categorical([0.3, 0.7]), 60), μ=[0.0, 1.0], w=[0.3, 0.7]),
+    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]
+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)

gibbs_example/hier_normal.jl

@@ -0,0 +1,73 @@
+using LogDensityProblems
+
+abstract type AbstractHierNormal end
+
+struct HierNormal <: AbstractHierNormal
+    data::NamedTuple
+end
+
+struct ConditionedHierNormal{conditioned_vars} <: AbstractHierNormal
+    data::NamedTuple
+    conditioned_values::NamedTuple{conditioned_vars}
+end
+
+function log_joint(; mu, tau2, x)
+    # mu is the mean
+    # tau2 is the variance
+    # x is data
+
+    # μ ~ Normal(0, 1)
+    # τ² ~ InverseGamma(1, 1)
+    # xᵢ ~ Normal(μ, √τ²)
+
+    logp = 0.0
+    mu = only(mu)
+    tau2 = only(tau2)
+
+    mu_prior = Normal(0, 1)
+    logp += logpdf(mu_prior, mu)
+
+    tau2_prior = InverseGamma(1, 1)
+    logp += logpdf(tau2_prior, tau2)
+
+    obs_prior = Normal(mu, sqrt(tau2))
+    logp += sum(logpdf(obs_prior, xi) for xi in x)
+
+    return logp
+end
+
+function condition(hn::HierNormal, conditioned_values::NamedTuple)
+    return ConditionedHierNormal(hn.data, conditioned_values)
+end
+
+function LogDensityProblems.logdensity(
+    hn::ConditionedHierNormal{names}, params::AbstractVector
+) where {names}
+    if Set(names) == Set([:mu]) # conditioned on mu, so params are tau2
+        return log_joint(; mu=hn.conditioned_values.mu, tau2=params, x=hn.data.x)
+    elseif Set(names) == Set([:tau2]) # conditioned on tau2, so params are mu
+        return log_joint(; mu=params, tau2=hn.conditioned_values.tau2, x=hn.data.x)
+    else
+        error("Unsupported conditioning configuration.")
+    end
+end
+
+function LogDensityProblems.capabilities(::HierNormal)
+    return LogDensityProblems.LogDensityOrder{0}()
+end
+
+function LogDensityProblems.capabilities(::ConditionedHierNormal)
+    return LogDensityProblems.LogDensityOrder{0}()
+end
+
+function flatten(nt::NamedTuple)
+    return only(values(nt))
+end
+
+function unflatten(vec::AbstractVector, group::Tuple)
+    return NamedTuple((only(group) => vec,))
+end
+
+function recompute_logprob!!(hn::ConditionedHierNormal, vals, state)
+    return setlogp!!(state, LogDensityProblems.logdensity(hn, vals))
+end