Add extensions with weak dependencies (#81)

* Add extensions with weak dependencies

* Fix test tolerance

* Do not overwrite method

* Fixes

* Fix typo

* Fix error hint

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedMH"
 uuid = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
-version = "0.7.2"
+version = "0.7.3"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -9,11 +9,26 @@ LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 
+[weakdeps]
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+
+[extensions]
+AdvancedMHForwardDiffExt = ["DiffResults", "ForwardDiff"]
+AdvancedMHMCMCChainsExt = "MCMCChains"
+AdvancedMHStructArraysExt = "StructArrays"
+
 [compat]
 AbstractMCMC = "4"
+DiffResults = "1"
 Distributions = "0.20, 0.21, 0.22, 0.23, 0.24, 0.25"
+ForwardDiff = "0.10"
 LogDensityProblems = "2"
+MCMCChains = "5, 6"
 Requires = "1"
+StructArrays = "0.6"
 julia = "1.6"
 
 [extras]

ext/AdvancedMHForwardDiffExt.jl

@@ -0,0 +1,20 @@
+module AdvancedMHForwardDiffExt
+
+if isdefined(Base, :get_extension)
+    import AdvancedMH
+    import DiffResults
+    import ForwardDiff
+else
+    import ..AdvancedMH
+    import ..DiffResults
+    import ..ForwardDiff
+end
+
+function AdvancedMH.logdensity_and_gradient(model::AdvancedMH.DensityModel, params)
+    res = DiffResults.GradientResult(params)
+    ForwardDiff.gradient!(res, model.logdensity, params)
+    return DiffResults.value(res), DiffResults.gradient(res)
+end
+
+
+end # module

src/mcmcchains-connect.jl renamed to ext/AdvancedMHMCMCChainsExt.jl

@@ -1,4 +1,12 @@
-import .MCMCChains: Chains
+module AdvancedMHMCMCChainsExt
+
+if isdefined(Base, :get_extension)
+    using AdvancedMH: AbstractMCMC, AbstractTransition, DensityModelOrLogDensityModel, Ensemble, MHSampler, Transition
+    using MCMCChains: Chains
+else
+    using ..AdvancedMH: AbstractMCMC, AbstractTransition, DensityModelOrLogDensityModel, Ensemble, MHSampler, Transition
+    using ..MCMCChains: Chains
+end
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
@@ -111,3 +119,5 @@ function AbstractMCMC.bundle_samples(
         start=discard_initial + 1, thin=thinning,
     )
 end
+
+end # module

src/structarray-connect.jl renamed to ext/AdvancedMHStructArraysExt.jl

@@ -1,4 +1,12 @@
-import .StructArrays: StructArray
+module AdvancedMHStructArraysExt
+
+if isdefined(Base, :get_extension)
+    using AdvancedMH: AbstractMCMC, AbstractTransition, DensityModelOrLogDensityModel, MHSampler
+    using StructArrays: StructArray
+else
+    using ..AdvancedMH: AbstractMCMC, AbstractTransition, DensityModelOrLogDensityModel, MHSampler
+    using ..StructArrays: StructArray
+end
 
 # A basic chains constructor that works with the Transition struct we defined.
 function AbstractMCMC.bundle_samples(
@@ -17,3 +25,6 @@ function AbstractMCMC.bundle_samples(
 end
 
 AbstractMCMC.chainscat(c::StructArray, cs::StructArray...) = vcat(c, cs...)
+
+
+end # module

src/AdvancedMH.jl

@@ -3,7 +3,6 @@ module AdvancedMH
 # Import the relevant libraries.
 using AbstractMCMC
 using Distributions
-using Requires
 
 using LogDensityProblems: LogDensityProblems
 
@@ -117,11 +116,23 @@ function AbstractMCMC.bundle_samples(
     return nts
 end
 
+if !isdefined(Base, :get_extension)
+    using Requires
+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("forwarddiff.jl")
+    # Better error message if users forget to load ForwardDiff
+    Base.Experimental.register_error_hint(MethodError) do io, exc, arg_types, kwargs
+        if exc.f === logdensity_and_gradient && length(arg_types) == 2 && first(arg_types) <: DensityModel && isempty(kwargs)
+            print(io, "\\nDid you forget to load ForwardDiff?")
+        end
+    end    
+    @static if !isdefined(Base, :get_extension)
+        @require MCMCChains="c7f686f2-ff18-58e9-bc7b-31028e88f75d" include("../ext/AdvancedMHMCMCChainsExt.jl")
+        @require StructArrays="09ab397b-f2b6-538f-b94a-2f83cf4a842a" include("../ext/AdvancedMHStructArraysExt.jl")
+        @require DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5" begin
+            @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" include("../ext/AdvancedMHForwardDiffExt.jl")
+        end
     end
 end
 

src/MALA.jl

@@ -78,19 +78,12 @@ function AbstractMCMC.step(
 end
 
 """
-    logdensity_and_gradient(model::DensityModel, params)
+    logdensity_and_gradient(model::AdvancedMH.DensityModelOrLogDensityModel, params)
 
 Return the value and gradient of the log density of the parameters `params` for the `model`.
 """
-function logdensity_and_gradient(model::DensityModel, params)
-    error("no implementation exist for `DensityModel`; try importing ForwardDiff.jl")
-end
-
-"""
-    logdensity_and_gradient(model::AbstractMCMC.LogDensityModel, params)
+logdensity_and_gradient(::DensityModelOrLogDensityModel, ::Any)
 
-Return the value and gradient of the log density of the parameters `params` for the `model`.
-"""
 function logdensity_and_gradient(model::AbstractMCMC.LogDensityModel, params)
     check_capabilities(model)
     return LogDensityProblems.logdensity_and_gradient(model.logdensity, params)

src/forwarddiff.jl

@@ -98,7 +98,7 @@ include("util.jl")
         )
         @test chain1b isa Chains
         @test range(chain1b) == range(26; step=4, length=10_000)
-        @test mean(chain1b["μ"]) ≈ 0.0 atol=0.1
+        @test mean(chain1b["μ"]) ≈ 0.0 atol=0.15
         @test mean(chain1b["σ"]) ≈ 1.0 atol=0.1
 
         # NamedTuple of parameters