Fix NUTS

[rlouf]

Update: The has_subtree_terminated condition was not used to stop the tree expansion which made the sampler biased. I corrected this (and simplified the sampling tests).

• [x] currently the algorithm only runs the integrator once. This is likely due to a problem initializing the termination state.
• [x] We are having similar shape issues in termination.py to the one we had in metrics.py

[codecov[bot]]

Codecov Report

Merging #45 (ff77a1c) into main (a6c1c5a) will not change coverage. The diff coverage is 100.00%.

@@            Coverage Diff            @@
##              main       #45   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           13        13           
  Lines          468       466    -2     
  Branches        25        27    +2     
=========================================
- Hits           468       466    -2     

Impacted Files	Coverage Δ
aehmc/integrators.py	100.00% <100.00%> (ø)
aehmc/metrics.py	100.00% <100.00%> (ø)
aehmc/nuts.py	100.00% <100.00%> (ø)
aehmc/proposals.py	100.00% <100.00%> (ø)
aehmc/termination.py	100.00% <100.00%> (ø)
aehmc/trajectory.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update a6c1c5a...ff77a1c. Read the comment docs.

[rlouf]

I ended up doing a step-by-step comparison with blackjax and it turns out there is one bug left in termination.py in the is_iterative_turning function. The function should return False when idx_max < idx_min but is currently returning True instead.

I am working on a quick fix, but the ideal would be to address this issue in aesara.

[rlouf]

Ah another instance of a test that fails when running pytest in the root folder, but passes when I do pytest tests/test_hmc.py. There's clearly something I don't understand about pytest.

[brandonwillard]

Ah another instance of a test that fails when running pytest in the root folder, but passes when I do pytest tests/test_hmc.py. There's clearly something I don't understand about pytest.

That could be due to dependencies between the RNG seeding and the test-running process. For example, if tests alter a shared/global RNG state, then they can fail when run in a given order.