Make MH proposal generation more sane (#1557)

Co-authored-by: David Widmann <[email protected]>

Author: cpfiffer

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.15.11"
+version = "0.15.12"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"

src/inference/mh.jl

@@ -7,7 +7,10 @@ struct MH{space, P} <: InferenceAlgorithm
 end
 
 proposal(p::AdvancedMH.Proposal) = p
+proposal(f::Function) = AdvancedMH.StaticProposal(f)
+proposal(d::Distribution) = AdvancedMH.StaticProposal(d)
 proposal(cov::AbstractMatrix) = AdvancedMH.RandomWalkProposal(MvNormal(cov))
+proposal(x) = error("proposals of type ", typeof(x), " are not supported")
 
 """
     MH(space...)
@@ -162,14 +165,7 @@ function MH(space...)
             # Check to see whether it's a pair that specifies a kernel
             # or a specific proposal distribution.
             push!(prop_syms, s[1])
-
-            if s[2] isa AMH.Proposal
-                push!(props, s[2])
-            elseif s[2] isa Distribution
-                push!(props, AMH.StaticProposal(s[2]))
-            elseif s[2] isa Function
-                push!(props, AMH.StaticProposal(s[2]))
-            end
+            push!(props, proposal(s[2]))
         elseif length(space) == 1
             # If we hit this block, check to see if it's 
             # a run-of-the-mill proposal or covariance
@@ -178,11 +174,17 @@ function MH(space...)
 
             # Return early, we got a covariance matrix. 
             return MH{(), typeof(prop)}(prop)
+        else
+            # Try to convert it to a proposal anyways, 
+            # throw an error if not acceptable.
+            prop = proposal(s)
+            push!(props, prop)
         end
     end
 
     proposals = NamedTuple{tuple(prop_syms...)}(tuple(props...))
     syms = vcat(syms, prop_syms)
+
     return MH{tuple(syms...), typeof(proposals)}(proposals)
 end
 

test/inference/gibbs.jl

@@ -127,7 +127,7 @@
             end
         end
         model = imm(randn(100), 1.0);
-        sample(model, Gibbs(MH(10, :z), HMC(0.01, 4, :m)), 100);
+        sample(model, Gibbs(MH(:z), HMC(0.01, 4, :m)), 100);
         sample(model, Gibbs(PG(10, :z), HMC(0.01, 4, :m)), 100);
     end
 end

test/inference/mh.jl

@@ -106,6 +106,48 @@
         check_gdemo(chain2)
     end
 
+    @turing_testset "gibbs MH proposal matrix" begin
+        # https://github.com/TuringLang/Turing.jl/issues/1556
+
+        # generate data
+        x = rand(Normal(5, 10), 20)
+        y = rand(LogNormal(-3, 2), 20)
+        
+        # Turing model
+        @model function twomeans(x, y)
+            # Set Priors
+            μ ~ MvNormal(2, 3)
+            σ ~ filldist(Exponential(1), 2)
+        
+            # Distributions of supplied data
+            x .~ Normal(μ[1], σ[1])
+            y .~ LogNormal(μ[2], σ[2])
+        
+        end
+        mod = twomeans(x, y)
+        
+        # generate covariance matrix for RWMH
+        # with small-valued VC matrix to check if we only see very small steps
+        vc_μ = convert(Array, 1e-4*I(2))
+        vc_σ = convert(Array, 1e-4*I(2))
+        
+        chn = sample(mod,
+            Gibbs(
+                MH((:μ, vc_μ)),
+                MH((:σ, vc_σ)),
+                ), 3_000 # draws
+            )
+        
+            
+        chn2 = sample(mod, MH(), 3_000)
+
+        # Test that the small variance version is actually smaller.
+        v1 = var(diff(Array(chn["μ[1]"]), dims=1))
+        v2 = var(diff(Array(chn2["μ[1]"]), dims=1))
+
+        @test v1 < v2
+    end
+
     @turing_testset "vector of multivariate distributions" begin
         @model function test(k)
             T = Vector{Vector{Float64}}(undef, k)