Skip to content

Commit

Permalink
Apply suggestions from code review
Browse files Browse the repository at this point in the history
Co-authored-by: Xianda Sun <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
  • Loading branch information
3 people authored Oct 10, 2024
1 parent d9f8585 commit 600d36c
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 64 deletions.
57 changes: 14 additions & 43 deletions docs/src/api.md
Original file line number Diff line number Diff line change
Expand Up @@ -112,8 +112,8 @@ AbstractMCMC.chainsstack

To make it a bit easier to interact with some arbitrary sampler state, we encourage implementations of `AbstractSampler` to implement the following methods:
```@docs
AbstractMCMC.realize
AbstractMCMC.realize!!
AbstractMCMC.getparams
AbstractMCMC.setparams!!
```
and optionally
```@docs
Expand All @@ -125,7 +125,7 @@ These methods can also be useful for implementing samplers which wraps some inne

In a `MixtureSampler` we need two things:
- `components`: collection of samplers.
- `weights`: collection of weights representing the probability of chosing the corresponding sampler.
- `weights`: collection of weights representing the probability of choosing the corresponding sampler.

```julia
struct MixtureSampler{W,C} <: AbstractMCMC.AbstractSampler
Expand All @@ -136,7 +136,6 @@ end

To implement the state, we need to keep track of a couple of things:
- `index`: the index of the sampler used in this `step`.
- `transition`: the transition resulting from this `step`.
- `states`: the current states of _all_ the components.
Two aspects of this might seem a bit strange:
1. We need to keep track of the states of _all_ components rather than just the state for the sampler we used previously.
Expand All @@ -146,11 +145,9 @@ The reason for (1) is that lots of samplers keep track of more than just the pre

For (2) the reason is similar: some samplers might keep track of the variables _in the state_ differently, e.g. you might have a sampler which is _independent_ of the current realizations and the state is simply `nothing`.

Hence, we need the `transition`, which should always contain the realizations, to make sure we can resume from the same point in the space in the next `step`.
```julia
struct MixtureState{T,S}
struct MixtureState{S}
index::Int
transition::T
states::S
end
```
Expand All @@ -162,15 +159,16 @@ X_t &\sim \mathcal{K}_i(\cdot \mid X_{t - 1})
\end{aligned}
```
where ``\mathcal{K}_i`` denotes the i-th kernel/sampler, and ``w_i`` denotes the weight/probability of choosing the i-th sampler.
[`AbstractMCMC.updatestate!!`](@ref) comes into play in defining/computing ``\mathcal{K}_i(\cdot \mid X_{t - 1})`` since ``X_{t - 1}`` could be coming from a different sampler.
[`AbstractMCMC.getparams`](@ref) and [`AbstractMCMC.setparams!!`](@ref) comes into play in defining/computing ``\mathcal{K}_i(\cdot \mid X_{t - 1})`` since ``X_{t - 1}`` could be coming from a different sampler.

If we let `state` be the current `MixtureState`, `i` the current component, and `i_prev` is the previous component we sampled from, then this translates into the following piece of code:

```julia
# Update the corresponding state, i.e. `state.states[i]`, using
# the state and transition from the previous iteration.
state_current = AbstractMCMC.updatestate!!(
state.states[i], state.states[i_prev], state.transition
state_current = AbstractMCMC.setparams!!(
state.states[i],
AbstractMCMC.getparams(state.states[i_prev]),
)

# Take a `step` for this sampler using the updated state.
Expand All @@ -191,8 +189,9 @@ function AbstractMCMC.step(rng, model::AbstractMCMC.AbstractModel, sampler::Mixt
# Update the corresponding state, i.e. `state.states[i]`, using
# the state and transition from the previous iteration.
i_prev = state.index
state_current = AbstractMCMC.updatestate!!(
model, state.states[i], state.states[i_prev], state.transition
state_current = AbstractMCMC.setparams!!(
state.states[i],
AbstractMCMC.getparams(state.states[i_prev]),
)

# Take a `step` for this sampler using the updated state.
Expand All @@ -217,7 +216,7 @@ function AbstractMCMC.step(rng, model::AbstractMCMC.AbstractModel, sampler::Mixt
end

# Create the new `MixtureState`.
state_new = MixtureState(i, transition, states_new)
state_new = MixtureState(i, states_new)

return transition, state_new
end
Expand All @@ -239,20 +238,14 @@ function AbstractMCMC.step(rng, model::AbstractMCMC.AbstractModel, sampler::Mixt
# Extract states.
states = map(last, transitions_and_states)
# Create new `MixtureState`.
state = MixtureState(i, transition, states)
state = MixtureState(i, states)

return transition, state
end
```

Suppose we then wanted to use this with some of the packages which implements AbstractMCMC.jl's interface, e.g. [`AdvancedMH.jl`](https://github.com/TuringLang/AdvancedMH.jl), then we'd simply have to implement `realize` and `realize!!`:
Suppose we then wanted to use this with some of the packages which implements AbstractMCMC.jl's interface, e.g. [`AdvancedMH.jl`](https://github.com/TuringLang/AdvancedMH.jl), then we'd simply have to implement `getparams` and `setparams!!`.

```julia
function AbstractMCMC.updatestate!!(model, ::AdvancedMH.Transition, state_prev::AdvancedMH.Transition)
# Let's `deepcopy` just to be certain.
return deepcopy(state_prev)
end
```

To use `MixtureSampler` with two samplers `sampler1` and `sampler2` from `AdvancedMH.jl` as components, we'd simply do

Expand All @@ -263,25 +256,3 @@ while ...
transition, state = AbstractMCMC.step(rng, model, sampler, state)
end
```

As a final note, there is one potential issue we haven't really addressed in the above implementation: a lot of samplers have their own implementations of `AbstractMCMC.AbstractModel` which means that we would also have to ensure that all the different samplers we are using would be compatible with the same model. A very easy way to fix this would be to just add a struct called `ManyModels` supporting `getindex`, e.g. `models[i]` would return the i-th `model`:

```julia
struct ManyModels{M} <: AbstractMCMC.AbstractModel
models::M
end

Base.getindex(model::ManyModels, I...) = model.models[I...]
```

Then the above `step` would just extract the `model` corresponding to the current sampler:

```julia
# Take a `step` for this sampler using the updated state.
transition, state_current = AbstractMCMC.step(
rng, model[i], sampler_current, state_current;
kwargs...
)
```

This issue should eventually disappear as the community moves towards a unified approach to implement `AbstractMCMC.AbstractModel`.
32 changes: 11 additions & 21 deletions src/AbstractMCMC.jl
Original file line number Diff line number Diff line change
Expand Up @@ -80,35 +80,25 @@ The `MCMCSerial` algorithm allows users to sample serially, with no thread or pr
struct MCMCSerial <: AbstractMCMCEnsemble end

"""
updatestate!!(model, state, transition_prev[, state_prev])
getparams(state[; kwargs...])
Return new instance of `state` using information from `model`, `transition_prev` and, optionally, `state_prev`.
Defaults to `realize!!(state, realize(transition_prev))`.
Retrieve the values of parameters from the sampler's `state` as a `Vector{<:Real}`.
"""
function updatestate!!(model, state, transition_prev, state_prev)
return updatestate!!(state, transition_prev)
end
updatestate!!(model, state, transition) = realize!!(state, realize(transition))
function getparams end

"""
realize!!(state, realization)
Update the realization of the `state` with `realization` and return it.
setparams!!(state, params)
If `state` can be updated in-place, it is expected that this function returns `state` with updated
realize. Otherwise a new `state` object with the new `realization` is returned.
"""
function realize!! end
Set the values of parameters in the sampler's `state` from a `Vector{<:Real}`.
"""
realize(transition)
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.
Return the realization of the random variables present in `transition`.
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
and the vector representation of the parameter values.
"""
function realize end


function setparams!! end
include("samplingstats.jl")
include("logging.jl")
include("interface.jl")
Expand Down

0 comments on commit 600d36c

Please sign in to comment.