Merge pull request #29 from briandepasquale/master

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedMH"
 uuid = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
-version = "0.5.4"
+version = "0.5.5"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -18,6 +18,8 @@ julia = "1"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 
 [targets]
-test = ["StructArrays", "MCMCChains", "Test"]
+test = ["StructArrays", "MCMCChains", "Test", "ForwardDiff", "DiffResults"]

README.md

@@ -118,3 +118,37 @@ in parallel for free:
 # Sample 4 chains from the posterior.
 chain = psample(model, RWMH(init_params), 100000, 4; param_names=["μ","σ"], chain_type=Chains)
 ```
+
+## Metropolis-adjusted Langevin algorithm (MALA)
+
+AdvancedMH.jl also offers an implementation of [MALA](https://en.wikipedia.org/wiki/Metropolis-adjusted_Langevin_algorithm) if the `ForwardDiff` and `DiffResults` packages are available. 
+
+A `MALA` sampler can be constructed by `MALA(proposal)` where `proposal` is a function that
+takes the gradient computed at the current sample. It is required to specify an initial sample `init_params` when calling `sample`.
+
+```julia
+# Import the package.
+using AdvancedMH
+using Distributions
+using MCMCChains
+using DiffResults
+using ForwardDiff
+
+# Generate a set of data from the posterior we want to estimate.
+data = rand(Normal(0, 1), 30)
+
+# Define the components of a basic model.
+insupport(θ) = θ[2] >= 0
+dist(θ) = Normal(θ[1], θ[2])
+density(θ) = insupport(θ) ? sum(logpdf.(dist(θ), data)) : -Inf
+
+# Construct a DensityModel.
+model = DensityModel(density)
+
+# Set up the sampler with a multivariate Gaussian proposal.
+sigma = 1e-1
+spl = MALA(x -> MvNormal((sigma^2 / 2) .* x, sigma))
+
+# Sample from the posterior.
+chain = sample(model, spl, 100000; init_params=ones(2), chain_type=StructArray, param_names=["μ", "σ"])
+```

src/AdvancedMH.jl

@@ -9,14 +9,17 @@ import Random
 
 # Exports
 export MetropolisHastings, DensityModel, RWMH, StaticMH, StaticProposal, 
-    RandomWalkProposal, Ensemble, StretchProposal
+    RandomWalkProposal, Ensemble, StretchProposal, MALA
 
 # Reexports
 export sample, MCMCThreads, MCMCDistributed
 
 # Abstract type for MH-style samplers. Needs better name? 
 abstract type MHSampler <: AbstractMCMC.AbstractSampler end
 
+# Abstract type for MH-style transitions.
+abstract type AbstractTransition end
+
 # Define a model type. Stores the log density function and the data to 
 # evaluate the log density on.
 """
@@ -37,7 +40,7 @@ end
 
 # Create a very basic Transition type, only stores the
 # parameter draws and the log probability of the draw.
-struct Transition{T<:Union{Vector, Real, NamedTuple}, L<:Real}
+struct Transition{T<:Union{Vector, Real, NamedTuple}, L<:Real} <: AbstractTransition
     params :: T
     lp :: L
 end
@@ -51,7 +54,7 @@ logdensity(model::DensityModel, t::Transition) = t.lp
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
-    ts::Vector{<:Transition},
+    ts::Vector{<:AbstractTransition},
     model::DensityModel,
     sampler::MHSampler,
     state,
@@ -98,6 +101,9 @@ end
 function __init__()
     @require MCMCChains="c7f686f2-ff18-58e9-bc7b-31028e88f75d" include("mcmcchains-connect.jl")
     @require StructArrays="09ab397b-f2b6-538f-b94a-2f83cf4a842a" include("structarray-connect.jl")
+    @require DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" begin
+        @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" include("MALA.jl")
+    end
 end
 
 # Include inference methods.

src/MALA.jl

@@ -0,0 +1,61 @@
+using .ForwardDiff: gradient!
+using .DiffResults: GradientResult, value, gradient
+
+struct MALA{D} <: MHSampler
+    proposal::D
+end
+
+
+# Create a RandomWalkProposal if we weren't given one already.
+MALA(d) = MALA(RandomWalkProposal(d))
+
+# If we were given a RandomWalkProposal, just use that instead.
+MALA(d::RandomWalkProposal) = MALA{typeof(d)}(d)
+
+
+struct GradientTransition{T<:Union{Vector, Real, NamedTuple}, L<:Real, G<:Union{Vector, Real, NamedTuple}} <: AbstractTransition
+    params::T
+    lp::L
+    gradient::G
+end
+
+transition(::MALA, model, params) = GradientTransition(model, params)
+
+# Store the new draw, its log density and its gradient
+GradientTransition(model::DensityModel, params) = GradientTransition(params, logdensity_and_gradient(model, params)...)
+
+propose(rng::Random.AbstractRNG, ::MALA, model) = error("please specify initial parameters")
+
+function propose(
+    rng::Random.AbstractRNG,
+    spl::MALA{<:Proposal},
+    model::DensityModel,
+    params_prev::GradientTransition
+)
+    proposal = propose(rng, spl.proposal(params_prev.gradient), model, params_prev.params)
+    return GradientTransition(model, proposal)
+end
+
+
+function q(
+    spl::MALA{<:Proposal},
+    t::GradientTransition,
+    t_cond::GradientTransition
+)
+    return q(spl.proposal(-t_cond.gradient), t.params, t_cond.params)
+end
+
+
+"""
+    logdensity_and_gradient(model::DensityModel, params)
+
+Efficiently returns the value and gradient of the model
+"""
+function logdensity_and_gradient(model::DensityModel, params)
+    res = GradientResult(params)
+    gradient!(res, model.logdensity, params)
+    return (value(res), gradient(res))
+end
+
+
+logdensity(model::DensityModel, t::GradientTransition) = t.lp

src/mh-core.jl

@@ -169,20 +169,23 @@ function q(
     end
 end
 
+transition(sampler, model, params) = transition(model, params)
+transition(model, params) = Transition(model, params)
+
 # Define the first sampling step.
 # Return a 2-tuple consisting of the initial sample and the initial state.
 # In this case they are identical.
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DensityModel,
-    spl::MetropolisHastings;
+    spl::MHSampler;
     init_params=nothing,
     kwargs...
 )
     if init_params === nothing
         transition = propose(rng, spl, model)
     else
-        transition = Transition(model, init_params)
+        transition = AdvancedMH.transition(spl, model, init_params)
     end
 
     return transition, transition
@@ -195,8 +198,8 @@ end
 function AbstractMCMC.step(
     rng::Random.AbstractRNG,
     model::DensityModel,
-    spl::MetropolisHastings,
-    params_prev::Transition;
+    spl::MHSampler,
+    params_prev::AbstractTransition;
     kwargs...
 )
     # Generate a new proposal.

test/runtests.jl

@@ -5,6 +5,8 @@ using MCMCChains
 
 using Random
 using Test
+using DiffResults
+using ForwardDiff
 
 @testset "AdvancedMH" begin
     # Set a seed
@@ -102,7 +104,20 @@ using Test
 
         @test chain1[1].params == val
     end
+    
+    @testset "MALA" begin
+        
+        # Set up the sampler.
+        sigma = 1e-1
+        spl1 = MALA(x -> MvNormal((sigma^2 / 2) .* x, sigma))
+
+        # Sample from the posterior with initial parameters.
+        chain1 = sample(model, spl1, 100000; init_params=ones(2), chain_type=StructArray, param_names=["μ", "σ"])
+
+        @test mean(chain1.μ) ≈ 0.0 atol=0.1
+        @test mean(chain1.σ) ≈ 1.0 atol=0.1 
+    end
 
     @testset "EMCEE" begin include("emcee.jl") end
+  
 end
-