Adopt AbstractMCMC.jl interface #259

torfjelde · 2021-04-06T14:06:15Z

Things to discuss

AFAIK, the way to customize the logging in AbstractMCMC.jl is to pass progress=false to the underlying AbstractMCMC.mcmcsample and then use the callback keyword argument to log the progress. So the question is: should we do this so as to preserve the current logging functionality?
To replicate the current summarization functionality (e.g. inform the user of average acceptance rates and EBFMI) as a post-sample step, we can overload StatsBase.sample and then perform this step after the call to AbstractMCMC.mcmcsample. Should we do this?

cpfiffer · 2021-04-06T15:19:49Z

xukai92 · 2021-04-06T20:39:41Z

src/trajectory.jl

@@ -211,7 +213,7 @@ nsteps(τ::Trajectory{TS, I, TC}) where {TS, I, TC<:FixedIntegrationTime} =
 ## Kernel interface
 ##

-struct HMCKernel{R, T<:Trajectory} <: AbstractMCMCKernel 
+struct HMCKernel{R, T<:Trajectory} <: AbstractMCMCKernel


Just to double check this is just the removal of a space?

Yep. But also, now AbstractMCMCKernel is AbstractMCMC.AbstractSampler (I'll remove this alias, and make it explicit).

Is there any side effect it might cause?

Don't think so. Anything particular in mind or just general question?

The only thing is that if we remove AbstractMCMCKernel completely, not even aliasing, e.g. CoupledHMC.jl won't work.

Maybe best approach is to just make AbstractMCMCKernel <: AbstractMCMC.AbstractSampler?

I think it's fine as long as CoupledHMC is fine (with sutff like mixture kernels still working).

xukai92 · 2021-04-06T20:41:49Z

src/sampler.jl

+    return AbstractMCMC.step(rng, model, spl, state; kwargs...)
+end
+
+function AbstractMCMC.step(


I think we should keep the old signature, which would construct HamiltonianModel and HMCState then call this function.

Hmm, maybe. We have already introduced breaking changes on master with the introduction of HMCKernel, etc., so it seems like it's as good a time as any to also transition completely over to AbstractMCMC.jl interface?

And it's not like there's anything in the current impl that isn't supported by the new implementation:)

There wasn't actually breaking changes to the outmost intefaces (e.g. those need to call sample).

Hmm, what do you mean? E.g. the example in the README doesn't work now since proposal isn't constructed in the same way (might also be other issues, but that's the immediate one that came to mind).

Sure, I'm cool with adding the old sample back and deprecating:) That seems like a clean approach 👍

Also it's a bit cumbersome to write sample(..., HamiltonianModel(hamiltonian), ..) instead of sample(..., hamiltonian, ..), no?

I'm also cool with either of the following:

Making AbstractHamiltonian <: AbstractMCMC.AbstractModel, which IMO seems sensible.

Additional overload that will automagically wrap hamiltonian in HamiltonianModel.

It seems that option 1 is more sensible, as we already making kernel sub-typed of samples from AbstractMCMC.

Yeaah I've realized I need to make quite a few changes. Because everything is mutated somehow, we essentially need to put everything in the sampler state.

Sooo the question is what we actually use for the model, etc. As of right now they would just be "fake" in the sense that they only contain the initial Hamiltonian and the initial AbstractHMCKernel 😕 We could mutate them, but that seems very much against the AbstractMCMC.jl interface...

Hmm. Do you mean it actually only make sense to make the log density and gradient as model but treating metric differently?

"Only" is maybe strong, but yeah, I think it makes more sense to keep the hamiltonian and the metric separate. Potentially do away completely with the hamiltonian and only keep metric or, alternatively, make hamiltonian nothing but a wrapper around metric (in case there are future cases we might want to cover) and just make a DifferentiableDensityModel <: AbstractModel (could even move this + DensityModel from AbstractMH.jl into AbstractMCMC.jl at some point, so these can be re-used e.g. for MALA which is present in AdvancedHMC.jl)

EDIT: I think it's best to keep hamiltonian, but just reconstruct it from a DifferentiableDensityModel at every call to step. And we just pass the "state" of hamiltonian around, e.g. metric.

xukai92 · 2021-04-06T20:43:25Z

src/sampler.jl

+    rng::AbstractRNG,
+    model::HamiltonianModel,
+    spl::HMCKernel;
+    init_params, # TODO: implement this. Just do `rand`? Need dimensionality though.


I can see multiple ways to do this but it's unclear what's the best.
Can you file a PR/discussion to chat on this?

Will do! But for now I'll just extract the dimensionality from the metric and use rand. But yeah, I'll open an issue after this PR so we can figure out potential better defaults.

xukai92 · 2021-04-06T20:47:07Z

So the question is: should we do this so as to preserve the current logging functionality?

I think we should do this as to keep the functionality unchanged.
We can decide later if that's hard to maintain or not being as useful.

To replicate the current summarization functionality ...

Does it make sense to add this functionality in AbtractMCMC first? Seems to be a useful thing to have in general.

cpfiffer · 2021-04-06T21:15:54Z

To replicate the current summarization functionality (e.g. inform the user of average acceptance rates and EBFMI) as a post-sample step, we can overload StatsBase.sample and then perform this step after the call to AbstractMCMC.mcmcsample. Should we do this?

Yes, that's the intended use case. 👍

AFAIK, the way to customize the logging in AbstractMCMC.jl is to pass progress=false to the underlying AbstractMCMC.mcmcsample and then use the callback keyword argument to log the progress. So the question is: should we do this so as to preserve the current logging functionality?

Yeah, I think the AHMC logging stuff should be moved to a callback.

torfjelde · 2021-04-06T21:28:31Z

Does it make sense to add this functionality in AbtractMCMC first? Seems to be a useful thing to have in general.

Not needed; with TuringLang/AbstractMCMC.jl#56 merged we can implement this as a callback very easily, and just make this callback the default.

xukai92 · 2021-04-09T13:05:49Z

Project.toml

@@ -3,6 +3,7 @@ uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 version = "0.2.28"


Please also increment the version here.

Should we make a minor version bump? Seems like quite a significant change given the introduction of HMCKernel too

Yeah a minor version bump sounds suitable.

torfjelde · 2021-04-09T13:50:40Z

Btw, waiting for TuringLang/Turing.jl#1579 since tests depend on Turing.

yebai · 2021-05-25T11:31:46Z

Is there anything missing from this PR? If there is nothing major, let's try to merge this PR since it will benefit several downstream tasks.

HarrisonWilde · 2021-05-25T13:54:04Z

I've been working off the branch for this PR to support tempering on AHMC via the new MCMCTempering package and noticed that upon finishing sampling, even if I have using MCMCChains, I still just get back a vector of transitions from sampling, I think it would be nice to add a function similar to the one in Turing/Inference that implements AbstractMCMC.bundle_samples to actually make use of Chains, what do others think @torfjelde @yebai ?

torfjelde · 2021-05-25T15:10:55Z

Is there anything missing from this PR? If there is nothing major, let's try to merge this PR since it will benefit several downstream tasks.

A couple of things:

Not happy with how we have a DifferentiableModel in this repo. This should preferably go somewhere else, e.g. AMCMC. Alternatively we can make use of LogDensityProblems.jl, which might be a good approach (me and @devmotion discussed this on slack a couple of weeks ago).
I've noticed something that might be a bug, so need to debug that first.

I've been working off the branch for this PR to support tempering on AHMC via the new MCMCTempering package and noticed that upon finishing sampling, even if I have using MCMCChains, I still just get back a vector of transitions from sampling, I think it would be nice to add a function similar to the one in Turing/Inference that implements AbstractMCMC.bundle_samples to actually make use of Chains, what do others think @torfjelde @yebai ?

I'm also using this branch heavily in some work, and have an example of such a method implemented locally. But my implementation is using information from VarInfo after a model-run in Turing, so it's not independent of Turing.jl. This is because the MCMCChains.jl-framework is mainly nice if you have several different variables while in AHMC there's no separation/labelling of the different parameters and thus we don't get much from MCMCChains other than statistics (which are already implemented for arrays, i.e. no need to make a Chain first). The issue is that we can't really put this in AHMC, as it we don't want to depend on MCMCChains here.

How are you using it? Are you using it with Turing or nah?

HarrisonWilde · 2021-05-25T15:52:43Z

I see yes, that is a fair assessment, I am using just AHMC, haven't tried with Turing but I agree it perhaps makes less sense to have something specific here for working with AHMC. On another note, I agree regarding Tor's comments on the state of the model currently, it would be better perhaps to rely on some standard definition in AMCMC.

What is the minimal expectation for the return from sample though? As currently it appears to be this vector of transitions rather than the more standard vector of thetas you see on master?

torfjelde · 2021-05-25T16:35:11Z

What is the minimal expectation for the return from sample though? As currently it appears to be this vector of transitions rather than the more standard vector of thetas you see on master?

Glad you brought this up! I'm not certain tbh 😕 @xukai92 @yebai What do you think? Should it return PhasePoint or just thetas? I'm with whatever, but made it return PhasePoint initially.

xukai92 · 2021-05-25T21:32:21Z

What do you think? Should it return PhasePoint or just thetas?

I think we should not change the return at least for this PR.
Plus now stuff like ComponentArrays is supported, you can even have simple "variable bindings".

Not happy with how we have a DifferentiableModel in this repo. This should preferably go somewhere else, e.g. AMCMC. Alternatively we can make use of LogDensityProblems.jl, which might be a good approach (me and @devmotion discussed this on slack a couple of weeks ago).

Both are OK but I don't think we have to make it perfect within this PR as long as it's used behind the scene (i.e. users don't need to create it explicitly).
Just hope this PR doesn't keep growing than interface binding.

I've noticed something that might be a bug, so need to debug that first.

What's the potential bug?

xukai92 · 2021-07-15T00:17:45Z

The CI workflow also fails here. Any idea?

torfjelde · 2021-07-15T00:23:11Z

The CI workflow also fails here. Any idea?

It might just be that the settings I have for the sampler is bad because it passes most of the time.

xukai92 · 2021-07-15T00:26:37Z

It might just be that the settings I have for the sampler is bad because it passes most of the time.

They were fine in the other PR so I suppose you mean it's due to randomness?

Seems like there are also conflicts.
Let's resolve them and see if the CI passes with a re-run.

torfjelde · 2021-07-15T00:29:07Z

Btw, one annoying thing is that find_good_stepsize isn't supported by DifferentiableDensityModel yet. There's a bit too much functionality that uses the hamiltonian in there atm, so I haven't bothered trying to refactor that part yet. But this is something we can always do later anyways.

xukai92 · 2021-07-15T01:12:19Z

Btw, one annoying thing is that find_good_stepsize isn't supported by DifferentiableDensityModel yet. There's a bit too much functionality that uses the hamiltonian in there atm, so I haven't bothered trying to refactor that part yet. But this is something we can always do later anyways.

Doing it in another PR sounds more reasonable. Can you create an issue for this along with other missing thing from DifferentiableDensityModel or even other alternatives (I believe we discussed a few).

xukai92 · 2021-07-15T01:13:08Z

How could IntegrationTests also pass even Turing is not compatible?

torfjelde · 2021-07-15T01:16:23Z

Doing it in another PR sounds more reasonable. Can you create an issue for this along with other missing thing from DifferentiableDensityModel or even other alternatives (I believe we discussed a few).

Will do 👍

How could IntegrationTests also pass even Turing is not compatible?

Because this is a new minor version for AHMC, which won't be compat with Turing.jl since it has compat-bounds on AHMC + it indicates it's a breaking release and so Turing.jl should be expected to have to make some changes to be compatible with the new version. If you look at the integration tests you see Info: Not compatible with this release. No problem..

So it's just a question of whether we make the tests pass or fail in this case; I went with letting it pass since this is also what we do in other Turing.jl-packages, e.g. DPPL.

torfjelde · 2021-07-15T01:16:56Z

Okay this is just strange. It fails on macOS but seems to do fine otherwise.

xukai92 · 2021-07-15T01:29:39Z

So it's just a question of whether we make the tests pass or fail in this case; I went with letting it pass since this is also what we do in other Turing.jl-packages, e.g. DPPL.

I see. Makes sense to me now.

Okay this is just strange. It fails on macOS but seems to do fine otherwise.

Could be randomness again but I feel using the same parameters as the non-AMCMC version is enough. Otherwise it might indicate some issues.

If I give same initial position and same random seed, do I expect to get exact same samples from the interface implementation? If so maybe we can add a test to check if samples are exactly the same.

torfjelde · 2021-07-15T12:20:16Z

I realized what was causing the issue: I was including the adapted samples in the mean-computation:)

Now that those are dropped, it works. But now there's a different issue? This one seems completely unrelated to the changes in this PR though. You got any clue @xukai92 ?

xukai92 · 2021-07-15T13:09:41Z

Looks like the tests are fine now. Is turning off the progress (the last commit) the fix?

torfjelde · 2021-07-15T13:36:10Z

Looks like the tests are fine now. Is turning off the progress (the last commit) the fix?

No, that was just to make the logs nicer. It seems like one of the adaptation tests can fail if we're unlucky with the random seed, but as I said, this has nothing to do with this PR so IMO we merge and I'll make an issue reminding us that the adaptation test is fickle.

initial work on adopting the AbstractMCMC interface

f36a3cc

removed old sample method since is now redundant

16a233c

torfjelde requested a review from xukai92 April 6, 2021 18:42

xukai92 reviewed Apr 6, 2021

View reviewed changes

torfjelde added 8 commits April 7, 2021 00:02

AbstractMCMCKernel is now a subtype of AbstractSampler

a1494ac

fixed impl and added back old one

5f01c5d

overload AbstractMCMC.sample to make interface a bit nicer

fc9e747

reverted renaming of abstract kernel type

ba0d623

replicate logging behavior of current AHMC when using AbstractMCMC

b6ad39c

initial work on tests, but require update to testing deps

66b37b5

Merge branch 'master' into tor/abstractmcmc

ea83745

moved abstractmcmc interface into separate file

f041ff6

xukai92 reviewed Apr 9, 2021

View reviewed changes

torfjelde added 4 commits April 12, 2021 15:32

version bump

dad71a8

only bump patch version to see if tests succeed

07be7ac

Merge branch 'master' into tor/abstractmcmc

ae16796

Merge branch 'master' into tor/abstractmcmc

d3e8542

HarrisonWilde marked this pull request as ready for review April 28, 2021 14:46

HarrisonWilde marked this pull request as draft April 28, 2021 14:47

torfjelde added 2 commits July 14, 2021 15:42

bumped minor version

de4e33a

increased number of samples used in abstractmcmc tests

fa17e41

remove thinning from tests

8f8e64c

make initial Leapfrog step size smaller

8d13ff5

Merge branch 'master' into tor/abstractmcmc

ebaec71

torfjelde marked this pull request as ready for review July 15, 2021 00:27

mistakenly removed AbstractMCMC as a test dep in previous commit

ccdc832

torfjelde added 2 commits July 15, 2021 02:17

increase adaptation to see if it helps

5a417bc

ensure we drop the adaptation samples in the test

32f6ff8

torfjelde added 2 commits July 15, 2021 03:53

made a mistake apparently

13a07e8

think I finally fixed the tests

b82a965

disable progress in test

38408cf

yebai approved these changes Jul 15, 2021

View reviewed changes

torfjelde merged commit 7cad9f0 into master Jul 15, 2021

delete-merged-branch bot deleted the tor/abstractmcmc branch July 15, 2021 13:49

This was referenced Dec 16, 2021

Interface with AbstractMCMC.jl #211

Closed

Pair values with metadata? #101

Closed

yebai mentioned this pull request Jan 27, 2023

Integration of TBLogger or WBLogger #171

Closed

		@@ -3,6 +3,7 @@ uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
		version = "0.2.28"

Adopt AbstractMCMC.jl interface #259

Adopt AbstractMCMC.jl interface #259

Conversation

torfjelde commented Apr 6, 2021 • edited Loading

Things to discuss

cpfiffer commented Apr 6, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

torfjelde Apr 7, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

xukai92 commented Apr 6, 2021

cpfiffer commented Apr 6, 2021 • edited Loading

torfjelde commented Apr 6, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

torfjelde commented Apr 9, 2021

yebai commented May 25, 2021

HarrisonWilde commented May 25, 2021

torfjelde commented May 25, 2021

HarrisonWilde commented May 25, 2021 • edited Loading

torfjelde commented May 25, 2021

xukai92 commented May 25, 2021

xukai92 commented Jul 15, 2021

torfjelde commented Jul 15, 2021

xukai92 commented Jul 15, 2021

torfjelde commented Jul 15, 2021

xukai92 commented Jul 15, 2021

xukai92 commented Jul 15, 2021

torfjelde commented Jul 15, 2021

torfjelde commented Jul 15, 2021

xukai92 commented Jul 15, 2021

torfjelde commented Jul 15, 2021

xukai92 commented Jul 15, 2021

torfjelde commented Jul 15, 2021

torfjelde commented Apr 6, 2021 •

edited

Loading

torfjelde Apr 7, 2021 •

edited

Loading

cpfiffer commented Apr 6, 2021 •

edited

Loading

HarrisonWilde commented May 25, 2021 •

edited

Loading