Add `state_from_transition`, `parameters` and `setparameters!!`

[torfjelde]

See #85.

Currently, almost every package implementing a AbstractSampler comes with their custom state and transition types. Since there currently are no "interface" or "guidelines" about what the types of state and transition should implement, it means that interaction across different samplers, e.g. composing two samplers from different packages, is annoying.

This PR introduces some methods which aims to alleviate this issue. In particular it sampler-implementations to also implement

• parameters(transition) for the transition-type which returns the parameters in the sample, where by "parameters" I mean the realizations of the random variables we're targeting present in the transition, and
• setparameters!! for the state-type which updates the current state of the Markov chain, thus likely (this might depend on the sampler) conditioning the sampler/kernel on the provided parameters.

Why parameters for transition but setparameters!! for state? A transition should (AFAIK) always contain a realization of the random variables we're sampling while the state might not, e.g. MH with a MvNormal(zeros(n), I) as proposal distribution (not dependent on current position/state/parameters). Of course there might be more properties one would like to use when converting from PackageA.Transition to PackageB.State, but this seems like the lowest common denominator that we can come up with. Maaaybe we should have the same for logdensity, i.e. logdensity(transition) and setlogdensity!!(state), but this might not always be available nor correspond directly to the same thing for both samplers. Might be worth including though. They could also obviously be implemented for both if the implementer so chooses.

But for the cases where you want to make use of more than just the parameters (and potentially logdensity), one can overload the state_from_transition directly (which is what should be used by packages making these "meta-samplers", e.g. composition), e.g. AdvancedHMC might want to implement state_from_transition(::AdvancedHMC.HMCState, ::AdvancedHMC.Transition) to not only update the parameters in the state but the entire PhasePoint from AdvnacedHMC.Transition, etc.

Thoughts?

[codecov[bot]]

Codecov Report

Attention: 3 lines in your changes are missing coverage. Please review.

Comparison is base (4dbcb3f) 97.37% compared to head (d9f8585) 96.42%. Report is 4 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master      #86      +/-   ##
==========================================
- Coverage   97.37%   96.42%   -0.95%     
==========================================
  Files           8        8              
  Lines         305      308       +3     
==========================================
  Hits          297      297              
- Misses          8       11       +3     

Files	Coverage Δ
src/AbstractMCMC.jl	57.14% <0.00%> (-42.86%)	:arrow_down:

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[devmotion]

parameters(transition) for the transition-type which returns the parameters in the sample, where by "parameters" I mean the realizations of the random variables we're targeting present in the transition,

Maybe call it samples instead?

[devmotion]

Maaaybe we should have the same for logdensity, i.e. logdensity(transition) and setlogdensity!!(state)

As you mention, the main problem is that it is not clear which logdensity different transitions and states refer to, so one needs some more fine-grained system e.g. some additional trait to distinguish between e.g. logdensity of prior or joint. I guess it's best to leave this for a separate PR and/or issue.

[torfjelde]

Maybe call it samples instead?

Though I agree with "parameters" not being a good choice, "samples" can be a bit confusing IMO given all the other sample-related functions we have. Would values make sense, given that in this context it's often used synonymously with realizations of random variables?

As you mention, the main problem is that it is not clear which logdensity different transitions and states refer to, so one needs some more fine-grained system e.g. some additional trait to distinguish between e.g. logdensity of prior or joint. I guess it's best to leave this for a separate PR and/or issue.

Happy with this :+1:

[torfjelde]

Do you have ideas for alternatives to the use of parameters @devmotion and @cpfiffer ? As I said in my previous comment, I think "samples" is kind of confusing too. Really what we want is "set the realizations of the random variables" but I think "realizations" is too unfamiliar for people.

I think I'm leaning towards values and setvalues!!.

[cpfiffer]

I dunno I feel like I wouldn't mind using realize or realize!!. Short and sweet, and I don't think it's that out of the way. If that doesn't seem right, I don't mind using values (though I will note that it's a base function override, which I don't care about overmuch).

[torfjelde]

I've added some docs + motivating example: https://turinglang.github.io/AbstractMCMC.jl/previews/PR86/api/#Interacting-with-states-of-samplers

(Btw, I love the fact that I can preview docs in the PRs; I'm guessing this is your doing @devmotion ? :heart: )

[torfjelde]

I think the only thing that's left is the discussion of parameters vs. realizations vs. values

[cpfiffer]

I prefer realizations but values is more Julian -- let's go with that.

[torfjelde]

On additional thing we might want to consider adding: logdensity and setlogdensity!!, or something along those lines. This is because:

1. Almost all samplers will have a cached value of the logdensity (up-to some constant) of their target in their state.
2. Some samplers will make use of said cached value, e.g. AdvancedMH uses this in its accept/reject computation.

Of course, this can be dealt with on a case-by-case basis by simply defining updatestate!! explicitly, but this will be unnecessary in most cases if we have something along the lines of logdensity and setlogdensity!!.

What do you say @devmotion @cpfiffer ? Add or not add, that is the question.

[torfjelde]

Also, @devmotion you mentioned tests. Do you have anything particular in mind?

I have a test for the MixtureSampler outlined in the docs using AdvancedMH.jl, but that seems a bit "too heavy". At the same time, doing something like:

struct TestState
    params
    logp
end

setvalues!!(s::TestState) = ...

etc. seems quite useless :shrug:

We could potentially take the test I have for the MixtureSampler and add a super-simple implementation of RWMH to the test-suite of AbstractMCMC? Not sure I'm too big a fan of this either though, as it'll be more like an example use-case rather than testing the functionality provided by updatestate!! etc.

[torfjelde]

On additional thing we might want to consider adding: logdensity and setlogdensity!!, or something along those lines. This is because:

1. Almost all samplers will have a cached value of the logdensity (up-to some constant) of their target in their state.
2. Some samplers will make use of said cached value, e.g. AdvancedMH uses this in its accept/reject computation.

Of course, this can be dealt with on a case-by-case basis by simply defining updatestate!! explicitly, but this will be unnecessary in most cases if we have something along the lines of logdensity and setlogdensity!!.

What do you say @devmotion @cpfiffer ? Add or not add, that is the question.

Yo boys @devmotion @cpfiffer! You got any thoughts on the above? AFAIK this is the only thing that's holding this PR back.

[cpfiffer]

Yo boys @devmotion @cpfiffer! You got any thoughts on the above? AFAIK this is the only thing that's holding this PR back.

Woops, sorry, missed this one. Full support from me on this one!

[devmotion]

Ah I also missed this question. I think it would probably be useful but I would leave it for a separate PR. There were quite many discussions about logdensity functions recently and there was a push for unifying them throughout the Julia ecosystem. DensityInterface.jl defines a common interface for log density functions and their evaluations now that is used by e.g. Distributions, AbstractPPL, and (AFAIK) in the future also MeasureTheory. Hence IMO we should be a bit careful and not introduce a different logdensity function here, now that we unified them, but also make sure that it corresponds to the DensityInterface assumptions (if we use its interface).

[torfjelde]

Yeah I was actually looking at the newly introduced logdensity stuff today, so I agree that maybe this should be done in a separate PR :+1:

[torfjelde]

I renamed values and setvalues!! to realize and realize!! as mentioned earlier.

Even if we're not completely happy with these names, we can instead just change these at a later stage.

[torfjelde]

Ah I also missed this question. I think it would probably be useful but I would leave it for a separate PR. There were quite many discussions about logdensity functions recently and there was a push for unifying them throughout the Julia ecosystem. DensityInterface.jl defines a common interface for log density functions and their evaluations now that is used by e.g. Distributions, AbstractPPL, and (AFAIK) in the future also MeasureTheory. Hence IMO we should be a bit careful and not introduce a different logdensity function here, now that we unified them, but also make sure that it corresponds to the DensityInterface assumptions (if we use its interface).

After looking into DensityInterface.jl I don't think there's anything that should be relevant for AbstractMCMC's functionality for setting and extracting log-probabilities from states/transitions. For the AbstractModel, etc., then it's def relevant though!

So maybe we should just introduce logprob or logtargetprob or maybe even just logtarget, and its corresponding set..!!? I'm realizing more and more that it's necessary to be able to transfer the log target-prob to the next state if we want to do something interesting. In fact, I think update!! state should optionally be able to take the AbstractModel that we're sampling from too so that it can re-evaluate the current state if need be.

EDIT: I added model as an argument to updatestate!! according to the last paragraph above.

EDIT 2: Maybe something like

"""
    target_logdensity(model, transition)

Return the log-density of the target `model` represented at `transition`.

Note that this does not necessarily compute the log-density of `model` at
`transition`, but may simply return a cached result from `transition`.
"""
function target_logdensity(model::AbstractModel, transition)
    return DensityInterface.logdensityof(model, realize(transition))
end

"""
    set_target_logdensity!!(state, logp)

Update the log-density in `state` by setting it to `logp`.

If `state` is mutable, this function mutates `state` directly and returns `state`.
Otherwise a new and updated instance of `state` is returned.
"""
function set_target_logdensity!! end