Skip to content

Consolidate RHMC Changes#481

Closed
J-Price-3 wants to merge 30 commits intoqqy/NEW_RMHMCfrom
Jamie_RHMC
Closed

Consolidate RHMC Changes#481
J-Price-3 wants to merge 30 commits intoqqy/NEW_RMHMCfrom
Jamie_RHMC

Conversation

@J-Price-3
Copy link

Changes include:

  • Type stability (updated core RHMC functions as well as utility functions in research folder)
  • Small optimisations
  • Addition of Implicit Midpoint Integrator

We should consolidate all the RHMC changes so that there is one PR to the main branch.

@J-Price-3 J-Price-3 self-assigned this Dec 2, 2025
@github-actions
Copy link
Contributor

github-actions bot commented Dec 4, 2025

AdvancedHMC.jl documentation for PR #481 is available at:
https://TuringLang.github.io/AdvancedHMC.jl/previews/PR481/

J-Price-3 and others added 5 commits December 4, 2025 13:18
… gradients, but not changing the schedule). (#473)

* initial changes to get a working demo

* fix tests, add new ones, and add documentation

* fix type

* fix some stray tests

* delete tmp folder with demo

* address review comments

* reference interface refactor issue

* refactor and fix test rng handling

* improve docstring for NutpieVar

* remove superfluous white space

* fix JET tests

* add entry to history, bump version

* fix NutpieVar docstring

* increase number of tests for mass matrix adaptation

---------

Co-authored-by: Markus Hauru <markus@mhauru.org>
Comment on lines +6 to +8
currently to be initialized for a `metric` of type `DiagEuclideanMetric`
via `mma = AdvancedHMC.NutpieVar(size(metric); var=copy(metric.M⁻¹))`
until a new interface is introduced in an upcoming breaking release to specify the method of adaptation.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
currently to be initialized for a `metric` of type `DiagEuclideanMetric`
via `mma = AdvancedHMC.NutpieVar(size(metric); var=copy(metric.M⁻¹))`
until a new interface is introduced in an upcoming breaking release to specify the method of adaptation.
currently to be initialized for a `metric` of type `DiagEuclideanMetric`
via `mma = AdvancedHMC.NutpieVar(size(metric); var=copy(metric.M⁻¹))`
until a new interface is introduced in an upcoming breaking release to specify the method of adaptation.


- Adapt the mass matrix `metric` of the Hamiltonian dynamics: `mma = MassMatrixAdaptor(metric)`

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change

Comment on lines +38 to +40
currently to be initialized for a `metric` of type `DiagEuclideanMetric`
via `mma = AdvancedHMC.NutpieVar(size(metric); var=copy(metric.M⁻¹))`
until a new interface is introduced in an upcoming breaking release to specify the method of adaptation.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
currently to be initialized for a `metric` of type `DiagEuclideanMetric`
via `mma = AdvancedHMC.NutpieVar(size(metric); var=copy(metric.M⁻¹))`
until a new interface is introduced in an upcoming breaking release to specify the method of adaptation.
currently to be initialized for a `metric` of type `DiagEuclideanMetric`
via `mma = AdvancedHMC.NutpieVar(size(metric); var=copy(metric.M⁻¹))`
until a new interface is introduced in an upcoming breaking release to specify the method of adaptation.


- Adapt the step size of the leapfrog integrator `integrator`: `ssa = StepSizeAdaptor(δ, integrator)`

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change


- The randomness is controlled by `rng`.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change

Comment on lines +95 to +97
z = PhasePoint(
θ, θ, DualValue(0., θ), DualValue(0., θ)
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
z = PhasePoint(
θ, θ, DualValue(0., θ), DualValue(0., θ)
)
z = PhasePoint(θ, θ, DualValue(0.0, θ), DualValue(0.0, θ))

pc1 = MassMatrixAdaptor(UnitEuclideanMetric) # default dim = 2
pc2 = MassMatrixAdaptor(DiagEuclideanMetric)
# Constructing like this until we've settled on a different interface
pc2_nutpie = NutpieVar{Float64}((2, ))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
pc2_nutpie = NutpieVar{Float64}((2, ))
pc2_nutpie = NutpieVar{Float64}((2,))

)
# Constructing like this until we've settled on a different interface
adaptor2_nutpie = StanHMCAdaptor(
NutpieVar{Float64}((2, )), NesterovDualAveraging(0.8, 0.5)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
NutpieVar{Float64}((2, )), NesterovDualAveraging(0.8, 0.5)
NutpieVar{Float64}((2,)), NesterovDualAveraging(0.8, 0.5)

# For this target, Nutpie (without regularization) will arrive at the true variances after two draws.
res_nutpie = runnuts_nutpie(ℓπ, DiagEuclideanMetric(D))
@test res.adaptor.pc.var ≈ σ² rtol = 0.2
@test preconditioned_cond(res_nutpie.adaptor.pc, Σ) < preconditioned_cond(res.adaptor.pc, Σ)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
@test preconditioned_cond(res_nutpie.adaptor.pc, Σ) < preconditioned_cond(res.adaptor.pc, Σ)
@test preconditioned_cond(res_nutpie.adaptor.pc, Σ) <
preconditioned_cond(res.adaptor.pc, Σ)

# find the best preconditioner for the target.
# As these are statistical algorithms, superiority is not always guaranteed, hence this way of testing.
res_nutpie = runnuts_nutpie(ℓπ, DiagEuclideanMetric(D))
n_nutpie_superior += preconditioned_cond(res_nutpie.adaptor.pc, Σ) < preconditioned_cond(res.adaptor.pc, Σ)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
n_nutpie_superior += preconditioned_cond(res_nutpie.adaptor.pc, Σ) < preconditioned_cond(res.adaptor.pc, Σ)
n_nutpie_superior +=
preconditioned_cond(res_nutpie.adaptor.pc, Σ) <
preconditioned_cond(res.adaptor.pc, Σ)

return -logZ - dot(r, h.metric._temp) / 2
#! Eq (14) of Girolami & Calderhead (2011)
function ∂H∂r(
h::Hamiltonian{<:AbstractRiemannianMetric}, θ::AbstractVecOrMat{T}, r::AbstractVecOrMat{T}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
h::Hamiltonian{<:AbstractRiemannianMetric}, θ::AbstractVecOrMat{T}, r::AbstractVecOrMat{T}
h::Hamiltonian{<:AbstractRiemannianMetric},
θ::AbstractVecOrMat{T},
r::AbstractVecOrMat{T},

# println("θ: ", θ)
# println("H: ", H)
# end
G = h.metric.map(H)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

map isn't defined for an arbitrary Riemannian metric. You'll need to introduce a layer of abstraction.

A abstract metric could have a get_G method which internally calls map if it is a Dense metric

@nsiccha
Copy link
Contributor

nsiccha commented Jan 16, 2026

This can be closed in favour of #485, right?

@J-Price-3 J-Price-3 closed this Jan 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants