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)
 
@@ -32,17 +47,17 @@ where `ϵ` is the step size of leapfrog integration.
 ### Adaptor (`adaptor`)
 
   - Adapt the mass matrix `metric` of the Hamiltonian dynamics: `mma = MassMatrixAdaptor(metric)`
-
       + This is lowered to `UnitMassMatrix`, `WelfordVar` or `WelfordCov` based on the type of the mass matrix `metric`
       + There is an experimental way to improve the *diagonal* mass matrix adaptation using gradient information (similar to [nutpie](https://github.com/pymc-devs/nutpie)),
-      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)`
-
       + It uses Nesterov's dual averaging with `δ` as the target acceptance rate.
+
   - Combine the two above *naively*: `NaiveHMCAdaptor(mma, ssa)`
+
   - Combine the first two using Stan's windowed adaptation: `StanHMCAdaptor(mma, ssa)`
 
 ## The `sample` functions
@@ -65,14 +80,15 @@ sample(
 Draw `n_samples` samples using the kernel `κ` under the Hamiltonian system `h`
 
   - The randomness is controlled by `rng`.
-
       + If `rng` is not provided, the default random number generator (`Random.default_rng()`) will be used.
 
   - The initial point is given by `θ`.
-  - The adaptor is set by `adaptor`, for which the default is no adaptation.
 
+  - The adaptor is set by `adaptor`, for which the default is no adaptation.
       + It will perform `n_adapts` steps of adaptation, for which the default is `1_000` or 10% of `n_samples`, whichever is lower.
+
   - `drop_warmup` specifies whether to drop samples.
+
   - `verbose` controls the verbosity.
   - `progress` controls whether to show the progress meter or not.
 

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