Unified Riemannian Metric API

docs/src/api.md

@@ -10,7 +10,22 @@ This modularity means that different HMC variants can be easily constructed by c
   - Dense metric: `DenseEuclideanMetric(dim)`
   - Rank update metric: `RankUpdateEuclideanMetric(dim)`
 
-where `dim` is the dimensionality of the sampling space.
+where `dim` is the dimension of the sampling space.
+
+Two experimental position-dependent (Riemannian) metrics are also available:
+
+  - `RiemannianMetric((dim,), calc_G, calc_∂G∂θ)` — for user-supplied positive-definite
+    metrics `G(θ)` (e.g. Fisher information). `calc_G` should return either a plain
+    `Matrix` or an `AbstractPDMat` (preferred — reuses the stored Cholesky). `calc_∂G∂θ`
+    returns the `(d, d, d)` tensor `∂G/∂θ`.
+  - `SoftAbsRiemannianMetric((dim,), calc_H, calc_∂H∂θ, α)` — for Hessian-based metrics
+    where `H(θ)` is not guaranteed to be positive definite. The SoftAbs transformation
+    `G = Q · diag(λ · coth(αλ)) · Qᵀ` (Betancourt, 2012) regularises `H`'s eigenvalues
+    to a strictly positive spectrum. `α` controls how closely SoftAbs approximates `|λ|`.
+
+The legacy `DenseRiemannianMetric(dim, G, ∂G∂θ[, map])` constructor is deprecated and
+forwards to the appropriate type above based on whether `map` is `IdentityMap()` or
+`SoftAbsMap(α)`.
 
 ### [Integrator (`integrator`)](@id integrator)
 

src/AdvancedHMC.jl

@@ -14,7 +14,11 @@ using LinearAlgebra:
     Diagonal,
     AbstractQ,
     qr,
-    lmul!
+    lmul!,
+    logdet,
+    tr,
+    eigen,
+    diagm
 using IrrationalConstants: loghalf
 using LogExpFunctions: logaddexp, logsumexp
 using Random: Random, AbstractRNG
@@ -71,6 +75,13 @@ export Leapfrog, JitteredLeapfrog, TemperedLeapfrog
 include("riemannian/integrator.jl")
 export GeneralizedLeapfrog
 
+include("riemannian/metric.jl")
+export RiemannianMetric, SoftAbsRiemannianMetric
+# Deprecated exports (for backward compatibility)
+export IdentityMap, SoftAbsMap, DenseRiemannianMetric
+
+include("riemannian/hamiltonian.jl")
+
 include("trajectory.jl")
 export Trajectory,
     HMCKernel,

src/hamiltonian.jl

@@ -113,7 +113,11 @@ function Base.similar(z::PhasePoint{<:AbstractVecOrMat{T}}) where {T<:AbstractFl
 end
 
 function phasepoint(
-    h::Hamiltonian, θ::T, r::T; ℓπ=∂H∂θ(h, θ), ℓκ=DualValue(neg_energy(h, r, θ), ∂H∂r(h, r))
+    h::Hamiltonian,
+    θ::T,
+    r::T;
+    ℓπ=∂H∂θ(h, θ),
+    ℓκ=DualValue(neg_energy(h, r, θ), ∂H∂r(h, θ, r)),
 ) where {T<:AbstractVecOrMat}
     return PhasePoint(θ, r, ℓπ, ℓκ)
 end
@@ -127,7 +131,7 @@ function phasepoint(
     _r::T2;
     r=safe_rsimilar(θ, _r),
     ℓπ=∂H∂θ(h, θ),
-    ℓκ=DualValue(neg_energy(h, r, θ), ∂H∂r(h, r)),
+    ℓκ=DualValue(neg_energy(h, r, θ), ∂H∂r(h, θ, r)),
 ) where {T1<:AbstractVecOrMat,T2<:AbstractVecOrMat}
     return PhasePoint(θ, r, ℓπ, ℓκ)
 end