Add a model argument to getparams and setparams!! functions (#150)

* add optional argument of logdensity_function

* version bump

* apply suggestions

* add some default implementations

* Apply suggestions from code review

Co-authored-by: Tor Erlend Fjelde <[email protected]>

* Update src/AbstractMCMC.jl

Co-authored-by: Tor Erlend Fjelde <[email protected]>

* Update Project.toml

* add some documentation improvement

* improve added doc

---------

Co-authored-by: Tor Erlend Fjelde <[email protected]>

Author: sunxd3

Project.toml

@@ -3,7 +3,7 @@ uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 keywords = ["markov chain monte carlo", "probabilistic programming"]
 license = "MIT"
 desc = "A lightweight interface for common MCMC methods."
-version = "5.5.0"
+version = "5.6.0"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"

docs/src/api.md

@@ -121,7 +121,13 @@ To make it a bit easier to interact with some arbitrary sampler state, we encour
 AbstractMCMC.getparams
 AbstractMCMC.setparams!!
 ```
-These methods can also be useful for implementing samplers which wraps some inner samplers, e.g. a mixture of samplers.
+`getparams` and `setparams!!` provide a generic interface for interacting with the parameters of a sampler's state, regardless of how that state is represented internally.
+
+This allows generic code to be written that works with any sampler implementing this interface. For example, a generic ensemble sampler could use `getparams` to extract the parameters from each of its component samplers' states, and `setparams!!` to initialize each component sampler with a different set of parameters.
+
+The optional `model` argument to these functions allows sampler implementations to customize their behavior based on the model being used. For example, some samplers may need to evaluate the log density at new parameter values when setting parameters, which requires access to the model. If access to `model` is not needed, the sampler only needs to implement the version without the `model` argument - the default implementations will then call those methods directly.
+
+These methods are particularly useful for implementing samplers which wrap some inner samplers, such as a mixture of samplers. In the next section, we will see how `getparams` and `setparams!!` can be used to implement a `MixtureSampler`.
 
 ### Example: `MixtureSampler`
 

src/AbstractMCMC.jl

@@ -81,26 +81,37 @@ The `MCMCSerial` algorithm allows users to sample serially, with no thread or pr
 struct MCMCSerial <: AbstractMCMCEnsemble end
 
 """
-    getparams(state[; kwargs...])
+    getparams([model::AbstractModel, ]state)
 
 Retrieve the values of parameters from the sampler's `state` as a `Vector{<:Real}`.
 """
 function getparams end
 
+function getparams(model::AbstractModel, state)
+    return getparams(state)
+end
+
 """
-    setparams!!(state, params)
+    setparams!!([model::AbstractModel, ]state, params)
 
 Set the values of parameters in the sampler's `state` from a `Vector{<:Real}`. 
 
 This function should follow the `BangBang` interface: mutate `state` in-place if possible and 
 return the mutated `state`. Otherwise, it should return a new `state` containing the updated parameters.
 
-Although not enforced, it should hold that `setparams!!(state, getparams(state)) == state`. In another
-word, the sampler should implement a consistent transformation between its internal representation
+Although not enforced, it should hold that `setparams!!(state, getparams(state)) == state`. In other
+words, the sampler should implement a consistent transformation between its internal representation
 and the vector representation of the parameter values.
+
+Sometimes, to maintain the consistency of the log density and parameter values, a `model`
+should be provided. This is useful for samplers that need to evaluate the log density at the new parameter values.
 """
 function setparams!! end
 
+function setparams!!(model::AbstractModel, state, params)
+    return setparams!!(state, params)
+end
+
 include("samplingstats.jl")
 include("logging.jl")
 include("interface.jl")