Make Gibbs constructor more flexible

Author: mhauru

src/mcmc/gibbs.jl

@@ -320,21 +320,30 @@ struct Gibbs{V,A} <: InferenceAlgorithm
     end
 end
 
+to_varname(vn::VarName) = vn
+to_varname(s::Symbol) = VarName{s}()
+# Any other value is assumed to be an iterable.
+to_varname(t) = map(to_varname, collect(t))
+
 # NamedTuple
 Gibbs(; algs...) = Gibbs(NamedTuple(algs))
 function Gibbs(algs::NamedTuple)
     return Gibbs(
-        map(s -> VarName{s}(), keys(algs)),
-        map(wrap_algorithm_maybe ∘ drop_space, values(algs)),
+        map(to_varname, keys(algs)), map(wrap_algorithm_maybe ∘ drop_space, values(algs))
     )
 end
 
 # AbstractDict
 function Gibbs(algs::AbstractDict)
-    return Gibbs(collect(keys(algs)), map(wrap_algorithm_maybe ∘ drop_space, values(algs)))
+    return Gibbs(
+        map(to_varname, collect(keys(algs))),
+        map(wrap_algorithm_maybe ∘ drop_space, values(algs)),
+    )
 end
 function Gibbs(algs::Pair...)
-    return Gibbs(map(first, algs), map(wrap_algorithm_maybe ∘ drop_space, map(last, algs)))
+    return Gibbs(
+        map(to_varname ∘ first, algs), map(wrap_algorithm_maybe ∘ drop_space ∘ last, algs)
+    )
 end
 
 # The below two constructors only provide backwards compatibility with the constructor of
@@ -383,6 +392,7 @@ end
 _maybevec(x) = vec(x)  # assume it's iterable
 _maybevec(x::Tuple) = [x...]
 _maybevec(x::VarName) = [x]
+_maybevec(x::Symbol) = [x]
 
 varinfo(state::GibbsState) = state.vi
 

test/mcmc/gibbs.jl

@@ -284,15 +284,15 @@ end
         N = 10
         # Two variables being sampled by one sampler.
         s1 = Gibbs((@varname(s), @varname(m)) => HMC(0.1, 5; adtype=adbackend))
-        s2 = Gibbs((@varname(s), @varname(m)) => PG(10))
+        s2 = Gibbs((@varname(s), :m) => PG(10))
         # One variable per sampler, using the keyword arg interface.
         s3 = Gibbs((; s=PG(3), m=HMC(0.4, 8; adtype=adbackend)))
         # As above but using a Dict of VarNames.
         s4 = Gibbs(Dict(@varname(s) => PG(3), @varname(m) => HMC(0.4, 8; adtype=adbackend)))
-        # As above but different samplers.
+        # As above but different samplers and using kwargs.
         s5 = Gibbs(; s=CSMC(3), m=HMC(0.4, 8; adtype=adbackend))
         s6 = Gibbs(; s=HMC(0.1, 5; adtype=adbackend), m=ESS())
-        s7 = Gibbs((@varname(s), @varname(m)) => PG(10))
+        s7 = Gibbs(Dict((:s, @varname(m)) => PG(10)))
         # Multiple instnaces of the same sampler. This implements running, in this case,
         # 3 steps of HMC on m and 2 steps of PG on m in every iteration of Gibbs.
         s8 = begin