new sampler: drghmc

[gil2rok]

Implemented sampler for (Probabilistic) Delayed Rejection Generalized HMC with basic test.

This pull request is a cleaner version of my previous pull request #38 , which contained a messy commit history.

[codecov-commenter]

Codecov Report

Merging #39 (403c8ca) into main (d32dcb5) will increase coverage by 0.57%. The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main      #39      +/-   ##
==========================================
+ Coverage   98.01%   98.59%   +0.57%     
==========================================
  Files          11       12       +1     
  Lines         302      426     +124     
==========================================
+ Hits          296      420     +124     
  Misses          6        6              

Files	Coverage Δ
bayes_kit/__init__.py	100.00% <100.00%> (ø)
bayes_kit/drghmc.py	100.00% <100.00%> (ø)

:mega: We’re building smart automated test selection to slash your CI/CD build times. Learn more

[gil2rok]

Thanks for the (super detailed!) feedback.

I plan to:

1. change the documentation as per @bob-carpenter 's comments
2. change my use of | as a Union operator as per @WardBrian's comment
3. add support for gradient cacheing as per @modichirag 's in-person discussion.

FYI: over the next few weeks I will be prioritizing my actual research project so progress on these changes will be a bit slower. However, I am still excited to continue working on this.

[gil2rok]

I have addressed every comment by (1) resolving it after incorporating the feedback or (2) responding with a question or clarifying what I meant. Feel free to resolve the comment if, after my response, there is nothing left to discuss.

A few more questions:

1. I want to clarify notation:

   • theta represents model parameters and rho represents auxiliary momentum
   • a draw is a tuple of model parameters theta and its log density
   • a sample is a sequence of draws (as @bob-carpenter commented above)

   Then what term should I use to refer to (theta, rho) in the extended phase space when writing the DRGHMC docstrings? I want to be able to refer to the current (theta, rho), proposed (theta, rho), and ghost (theta, rho).

2. What tests do you recommend I include? I am already testing for:

   • The number of density + gradient evaluations
   • Dummy examples with standard normal and binomial
   • Various errors are thrown correctly

3. I want to confirm the correct implementation of the leapfrog integrator in the format @bob-carpenter wrote in an earlier comment here.

   Bob wrote that his implementation should save a gradient evaluation and "avoids having to save that grad value out of the loop". However, I am not sure if my implementation does this.

   # pos = position
   # mo = momentum
   # grad = gradient
   
   def leapfrog(pos, mo):
       grad = get_grad(pos)
       mo_mid = mo + (0.5 * stepsize) * grad
       pos += stepsize * mo_mid
   
       for i in range(stepcount - 1):
           grad = get_grad(pos)
           mo_mid += stepsize * grad
           pos += stepsize * mo_mid
   
       grad = get_grad(pos)
       mo = mo_mid + (0.5 * stepsize) * grad
       return (pos, mo)

[bob-carpenter]

1a. theta and rho are fine with me 1b. A draw consists of both a theta and a rho. 1c. Yes, given the correction to 1b.

The point is that if we have draws of multiple variables, we can throw some of the variables away to get marginal draws from the remaining variable. Usually we only care about the theta draws for downstream Bayesian inference.

(theta, rho) is just the phase space (no "extended") and it's conventional to just use the pair and not try to reduce to a single variable.

2. You want to test all the boundary conditions of all your algorithms. And otherwise provide a range of tests that test all the behavior for "regular" inputs. I wouldn't call normal and binomial examples "dummy".

3. Yes, that looks good. You can reason through the algorithm and verify inductively that there's only one gradient eval for each position instead of an extra one as the original implementation had.

[gil2rok]

All comments are addressed, code coverage is 100%, everything is properly formatted and typed (thanks to @WardBrian for help on this) and all tests are passing (except for specific functions in test_iat or test_drghmc that sometimes fails depending on the rng seed).

I would love some more feedback: what else is needed for this PR to be merged?

[bob-carpenter]

If you're willing to wait, I can do this net week (Aug 21+).

[gil2rok]

That sounds great!

[gil2rok]

I believe this pull request is ready for review -- all the comments have been addressed.

Can @bob-carpenter and/or @WardBrian take a look at this when you have a chance? Would be much appreciated!

[gil2rok]

A couple general questions/notes:

1. What is the recommended line length? The black code formatter seems to suggest 88 characters, while Google's style guide recommends 80 characters in docstrings (link is found here ). My code currently adheres to the 88 character line limit.
2. The Google style guide recommends every function docstring starts with a one line summary. Some of @bob-carpenter's comments on my one line docstring summaries are much clearer than what I wrote, but are much longer than this 88 character limit. Wanted to briefly explain why I will not just be accepting these suggestions as is, and instead will edit them down.
3. I want to ensure I am using the correct terminology for draws. As per, https://github.com/flatironinstitute/bayes-kit/pull/39#issuecomment-1671758891, @bob-carpenter wrote that a draw is a tuple of a position variable and momentum variable, denoted by (theta, rho). If this is so, why does the DrawAndLogP type define a draw as a single VectorType variable theta instead of the tuple (theta, rho)? This seems inconsistent.

[WardBrian]

My two cents:

1. recommended line length

I think 88 characters is good for code and comments

2. one line summary.

I tend to focus more on it being one (preferably simple) sentence rather than exactly one line. I think the spirit of the guideline is just that there should be a summary to start

[gil2rok]

Yesterday I met with @WardBrian (huge thank you!) to finalize some changes, summarized below, that address the remaining comments from @bob-carpenter :

1. Bayesian Analysis journal name is now italisized.
2. leapfrog_stepsizes is renamed to leapfrog_step_sizes. Also, leapfrog_stepcounts is renamed to leapfrog_step_counts.
3. Validation logic is reworked such that my validation functions no longer return their input unnecessarily.
4. leapfrog_step_sizes can now be a list containingint or float values. This is because I changed its container data type from a list to sequence.

I am now formatting the code + updating my tests. Afterwards, can @bob-carpenter review it one last time to confirm it is ready to merge?

[bob-carpenter]

Yes, happy to review. Just let me know when it's ready. This is something where it's probably good for both me and @WardBrian to review for the sampler and Python style respectively.

[gil2rok]

Quick question: for typing reasons, leapfrog_step_sizes and leapfrog_step_counts are of type Sequence[float] and Sequence[int] respectively.

In the documentation should I refer to these as lists or sequences? The term list is simple to understand, but the term sequence is more technically precise.

For example, should the error message read "expected leapfrog_step_counts to be a list" or should the error message read "expected leapfrog_step_counts to be a sequence"?

[bob-carpenter]

I would use "sequence" in the written doc (not capitalized---common nouns are never capitalized in English). The type will show up in the doc as Sequence to reinforce that. List is just one of many types of Sequence. But tuples and numpy arrays are also of type Sequence.

[WardBrian]

(@gil2rok : Your tests are raising a type error currently, which should probably be fixed before actually merging 😉 )

[gil2rok]

@WardBrian I am aware of the type issues!

I am fixing them on my local machine and also updating the tests. I should push these changes to the repo later today!

[bob-carpenter]

@WardBrian says:

The few comments I think I would leave at this point end up going against some of the things @bob-carpenter suggested earlier

@WardBrian is better than me at Python, so if he has conflicting advice w.r.t. Python, you should go with his suggestions.

[gil2rok]

Tests are updated and passing (except for stochasticity that makes them occasionally fail). All files are properly formatted, code coverage is 100%, and all comments are addressed.

@WardBrian: are there any Python style changes you would recommend, as per Bob's comment above? If not, is there anything else to do before merging this pull request?

[WardBrian]

@gil2rok @bob-carpenter

I am happy with the PR as it stands now. The things I have in mind are more about conventions around e.g. input checking or error handling that would apply to both this but also the existing code in the repo.

I'd like to discuss it (at some point), but a follow on PR would be necessary in either case (either to align the rest of the repo to what was done here, or to adjust this code to whatever we decide)

[gil2rok]

Very excited the pull request was merged!

@bob-carpenter We should also change the README to include DRGHMC. Is it all right if I make that change?

[bob-carpenter]

Yes, please do change the README. You can just push that w/o review.