diff --git a/docs/src/api.md b/docs/src/api.md
index 663973a0..13cb147c 100644
--- a/docs/src/api.md
+++ b/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)
 
diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl
index 17b4ceac..e4f0dda9 100644
--- a/src/AdvancedHMC.jl
+++ b/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,
diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl
index c57f3be8..e5376f78 100644
--- a/src/hamiltonian.jl
+++ b/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
diff --git a/src/riemannian/hamiltonian.jl b/src/riemannian/hamiltonian.jl
index 6f051ffb..6d134d81 100644
--- a/src/riemannian/hamiltonian.jl
+++ b/src/riemannian/hamiltonian.jl
@@ -1,358 +1,242 @@
-using Random
+"""
+    tr_product(A, B)
 
-### integrator.jl
+Compute `tr(A * B)` for square matrices in O(n²) without forming the product.
+Uses the identity: tr(A * B) = sum(A' .* B)
+"""
+tr_product(A::AbstractMatrix, B::AbstractMatrix) = sum(Base.broadcasted(*, A', B))
 
-import AdvancedHMC: ∂H∂θ, ∂H∂r, DualValue, PhasePoint, phasepoint, step
-using AdvancedHMC: TYPEDEF, TYPEDFIELDS, AbstractScalarOrVec, AbstractLeapfrog, step_size
+"""
+    tr_product(A, v)
 
+Compute `tr(A * v * v')` = v' * A * v efficiently.
 """
-$(TYPEDEF)
+tr_product(A::AbstractMatrix, v::AbstractVector) = dot(v, A * v)
 
-Generalized leapfrog integrator with fixed step size `ϵ`.
+####
+#### Gradient cache for θ-dependent computations
+####
 
-# Fields
+"""
+    RiemannianGradCache{T, TG, TP}
+
+Cache for θ-dependent computations in Riemannian HMC gradient calculation.
+This allows reusing expensive eigendecomposition/factorization across fixed-point iterations.
 
-$(TYPEDFIELDS)
+# Fields
+- `G_eval`: Evaluated metric (SoftAbsEval or matrix)
+- `∂P∂θ`: Pre-metric sensitivities, shape (d, d, d)
+- `ℓπ`: Log density value at θ
+- `∂ℓπ∂θ`: Log density gradient at θ
+- `logdet_terms`: Precomputed 0.5 * tr(M_logdet * ∂P∂θ[:,:,i]) for each i
 """
-struct GeneralizedLeapfrog{T<:AbstractScalarOrVec{<:AbstractFloat}} <: AbstractLeapfrog{T}
-    "Step size."
-    ϵ::T
-    n::Int
-end
-function Base.show(io::IO, l::GeneralizedLeapfrog)
-    return print(io, "GeneralizedLeapfrog(ϵ=", round.(l.ϵ; sigdigits=3), ", n=", l.n, ")")
+struct RiemannianGradCache{T,TG,TP}
+    G_eval::TG
+    ∂P∂θ::TP
+    ℓπ::T
+    ∂ℓπ∂θ::Vector{T}
+    logdet_terms::Vector{T}
 end
 
-# Fallback to ignore return_cache & cache kwargs for other ∂H∂θ
-function ∂H∂θ_cache(h, θ, r; return_cache=false, cache=nothing) where {T}
-    dv = ∂H∂θ(h, θ, r)
-    return return_cache ? (dv, nothing) : dv
-end
+"""
+    build_grad_cache(h::Hamiltonian{<:AbstractRiemannianMetric}, θ)
 
-# TODO Make sure vectorization works
-# TODO Check if tempering is valid
-function step(
-    lf::GeneralizedLeapfrog{T},
-    h::Hamiltonian,
-    z::P,
-    n_steps::Int=1;
-    fwd::Bool=n_steps > 0,  # simulate hamiltonian backward when n_steps < 0
-    full_trajectory::Val{FullTraj}=Val(false),
-) where {T<:AbstractScalarOrVec{<:AbstractFloat},P<:PhasePoint,FullTraj}
-    n_steps = abs(n_steps)  # to support `n_steps < 0` cases
-
-    ϵ = fwd ? step_size(lf) : -step_size(lf)
-    ϵ = ϵ'
-
-    res = if FullTraj
-        Vector{P}(undef, n_steps)
-    else
-        z
-    end
+Build cache for gradient computation at position θ.
+Computes all θ-dependent quantities that can be reused across r values.
+"""
+function build_grad_cache(
+    h::Hamiltonian{<:AbstractRiemannianMetric}, θ::AbstractVector{T}
+) where {T}
+    # Evaluate log density and gradient
+    ℓπ, ∂ℓπ∂θ = h.∂ℓπ∂θ(θ)
 
-    for i in 1:n_steps
-        θ_init, r_init = z.θ, z.r
-        # Tempering
-        #r = temper(lf, r, (i=i, is_half=true), n_steps)
-        #! Eq (16) of Girolami & Calderhead (2011)
-        r_half = copy(r_init)
-        local cache
-        for j in 1:(lf.n)
-            # Reuse cache for the first iteration
-            if j == 1
-                (; value, gradient) = z.ℓπ
-            elseif j == 2 # cache intermediate values that depends on θ only (which are unchanged)
-                retval, cache = ∂H∂θ_cache(h, θ_init, r_half; return_cache=true)
-                (; value, gradient) = retval
-            else # reuse cache
-                (; value, gradient) = ∂H∂θ_cache(h, θ_init, r_half; cache=cache)
-            end
-            r_half = r_init - ϵ / 2 * gradient
-            # println("r_half: ", r_half)
-        end
-        #! Eq (17) of Girolami & Calderhead (2011)
-        θ_full = copy(θ_init)
-        term_1 = ∂H∂r(h, θ_init, r_half) # unchanged across the loop
-        for j in 1:(lf.n)
-            θ_full = θ_init + ϵ / 2 * (term_1 + ∂H∂r(h, θ_full, r_half))
-            # println("θ_full :", θ_full)
-        end
-        #! Eq (18) of Girolami & Calderhead (2011)
-        (; value, gradient) = ∂H∂θ(h, θ_full, r_half)
-        r_full = r_half - ϵ / 2 * gradient
-        # println("r_full: ", r_full)
-        # Tempering
-        #r = temper(lf, r, (i=i, is_half=false), n_steps)
-        # Create a new phase point by caching the logdensity and gradient
-        z = phasepoint(h, θ_full, r_full; ℓπ=DualValue(value, gradient))
-        # Update result
-        if FullTraj
-            res[i] = z
-        else
-            res = z
-        end
-        if !isfinite(z)
-            # Remove undef
-            if FullTraj
-                res = res[isassigned.(Ref(res), 1:n_steps)]
-            end
-            break
-        end
-        # @assert false
+    # Evaluate metric and sensitivities
+    G_eval = metric_eval(h.metric, θ)
+    ∂P∂θ = metric_sensitivity(h.metric, θ)
+
+    # Get logdet gradient matrix and precompute logdet gradient terms
+    M_logdet = logdet_grad_matrix(G_eval)
+    d = size(∂P∂θ, 3)
+    logdet_terms = Vector{T}(undef, d)
+    @inbounds for i in 1:d
+        ∂Pᵢ = @view ∂P∂θ[:, :, i]
+        logdet_terms[i] = T(0.5) * tr_product(M_logdet, ∂Pᵢ)
     end
-    return res
-end
 
-# TODO Make the order of θ and r consistent with neg_energy
-∂H∂θ(h::Hamiltonian, θ::AbstractVecOrMat, r::AbstractVecOrMat) = ∂H∂θ(h, θ)
-∂H∂r(h::Hamiltonian, θ::AbstractVecOrMat, r::AbstractVecOrMat) = ∂H∂r(h, r)
+    return RiemannianGradCache(G_eval, ∂P∂θ, ℓπ, ∂ℓπ∂θ, logdet_terms)
+end
 
-### hamiltonian.jl
+"""
+    ∂H∂θ_from_cache(cache::RiemannianGradCache, r)
 
-import AdvancedHMC: refresh, phasepoint
-using AdvancedHMC: FullMomentumRefreshment, PartialMomentumRefreshment, AbstractMetric
+Compute Hamiltonian gradient ∂H/∂θ using cached θ-dependent values.
+Only performs r-dependent computation (kinetic gradient matrix and trace products).
+"""
+function ∂H∂θ_from_cache(cache::RiemannianGradCache{T}, r::AbstractVector) where {T}
+    # Compute kinetic gradient matrix (r-dependent)
+    M_kinetic = kinetic_grad_matrix(cache.G_eval, r)
+
+    # Compute full gradient
+    d = length(cache.∂ℓπ∂θ)
+    grad = Vector{T}(undef, d)
+
+    @inbounds for i in 1:d
+        ∂Pᵢ = @view cache.∂P∂θ[:, :, i]
+        # ∂H/∂θᵢ = -∂ℓπ/∂θᵢ + 0.5*tr(M_logdet*∂P/∂θᵢ) - 0.5*tr(M_kinetic*∂P/∂θᵢ)
+        kinetic_term = T(0.5) * tr_product(M_kinetic, ∂Pᵢ)
+        grad[i] = -cache.∂ℓπ∂θ[i] + cache.logdet_terms[i] - kinetic_term
+    end
 
-# To change L180 of hamiltonian.jl
-function phasepoint(
-    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
-    θ::AbstractVecOrMat{T},
-    h::Hamiltonian,
-) where {T<:Real}
-    return phasepoint(h, θ, rand_momentum(rng, h.metric, h.kinetic, θ))
+    return DualValue(cache.ℓπ, grad)
 end
 
-# To change L191 of hamiltonian.jl
-function refresh(
-    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
-    ::FullMomentumRefreshment,
-    h::Hamiltonian,
-    z::PhasePoint,
+####
+#### Main gradient interface
+####
+
+"""
+    ∂H∂θ(h::Hamiltonian{<:AbstractRiemannianMetric}, θ, r)
+
+Compute the gradient of the Hamiltonian with respect to position θ.
+Returns a DualValue containing (log_density, gradient).
+
+Ref: Eq (15) of Girolami & Calderhead (2011)
+"""
+function ∂H∂θ(
+    h::Hamiltonian{<:AbstractRiemannianMetric,<:GaussianKinetic},
+    θ::AbstractVector,
+    r::AbstractVector,
 )
-    return phasepoint(h, z.θ, rand_momentum(rng, h.metric, h.kinetic, z.θ))
+    cache = build_grad_cache(h, θ)
+    return ∂H∂θ_from_cache(cache, r)
 end
 
-# To change L215 of hamiltonian.jl
-function refresh(
-    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
-    ref::PartialMomentumRefreshment,
-    h::Hamiltonian,
-    z::PhasePoint,
+"""
+    ∂H∂θ_cache(h, θ, r; cache=nothing)
+
+Compute ∂H/∂θ with optional caching for fixed-point iterations.
+Returns (DualValue, cache) tuple.
+
+When cache is provided, reuses θ-dependent computations (eigendecomposition,
+logdet gradient terms) and only recomputes r-dependent terms.
+"""
+function ∂H∂θ_cache(
+    h::Hamiltonian{<:AbstractRiemannianMetric,<:GaussianKinetic},
+    θ::AbstractVector,
+    r::AbstractVector;
+    cache=nothing,
 )
-    return phasepoint(
-        h,
-        z.θ,
-        ref.α * z.r + sqrt(1 - ref.α^2) * rand_momentum(rng, h.metric, h.kinetic, z.θ),
-    )
+    cache = @something cache build_grad_cache(h, θ)
+    return ∂H∂θ_from_cache(cache, r), cache
 end
 
-### metric.jl
+####
+#### Momentum gradient ∂H/∂r
+####
 
-import AdvancedHMC: _rand
-using AdvancedHMC: AbstractMetric
-using LinearAlgebra: eigen, cholesky, Symmetric
+"""
+    ∂H∂r(h::Hamiltonian{<:AbstractRiemannianMetric}, θ, r; G_eval=nothing)
 
-abstract type AbstractRiemannianMetric <: AbstractMetric end
+Compute the gradient of the Hamiltonian with respect to momentum r.
+For Riemannian metrics: ∂H/∂r = G(θ)⁻¹ * r
 
-abstract type AbstractHessianMap end
+If `G_eval` is provided, uses it directly instead of recomputing the metric.
 
-struct IdentityMap <: AbstractHessianMap end
+Ref: Eq (14) of Girolami & Calderhead (2011)
+"""
+function ∂H∂r(
+    h::Hamiltonian{<:AbstractRiemannianMetric,<:GaussianKinetic},
+    θ::AbstractVector,
+    r::AbstractVector;
+    G_eval=nothing,
+)
+    G = @something G_eval metric_eval(h.metric, θ)
+    return G \ r
+end
 
-(::IdentityMap)(x) = x
+####
+#### Negative energy (log probability)
+####
 
-struct SoftAbsMap{T} <: AbstractHessianMap
-    α::T
-end
+"""
+    neg_energy(h::Hamiltonian{<:AbstractRiemannianMetric}, r, θ; G_eval=nothing)
 
-# TODO Register softabs with ReverseDiff
-#! The definition of SoftAbs from Page 3 of Betancourt (2012)
-function softabs(X, α=20.0)
-    F = eigen(X) # ReverseDiff cannot diff through `eigen`
-    Q = hcat(F.vectors)
-    λ = F.values
-    softabsλ = λ .* coth.(α * λ)
-    return Q * diagm(softabsλ) * Q', Q, λ, softabsλ
-end
+Compute the negative kinetic energy for Riemannian metrics.
+Includes the log-determinant normalization term since G depends on θ.
 
-(map::SoftAbsMap)(x) = softabs(x, map.α)[1]
-
-struct DenseRiemannianMetric{
-    T,
-    TM<:AbstractHessianMap,
-    A<:Union{Tuple{Int},Tuple{Int,Int}},
-    AV<:AbstractVecOrMat{T},
-    TG,
-    T∂G∂θ,
-} <: AbstractRiemannianMetric
-    size::A
-    G::TG # TODO store G⁻¹ here instead
-    ∂G∂θ::T∂G∂θ
-    map::TM
-    _temp::AV
-end
+If `G_eval` is provided, uses it directly instead of recomputing the metric.
 
-# TODO Make dense mass matrix support matrix-mode parallel
-function DenseRiemannianMetric(size, G, ∂G∂θ, map=IdentityMap()) where {T<:AbstractFloat}
-    _temp = Vector{Float64}(undef, size[1])
-    return DenseRiemannianMetric(size, G, ∂G∂θ, map, _temp)
-end
-# DenseEuclideanMetric(::Type{T}, D::Int) where {T} = DenseEuclideanMetric(Matrix{T}(I, D, D))
-# DenseEuclideanMetric(D::Int) = DenseEuclideanMetric(Float64, D)
-# DenseEuclideanMetric(::Type{T}, sz::Tuple{Int}) where {T} = DenseEuclideanMetric(Matrix{T}(I, first(sz), first(sz)))
-# DenseEuclideanMetric(sz::Tuple{Int}) = DenseEuclideanMetric(Float64, sz)
+K(r, θ) = 0.5 * (D*log(2π) + log|G(θ)| + r'G(θ)⁻¹r)
+neg_energy = -K = -0.5 * (D*log(2π) + log|G(θ)| + r'G(θ)⁻¹r)
+
+Ref: Eq (13) of Girolami & Calderhead (2011)
+"""
+function neg_energy(
+    h::Hamiltonian{<:AbstractRiemannianMetric,<:GaussianKinetic},
+    r::AbstractVector,
+    θ::AbstractVector;
+    G_eval=nothing,
+)
+    G = @something G_eval metric_eval(h.metric, θ)
+    D = length(r)
 
-# renew(ue::DenseEuclideanMetric, M⁻¹) = DenseEuclideanMetric(M⁻¹)
+    # Quadratic form: r' * G⁻¹ * r
+    G_inv_r = G \ r
+    quadform = dot(r, G_inv_r)
 
-Base.size(e::DenseRiemannianMetric) = e.size
-Base.size(e::DenseRiemannianMetric, dim::Int) = e.size[dim]
-Base.show(io::IO, dem::DenseRiemannianMetric) = print(io, "DenseRiemannianMetric(...)")
+    # Log normalization constant (position-dependent)
+    logZ = (D * log(2π) + logdet(G)) / 2
 
-function rand_momentum(
-    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
-    metric::DenseRiemannianMetric{T},
-    kinetic,
-    θ::AbstractVecOrMat,
-) where {T}
-    r = _randn(rng, T, size(metric)...)
-    G⁻¹ = inv(metric.map(metric.G(θ)))
-    chol = cholesky(Symmetric(G⁻¹))
-    ldiv!(chol.U, r)
-    return r
+    return -logZ - quadform / 2
 end
 
-### hamiltonian.jl
-
-import AdvancedHMC: phasepoint, neg_energy, ∂H∂θ, ∂H∂r
-using LinearAlgebra: logabsdet, tr
+####
+#### Phase point construction
+####
 
-# QUES Do we want to change everything to position dependent by default?
-# Add θ to ∂H∂r for DenseRiemannianMetric
+"""
+Create a PhasePoint for Riemannian metrics, computing position-dependent kinetic energy.
+Shares the metric evaluation between neg_energy and ∂H∂r to avoid redundant computation.
+"""
 function phasepoint(
-    h::Hamiltonian{<:DenseRiemannianMetric},
+    h::Hamiltonian{<:AbstractRiemannianMetric},
     θ::T,
     r::T;
-    ℓπ=∂H∂θ(h, θ),
-    ℓκ=DualValue(neg_energy(h, r, θ), ∂H∂r(h, θ, r)),
-) where {T<:AbstractVecOrMat}
-    return PhasePoint(θ, r, ℓπ, ℓκ)
-end
-
-# Negative kinetic energy
-#! Eq (13) of Girolami & Calderhead (2011)
-function neg_energy(
-    h::Hamiltonian{<:DenseRiemannianMetric}, r::T, θ::T
+    ℓπ=∂H∂θ(h, θ, r),
+    G_eval=nothing,
+    ℓκ=nothing,
 ) where {T<:AbstractVecOrMat}
-    G = h.metric.map(h.metric.G(θ))
-    D = size(G, 1)
-    # Need to consider the normalizing term as it is no longer same for different θs
-    logZ = 1 / 2 * (D * log(2π) + logdet(G)) # it will be user's responsibility to make sure G is SPD and logdet(G) is defined
-    mul!(h.metric._temp, inv(G), r)
-    return -logZ - dot(r, h.metric._temp) / 2
-end
-
-# QUES L31 of hamiltonian.jl now reads a bit weird (semantically)
-function ∂H∂θ(
-    h::Hamiltonian{<:DenseRiemannianMetric{T,<:IdentityMap}},
-    θ::AbstractVecOrMat{T},
-    r::AbstractVecOrMat{T},
-) where {T}
-    ℓπ, ∂ℓπ∂θ = h.∂ℓπ∂θ(θ)
-    G = h.metric.map(h.metric.G(θ))
-    invG = inv(G)
-    ∂G∂θ = h.metric.∂G∂θ(θ)
-    d = length(∂ℓπ∂θ)
-    return DualValue(
-        ℓπ,
-        #! Eq (15) of Girolami & Calderhead (2011)
-        -mapreduce(vcat, 1:d) do i
-            ∂G∂θᵢ = ∂G∂θ[:, :, i]
-            ∂ℓπ∂θ[i] - 1 / 2 * tr(invG * ∂G∂θᵢ) + 1 / 2 * r' * invG * ∂G∂θᵢ * invG * r
-            # Gr = G \ r
-            # ∂ℓπ∂θ[i] - 1 / 2 * tr(G \ ∂G∂θᵢ) + 1 / 2 * Gr' * ∂G∂θᵢ * Gr
-            # 1 / 2 * tr(invG * ∂G∂θᵢ)
-            # 1 / 2 * r' * invG * ∂G∂θᵢ * invG * r
-        end,
-    )
-end
-
-# Ref: https://www.wolframalpha.com/input?i=derivative+of+x+*+coth%28a+*+x%29
-#! Based on middle of the right column of Page 3 of Betancourt (2012) "Note that whenλi=λj, such as for the diagonal elementsor degenerate eigenvalues, this becomes the derivative"
-dsoftabsdλ(α, λ) = coth(α * λ) + λ * α * -csch(λ * α)^2
-
-#! J as defined in middle of the right column of Page 3 of Betancourt (2012)
-function make_J(λ::AbstractVector{T}, α::T) where {T<:AbstractFloat}
-    d = length(λ)
-    J = Matrix{T}(undef, d, d)
-    for i in 1:d, j in 1:d
-        J[i, j] = if (λ[i] == λ[j])
-            dsoftabsdλ(α, λ[i])
-        else
-            ((λ[i] * coth(α * λ[i]) - λ[j] * coth(α * λ[j])) / (λ[i] - λ[j]))
-        end
+    if isnothing(ℓκ)
+        # Compute G_eval once and share between neg_energy and ∂H∂r
+        G = @something G_eval metric_eval(h.metric, θ)
+        ℓκ = DualValue(neg_energy(h, r, θ; G_eval=G), ∂H∂r(h, θ, r; G_eval=G))
     end
-    return J
-end
-
-function ∂H∂θ(
-    h::Hamiltonian{<:DenseRiemannianMetric{T,<:SoftAbsMap}},
-    θ::AbstractVecOrMat{T},
-    r::AbstractVecOrMat{T},
-) where {T}
-    return ∂H∂θ_cache(h, θ, r)
+    return PhasePoint(θ, r, ℓπ, ℓκ)
 end
-function ∂H∂θ_cache(
-    h::Hamiltonian{<:DenseRiemannianMetric{T,<:SoftAbsMap}},
-    θ::AbstractVecOrMat{T},
-    r::AbstractVecOrMat{T};
-    return_cache=false,
-    cache=nothing,
-) where {T}
-    # Terms that only dependent on θ can be cached in θ-unchanged loops
-    if isnothing(cache)
-        ℓπ, ∂ℓπ∂θ = h.∂ℓπ∂θ(θ)
-        H = h.metric.G(θ)
-        ∂H∂θ = h.metric.∂G∂θ(θ)
-
-        G, Q, λ, softabsλ = softabs(H, h.metric.map.α)
 
-        R = diagm(1 ./ softabsλ)
+####
+#### Momentum refreshment
+####
 
-        # softabsΛ = diagm(softabsλ)
-        # M = inv(softabsΛ) * Q' * r
-        # M = R * Q' * r # equiv to above but avoid inv
-
-        J = make_J(λ, h.metric.map.α)
-
-        #! Based on the two equations from the right column of Page 3 of Betancourt (2012)
-        term_1_cached = Q * (R .* J) * Q'
-    else
-        ℓπ, ∂ℓπ∂θ, ∂H∂θ, Q, softabsλ, J, term_1_cached = cache
-    end
-    d = length(∂ℓπ∂θ)
-    D = diagm((Q' * r) ./ softabsλ)
-    term_2_cached = Q * D * J * D * Q'
-    g =
-        -mapreduce(vcat, 1:d) do i
-            ∂H∂θᵢ = ∂H∂θ[:, :, i]
-            # ∂ℓπ∂θ[i] - 1 / 2 * tr(term_1_cached * ∂H∂θᵢ) + 1 / 2 * M' * (J .* (Q' * ∂H∂θᵢ * Q)) * M # (v1)
-            # NOTE Some further optimization can be done here: cache the 1st product all together
-            ∂ℓπ∂θ[i] - 1 / 2 * tr(term_1_cached * ∂H∂θᵢ) + 1 / 2 * tr(term_2_cached * ∂H∂θᵢ) # (v2) cache friendly
-        end
-
-    dv = DualValue(ℓπ, g)
-    return return_cache ? (dv, (; ℓπ, ∂ℓπ∂θ, ∂H∂θ, Q, softabsλ, J, term_1_cached)) : dv
-end
-
-#! Eq (14) of Girolami & Calderhead (2011)
-function ∂H∂r(
-    h::Hamiltonian{<:DenseRiemannianMetric}, θ::AbstractVecOrMat, r::AbstractVecOrMat
+# PartialMomentumRefreshment + Riemannian is mathematically well-defined at stationarity 
+# (both terms in α·r + √(1-α²)·n live at the same θ), but we were unable to verify the
+# validity of the sampler step emprically. This may simply be due to poor performance of the
+# sampler, but to be safe we are marking this as untested for now.
+function refresh(
+    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
+    ref::PartialMomentumRefreshment,
+    h::Hamiltonian{<:AbstractRiemannianMetric},
+    z::PhasePoint,
 )
-    H = h.metric.G(θ)
-    # if !all(isfinite, H)
-    #     println("θ: ", θ)
-    #     println("H: ", H)
-    # end
-    G = h.metric.map(H)
-    # return inv(G) * r
-    # println("G \ r: ", G \ r)
-    return G \ r # NOTE it's actually pretty weird that ∂H∂θ returns DualValue but ∂H∂r doesn't
+    @warn (
+        "PartialMomentumRefreshment with Riemannian metrics is untested and may not " *
+        "target the correct posterior. Prefer FullMomentumRefreshment unless you have " *
+        "validated convergence for your model."
+    ) maxlog = 1
+    return @invoke refresh(
+        rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
+        ref::PartialMomentumRefreshment,
+        h::Hamiltonian,
+        z::PhasePoint,
+    )
 end
diff --git a/src/riemannian/integrator.jl b/src/riemannian/integrator.jl
index a6d2de5f..27729489 100644
--- a/src/riemannian/integrator.jl
+++ b/src/riemannian/integrator.jl
@@ -15,18 +15,14 @@ $(TYPEDFIELDS)
 struct GeneralizedLeapfrog{T<:AbstractScalarOrVec{<:AbstractFloat}} <: AbstractLeapfrog{T}
     "Step size."
     ϵ::T
+    "Number of fixed-point iterations for implicit steps."
     n::Int
 end
+
 function Base.show(io::IO, l::GeneralizedLeapfrog)
     return print(io, "GeneralizedLeapfrog(ϵ=", round.(l.ϵ; sigdigits=3), ", n=", l.n, ")")
 end
 
-# fallback to ignore return_cache & cache kwargs for other ∂H∂θ
-function ∂H∂θ_cache(h, θ, r; return_cache=false, cache=nothing)
-    dv = ∂H∂θ(h, θ, r)
-    return return_cache ? (dv, nothing) : dv
-end
-
 # TODO(Kai) make sure vectorization works
 # TODO(Kai) check if tempering is valid
 # TODO(Kai) abstract out the 3 main steps and merge with `step` in `integrator.jl`
@@ -55,42 +51,47 @@ function step(
 
     for i in 1:n_steps
         θ_init, r_init = z.θ, z.r
-        # Tempering
-        #r = temper(lf, r, (i=i, is_half=true), n_steps)
-        # eq (16) of Girolami & Calderhead (2011)
+
+        # Cache θ-dependent quantities at θ_init once per step; reused across the Eq (16)
+        # fixed-point iterations and again for Eq (17) term_1.
+        cache = build_grad_cache(h, θ_init)
+
+        # Eq (16) of Girolami & Calderhead (2011) - implicit momentum half-step
         r_half = r_init
-        local cache = nothing
-        for j in 1:(lf.n)
-            # Reuse cache for the first iteration
-            if j == 1
-                (; value, gradient) = z.ℓπ
-            elseif j == 2 # cache intermediate values that depends on θ only (which are unchanged)
-                retval, cache = ∂H∂θ_cache(h, θ_init, r_half; return_cache=true)
-                (; value, gradient) = retval
-            else # reuse cache
-                (; value, gradient) = ∂H∂θ_cache(h, θ_init, r_half; cache=cache)
-            end
+        for _ in 1:(lf.n)
+            gradient = ∂H∂θ_from_cache(cache, r_half).gradient
             r_half = r_init - ϵ / 2 * gradient
         end
-        # eq (17) of Girolami & Calderhead (2011)
+
+        # Eq (17) of Girolami & Calderhead (2011) - implicit position step
+        term_1 = ∂H∂r(h, θ_init, r_half; G_eval=cache.G_eval)
         θ_full = θ_init
-        term_1 = ∂H∂r(h, θ_init, r_half) # unchanged across the loop
         for j in 1:(lf.n)
             θ_full = θ_init + ϵ / 2 * (term_1 + ∂H∂r(h, θ_full, r_half))
         end
-        # eq (18) of Girolami & Calderhead (2011)
-        (; value, gradient) = ∂H∂θ(h, θ_full, r_half)
+
+        # Eq (18) of Girolami & Calderhead (2011) - explicit momentum half-step
+        # Use the cached G_eval at θ_full to avoid a redundant metric_eval in phasepoint
+        dv, cache = ∂H∂θ_cache(h, θ_full, r_half)
+        (; value, gradient) = dv
         r_full = r_half - ϵ / 2 * gradient
-        # Tempering
-        #r = temper(lf, r, (i=i, is_half=false), n_steps)
-        # Create a new phase point by caching the logdensity and gradient
-        z = phasepoint(h, θ_full, r_full; ℓπ=DualValue(value, gradient))
+
+        # Create a new phase point by caching the logdensity and gradient.
+        # The gradient stored in ℓπ here is `∂H/∂θ(θ_full, r_half)`, not the field's
+        # documented `-∂ℓπ/∂θ(θ_full)` contract, and is evaluated at `r_half` rather
+        # than `r_full`. No live consumer reads it on the Riemannian path; left as a
+        # placeholder pending a follow-up to clean up the `ℓπ`/`neg_energy` naming.
+        z = phasepoint(
+            h, θ_full, r_full; ℓπ=DualValue(value, gradient), G_eval=cache.G_eval
+        )
+
         # Update result
         if FullTraj
             res[i] = z
         else
             res = z
         end
+
         if !isfinite(z)
             # Remove undef
             if FullTraj
diff --git a/src/riemannian/metric.jl b/src/riemannian/metric.jl
new file mode 100644
index 00000000..fff6a456
--- /dev/null
+++ b/src/riemannian/metric.jl
@@ -0,0 +1,373 @@
+using LinearAlgebra: LinearAlgebra
+
+####
+#### Riemannian Metric Types
+####
+
+"""
+Abstract type for Riemannian (position-dependent) metrics.
+
+Subtypes must implement:
+- `metric_eval(metric, θ)` - evaluate metric at position θ
+- `metric_sensitivity(metric, θ)` - compute ∂P/∂θ where P is the "pre-metric"
+  (G itself for RiemannianMetric, or H the Hessian for SoftAbsRiemannianMetric)
+"""
+abstract type AbstractRiemannianMetric <: AbstractMetric end
+
+####
+#### SoftAbsEval - cached eigendecomposition for SoftAbs metrics
+####
+
+"""
+    SoftAbsEval{T}
+
+Cached result of evaluating a SoftAbs metric at a position θ.
+Stores eigendecomposition and precomputed matrices for efficient gradient computation.
+
+# Fields
+- `Q`: Eigenvectors (orthogonal matrix)
+- `softabsλ`: Transformed eigenvalues: λᵢ * coth(α * λᵢ)
+- `J`: Jacobian matrix encoding the derivative of softabs (divided difference formula)
+- `M_logdet`: Precomputed matrix Q * (R .* J) * Q' for logdet gradient
+"""
+struct SoftAbsEval{T<:AbstractFloat}
+    Q::Matrix{T}
+    softabsλ::Vector{T}
+    J::Matrix{T}
+    M_logdet::Matrix{T}
+end
+
+# Standard operations for SoftAbsEval
+function Base.:\(G::SoftAbsEval, p::AbstractVector)
+    return G.Q * ((G.Q' * p) ./ G.softabsλ)
+end
+
+function LinearAlgebra.logdet(G::SoftAbsEval)
+    return sum(log, G.softabsλ)
+end
+
+"""
+    unwhiten(G::SoftAbsEval, z)
+
+Transform z ~ N(0, I) to sample from N(0, G).
+"""
+function unwhiten(G::SoftAbsEval, z::AbstractVector)
+    return G.Q * (sqrt.(G.softabsλ) .* z)
+end
+
+####
+#### RiemannianMetric - for user-provided PD metrics
+####
+
+"""
+    RiemannianMetric{T, TG, T∂G}
+
+Riemannian metric where the user provides a function returning a positive-definite
+matrix (or AbstractPDMat subtype). `T` is the element type used for momentum sampling
+and defaults to `$(DEFAULT_FLOAT_TYPE)` if not specified.
+
+For best performance, return an `AbstractPDMat` so that the kinetic and logdet gradients can
+reuse the stored Cholesky factor instead of refactorising on every call.
+
+If you need to stabilise `G` (e.g. by adding `λI`), do so *inside* the `PDMat` constructor 
+since `PDMat(G) + λ*I` and other variants may fall back to returning Symmetric or even dense
+matrices.
+
+# Fields
+- `size`: Tuple{Int} giving the dimension
+- `calc_G`: Function θ → G(θ), returns a positive-definite matrix
+- `calc_∂G∂θ`: Function θ → ∂G/∂θ, returns Array{T,3} of shape (d, d, d)
+
+# Example
+```julia
+# Simple Fisher information metric
+calc_G = θ -> PDMat(fisher_information(θ))
+calc_∂G∂θ = θ -> ForwardDiff.jacobian(θ -> vec(fisher_information(θ)), θ) |> reshape_∂G∂θ
+metric = RiemannianMetric((d,), calc_G, calc_∂G∂θ)  # T = Float64
+metric_f32 = RiemannianMetric{Float32}((d,), calc_G, calc_∂G∂θ)
+```
+"""
+struct RiemannianMetric{T<:AbstractFloat,TG,T∂G} <: AbstractRiemannianMetric
+    size::Tuple{Int}
+    calc_G::TG       # θ → Matrix or AbstractPDMat
+    calc_∂G∂θ::T∂G   # θ → Array{T,3}
+end
+
+function RiemannianMetric{T}(
+    size::Tuple{Int}, calc_G::TG, calc_∂G∂θ::T∂G
+) where {T<:AbstractFloat,TG,T∂G}
+    return RiemannianMetric{T,TG,T∂G}(size, calc_G, calc_∂G∂θ)
+end
+
+function RiemannianMetric(size::Tuple{Int}, calc_G, calc_∂G∂θ)
+    return RiemannianMetric{DEFAULT_FLOAT_TYPE}(size, calc_G, calc_∂G∂θ)
+end
+
+Base.size(m::RiemannianMetric) = m.size
+Base.size(m::RiemannianMetric, dim::Int) = m.size[dim]
+Base.eltype(::RiemannianMetric{T}) where {T} = T
+
+function Base.show(io::IO, m::RiemannianMetric)
+    return print(io, "RiemannianMetric(size=", m.size, ")")
+end
+
+# Interface implementations for RiemannianMetric
+metric_eval(m::RiemannianMetric, θ) = m.calc_G(θ)
+metric_sensitivity(m::RiemannianMetric, θ) = m.calc_∂G∂θ(θ)
+
+####
+#### SoftAbsRiemannianMetric - for Hessian-based metrics with SoftAbs regularization
+####
+
+"""
+    SoftAbsRiemannianMetric{T, TH, T∂H}
+
+Riemannian metric based on the SoftAbs transformation of a Hessian.
+The Hessian may not be positive-definite; the SoftAbs transformation
+G = Q * diag(λ * coth(α*λ)) * Q' guarantees positive-definiteness.
+
+# Fields
+- `size`: Tuple{Int} giving the dimension
+- `calc_H`: Function θ → H(θ), returns the Hessian matrix (the "pre-metric")
+- `calc_∂H∂θ`: Function θ → ∂H/∂θ, returns Array{T,3} of shape (d, d, d)
+- `α`: SoftAbs regularization parameter (larger = closer to |λ|)
+- `canonicalize`: If `true`, eigenvector signs are fixed after each `eigen` call so that
+  the largest-magnitude element of each column is positive. This makes momentum sampling
+  via `unwhiten` reproducible across BLAS/LAPACK implementations that may return the same
+  eigenvectors with opposite sign conventions (e.g. OpenBLAS on x86 vs arm64). Has no
+  effect on the metric matrix G or any gradient — those are sign-invariant. Default `false`.
+
+# References
+- Betancourt, M. "A general metric for Riemannian manifold Hamiltonian Monte Carlo" (2012)
+"""
+struct SoftAbsRiemannianMetric{T<:AbstractFloat,TH,T∂H} <: AbstractRiemannianMetric
+    size::Tuple{Int}
+    calc_H::TH       # θ → Hessian matrix (pre-metric)
+    calc_∂H∂θ::T∂H   # θ → Array{T,3}
+    α::T
+    canonicalize::Bool
+end
+
+function SoftAbsRiemannianMetric(
+    size::Tuple{Int}, calc_H::TH, calc_∂H∂θ::T∂H, α::T; canonicalize::Bool=false
+) where {T<:AbstractFloat,TH,T∂H}
+    return SoftAbsRiemannianMetric{T,TH,T∂H}(size, calc_H, calc_∂H∂θ, α, canonicalize)
+end
+
+Base.size(m::SoftAbsRiemannianMetric) = m.size
+Base.size(m::SoftAbsRiemannianMetric, dim::Int) = m.size[dim]
+Base.eltype(::SoftAbsRiemannianMetric{T}) where {T} = T
+
+function Base.show(io::IO, m::SoftAbsRiemannianMetric)
+    print(io, "SoftAbsRiemannianMetric(size=", m.size, ", α=", m.α)
+    m.canonicalize && print(io, ", canonicalize=true")
+    return print(io, ")")
+end
+
+# Compile-time cutoffs for the SoftAbs stability switches. `@generated` ensures
+# `eps(T)^(1//n)` is evaluated once per specialisation, not on every call.
+@generated _xcothx_cutoff(::Type{T}) where {T<:AbstractFloat} = eps(T)^(1//6)
+@generated _make_J_cutoff(::Type{T}) where {T<:AbstractFloat} = eps(T)^(1//3)
+
+# Branch implementations of x·coth(x). _exact is correct away from zero; _taylor
+# is correct near zero. _xcothx dispatches between them at _xcothx_cutoff(T).
+@inline _xcothx_taylor(x::T) where {T<:AbstractFloat} = one(T) + x^2 / 3 - x^4 / 45
+@inline _xcothx_exact(x::T) where {T<:AbstractFloat} = x * coth(x)
+
+"""
+    _xcothx(x)
+
+Compute `x * coth(x)` (i.e. `α * softabs(λ)` for `x = α*λ`) in a way that is
+numerically stable as `x → 0`.
+
+Naively, `x * coth(x)` evaluates to `0 * Inf = NaN` at exactly `x = 0`, even though
+the true limit is `1`. We switch to the two-term Taylor expansion
+`1 + x²/3 − x⁴/45` below `|x| < eps(T)^(1//6)`, whose truncation error
+(`O(x⁶) ≈ eps`) is at machine precision by the time we hit the switch.
+"""
+function _xcothx(x::T) where {T<:AbstractFloat}
+    return abs(x) < _xcothx_cutoff(T) ? _xcothx_taylor(x) : _xcothx_exact(x)
+end
+
+# Branch implementations of d/dx[x·coth(x)] = coth(x) − x·csch(x)².
+@inline _xcothx_deriv_taylor(x::T) where {T<:AbstractFloat} = T(2) / 3 * x - T(4) / 45 * x^3
+@inline _xcothx_deriv_exact(x::T) where {T<:AbstractFloat} = coth(x) - x * csch(x)^2
+
+"""
+    _xcothx_deriv(x)
+
+Compute the derivative `coth(x) − x * csch(x)²` of `x * coth(x)`, stably as `x → 0`.
+
+Naively, both `coth(x)` and `x * csch(x)²` behave like `1/x` near zero, so direct
+subtraction suffers catastrophic cancellation (relative error `~eps/x²`); at `x = 0`
+the result is `Inf − Inf = NaN`. The Taylor expansion is `2x/3 − 4x³/45 + O(x⁵)`.
+Balancing Taylor truncation (`~x⁴` relative) against cancellation (`~eps/x²` relative)
+gives the optimal switch at `|x| < eps(T)^(1//6)`, yielding ~13 digits across the range.
+"""
+function _xcothx_deriv(x::T) where {T<:AbstractFloat}
+    return abs(x) < _xcothx_cutoff(T) ? _xcothx_deriv_taylor(x) : _xcothx_deriv_exact(x)
+end
+
+"""
+    make_J(λ, α)
+
+Construct the J matrix for softabs gradient computation.
+J encodes the derivative of the softabs transformation using the divided difference formula.
+
+For `λᵢ` well separated from `λⱼ`:
+    J[i,j] = (softabs(λᵢ) − softabs(λⱼ)) / (λᵢ − λⱼ)
+For `λᵢ ≈ λⱼ` (including the diagonal):
+    J[i,j] = d/dλ [λ coth(αλ)] |_{λ = (λᵢ + λⱼ)/2}
+
+The branches are switched when `|α(λᵢ − λⱼ)| < eps(T)^(1//3)`: at that point the
+divided-difference cancellation (`~eps/(αδ)`) and the midpoint-rule truncation
+(`~(αδ)²`) balance, each contributing `~eps^(2/3)` relative error.
+
+# References
+- Betancourt (2012)
+"""
+function make_J(λ::AbstractVector{T}, α::T) where {T<:AbstractFloat}
+    d = length(λ)
+    J = Matrix{T}(undef, d, d)
+    deg_tol = _make_J_cutoff(T)
+    @inbounds for j in 1:d, i in 1:d
+        xi = α * λ[i]
+        xj = α * λ[j]
+        if abs(xi - xj) < deg_tol
+            # Diagonal or near-degenerate: derivative at the midpoint (stable at x = 0).
+            J[i, j] = _xcothx_deriv((xi + xj) / 2)
+        else
+            # Divided difference written in terms of α·λ so _xcothx handles λ = 0 safely.
+            J[i, j] = (_xcothx(xi) - _xcothx(xj)) / (xi - xj)
+        end
+    end
+    return J
+end
+
+"""
+    metric_eval(m::SoftAbsRiemannianMetric, θ)
+
+Evaluate SoftAbs metric at position θ, returning a `SoftAbsEval` with cached matrices.
+"""
+function metric_eval(m::SoftAbsRiemannianMetric{T}, θ) where {T}
+    H = m.calc_H(θ)
+    F = eigen(Symmetric(H))
+    λ = F.values
+    Q = F.vectors
+
+    if m.canonicalize
+        @inbounds for j in axes(Q, 2)
+            col = view(Q, :, j)
+            if col[argmax(abs.(col))] < 0
+                col .*= -1
+            end
+        end
+    end
+
+    # SoftAbs transformation: G = Q * diag(softabsλ) * Q'.
+    # Use _xcothx to avoid `0 * Inf = NaN` at exactly λ = 0 (limit is 1/α).
+    softabsλ = _xcothx.(m.α .* λ) ./ m.α
+
+    # Compute J matrix for gradient chain rule
+    J = make_J(λ, m.α)
+
+    # Precompute M_logdet = Q * (R .* J) * Q' where R = diag(1 ./ softabsλ)
+    # This is used for: ∂log|G|/∂θᵢ = 0.5 * tr(M_logdet * ∂H/∂θᵢ)
+    R = Diagonal(one(T) ./ softabsλ)
+    M_logdet = Q * (R .* J) * Q'
+
+    return SoftAbsEval(Q, softabsλ, J, M_logdet)
+end
+
+metric_sensitivity(m::SoftAbsRiemannianMetric, θ) = m.calc_∂H∂θ(θ)
+
+####
+#### Gradient matrices for unified computation
+####
+
+"""
+    logdet_grad_matrix(G)
+
+Return the matrix M such that ∂log|G|/∂θᵢ = 0.5 * tr(M * ∂P/∂θᵢ), where P is the
+"pre-metric" (G itself for RiemannianMetric, or H the Hessian for SoftAbsRiemannianMetric).
+
+For dense matrices: M = G⁻¹
+For SoftAbsEval: M = Q * (R .* J) * Q' (precomputed in metric_eval)
+
+The J matrix in SoftAbsEval absorbs the chain rule through the softabs transformation,
+so the same formula works with ∂H/∂θ instead of ∂G/∂θ.
+"""
+logdet_grad_matrix(G::SoftAbsEval) = G.M_logdet
+logdet_grad_matrix(G::AbstractMatrix) = inv(G)
+
+"""
+    kinetic_grad_matrix(G, r)
+
+Return the matrix M such that ∂(r'G⁻¹r)/∂θᵢ = -tr(M * ∂P/∂θᵢ), where P is the
+"pre-metric" (G itself for RiemannianMetric, or H the Hessian for SoftAbsRiemannianMetric).
+
+For dense matrices: M = (G⁻¹r)(G⁻¹r)' (rank-1 outer product)
+For SoftAbsEval: M = Q * D * J * D * Q' where D = diag((Q'r) ./ softabsλ)
+
+For SoftAbsEval, the J matrix absorbs the chain rule through softabs, allowing
+the gradient to be computed with respect to ∂H/∂θ rather than ∂G/∂θ. This avoids
+recomputing J for each value of r during fixed-point iterations.
+"""
+function kinetic_grad_matrix(G::SoftAbsEval, r::AbstractVector)
+    # D = diag((Q'r) ./ softabsλ)
+    d = (G.Q' * r) ./ G.softabsλ
+    D = Diagonal(d)
+    return G.Q * D * G.J * D * G.Q'
+end
+
+function kinetic_grad_matrix(G::AbstractMatrix, r::AbstractVector)
+    v = G \ r
+    return v * v'  # Rank-1 outer product
+end
+
+####
+#### Momentum sampling
+####
+
+function rand_momentum(
+    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
+    metric::AbstractRiemannianMetric,
+    ::GaussianKinetic,
+    θ::AbstractVecOrMat,
+)
+    G = metric_eval(metric, θ)
+    z = _randn(rng, eltype(metric), size(metric)...)
+    return unwhiten(G, z)
+end
+
+# unwhiten for regular matrices (PDMat or dense)
+function unwhiten(G::AbstractMatrix, z::AbstractVector)
+    # G = L * L', so sample = L * z where L = chol(G).L
+    chol = cholesky(Symmetric(G))
+    return chol.L * z
+end
+
+####
+#### Deprecated API forwards: DenseRiemannianMetric → {RiemannianMetric, SoftAbsRiemannianMetric}
+####
+#### `IdentityMap` and `SoftAbsMap` are retained as minimal tag types so the old
+#### `DenseRiemannianMetric(size, G, ∂G∂θ, map)` signature still parses; they have no
+#### runtime role beyond dispatching the @deprecate forwards below.
+
+struct IdentityMap end
+
+struct SoftAbsMap{T}
+    α::T
+end
+
+@deprecate(
+    DenseRiemannianMetric(size::Tuple{Int}, G, ∂G∂θ), RiemannianMetric(size, G, ∂G∂θ),
+)
+@deprecate(
+    DenseRiemannianMetric(size::Tuple{Int}, G, ∂G∂θ, ::IdentityMap),
+    RiemannianMetric(size, G, ∂G∂θ),
+)
+@deprecate(
+    DenseRiemannianMetric(size::Tuple{Int}, G, ∂G∂θ, map::SoftAbsMap),
+    SoftAbsRiemannianMetric(size, G, ∂G∂θ, map.α),
+)
diff --git a/src/sampler.jl b/src/sampler.jl
index 1b282383..2aab6c65 100644
--- a/src/sampler.jl
+++ b/src/sampler.jl
@@ -1,22 +1,26 @@
 # Update of hamiltonian and proposal
 
+const MassMatrixAdaptors = Union{MassMatrixAdaptor,NaiveHMCAdaptor,StanHMCAdaptor}
+const StepSizeAdaptors = Union{StepSizeAdaptor,NaiveHMCAdaptor,StanHMCAdaptor}
+
 update(h::Hamiltonian, ::AbstractAdaptor) = h
-function update(
-    h::Hamiltonian, adaptor::Union{MassMatrixAdaptor,NaiveHMCAdaptor,StanHMCAdaptor}
-)
+function update(h::Hamiltonian, adaptor::MassMatrixAdaptors)
     metric = renew(h.metric, getM⁻¹(adaptor))
     return @set h.metric = metric
 end
 
 update(τ::Trajectory, ::AbstractAdaptor) = τ
-function update(
-    τ::Trajectory, adaptor::Union{StepSizeAdaptor,NaiveHMCAdaptor,StanHMCAdaptor}
-)
+function update(τ::Trajectory, adaptor::StepSizeAdaptors)
     # FIXME: this does not support change type of `ϵ` (e.g. Float to Vector)
     integrator = update_nom_step_size(τ.integrator, getϵ(adaptor))
     @set τ.integrator = integrator
 end
 
+# Error clearly if mass-matrix adaptation is attempted with a Riemannian metric
+function update(::Hamiltonian{<:AbstractRiemannianMetric}, ::MassMatrixAdaptors)
+    return error("Mass-matrix adaptation is not supported/required for Riemannian metrics")
+end
+
 function update(κ::AbstractMCMCKernel, adaptor::AbstractAdaptor)
     @set κ.τ = update(κ.τ, adaptor)
 end
diff --git a/src/trajectory.jl b/src/trajectory.jl
index 66246e74..93006ea4 100644
--- a/src/trajectory.jl
+++ b/src/trajectory.jl
@@ -141,8 +141,9 @@ $(TYPEDEF)
 Slice sampler for the starting single leaf tree.
 Slice variable is initialized.
 """
-SliceTS(rng::AbstractRNG, z0::PhasePoint) =
-    SliceTS(z0, neg_energy(z0) - Random.randexp(rng), 1)
+function SliceTS(rng::AbstractRNG, z0::PhasePoint)
+    return SliceTS(z0, neg_energy(z0) - Random.randexp(rng), 1)
+end
 
 """
 $(TYPEDEF)
@@ -292,7 +293,7 @@ function transition(
             hamiltonian_energy=H,
             hamiltonian_energy_error=H - H0,
             # check numerical error in proposed phase point. 
-            numerical_error=!all(isfinite, H′),
+            numerical_error=(!all(isfinite, H′)),
         ),
         stat(τ.integrator),
     )
@@ -552,7 +553,7 @@ function isterminated(::ClassicNoUTurn, h::Hamiltonian, t::BinaryTree)
     # z0 is starting point and z1 is ending point
     z0, z1 = t.zleft, t.zright
     Δθ = z1.θ - z0.θ
-    s = (dot(Δθ, ∂H∂r(h, -z0.r)) >= 0) || (dot(-Δθ, ∂H∂r(h, z1.r)) >= 0)
+    s = (dot(Δθ, ∂H∂r(h, z0.θ, -z0.r)) >= 0) || (dot(-Δθ, ∂H∂r(h, z1.θ, z1.r)) >= 0)
     return Termination(s, false)
 end
 
@@ -565,7 +566,9 @@ Ref: https://arxiv.org/abs/1701.02434
 """
 function isterminated(::GeneralisedNoUTurn, h::Hamiltonian, t::BinaryTree)
     rho = t.ts.rho
-    s = generalised_uturn_criterion(rho, ∂H∂r(h, t.zleft.r), ∂H∂r(h, t.zright.r))
+    s = generalised_uturn_criterion(
+        rho, ∂H∂r(h, t.zleft.θ, t.zleft.r), ∂H∂r(h, t.zright.θ, t.zright.r)
+    )
     return Termination(s, false)
 end
 
@@ -595,7 +598,9 @@ phase point of `tright`, the right subtree.
 """
 function check_left_subtree(h::Hamiltonian, t::T, tleft::T, tright::T) where {T<:BinaryTree}
     rho = tleft.ts.rho + tright.zleft.r
-    s = generalised_uturn_criterion(rho, ∂H∂r(h, t.zleft.r), ∂H∂r(h, tright.zleft.r))
+    s = generalised_uturn_criterion(
+        rho, ∂H∂r(h, t.zleft.θ, t.zleft.r), ∂H∂r(h, tright.zleft.θ, tright.zleft.r)
+    )
     return Termination(s, false)
 end
 
@@ -608,7 +613,9 @@ function check_right_subtree(
     h::Hamiltonian, t::T, tleft::T, tright::T
 ) where {T<:BinaryTree}
     rho = tleft.zright.r + tright.ts.rho
-    s = generalised_uturn_criterion(rho, ∂H∂r(h, tleft.zright.r), ∂H∂r(h, t.zright.r))
+    s = generalised_uturn_criterion(
+        rho, ∂H∂r(h, tleft.zright.θ, tleft.zright.r), ∂H∂r(h, t.zright.θ, t.zright.r)
+    )
     return Termination(s, false)
 end
 
@@ -727,7 +734,7 @@ function transition(
         (
             n_steps=tree.nα,
             is_accept=true,
-            acceptance_rate=tree.sum_α / tree.nα,
+            acceptance_rate=(tree.sum_α / tree.nα),
             log_density=zcand.ℓπ.value,
             hamiltonian_energy=H,
             hamiltonian_energy_error=H - H0,
diff --git a/test/Project.toml b/test/Project.toml
index 8dfad5c7..cbe5bb8f 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -7,6 +7,7 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
diff --git a/test/riemannian.jl b/test/riemannian.jl
index 67b1cad0..a04c2de2 100644
--- a/test/riemannian.jl
+++ b/test/riemannian.jl
@@ -1,69 +1,629 @@
-using ReTest, AdvancedHMC
-
-include("../src/riemannian_hmc.jl")
-include("../src/riemannian_hmc_utility.jl")
-
+using ReTest, Random
+using AdvancedHMC, ForwardDiff, AbstractMCMC
+using LinearAlgebra
+using Distributions: MvNormal, logpdf
 using FiniteDiff:
     finite_difference_gradient, finite_difference_hessian, finite_difference_jacobian
-using Distributions: MvNormal
-using AdvancedHMC: neg_energy, energy
+using AdvancedHMC:
+    neg_energy,
+    energy,
+    ∂H∂θ,
+    ∂H∂r,
+    metric_eval,
+    metric_sensitivity,
+    logdet_grad_matrix,
+    kinetic_grad_matrix,
+    SoftAbsEval,
+    RiemannianMetric,
+    SoftAbsRiemannianMetric,
+    make_J,
+    _xcothx,
+    _xcothx_taylor,
+    _xcothx_exact,
+    _xcothx_deriv,
+    _xcothx_deriv_taylor,
+    _xcothx_deriv_exact,
+    _xcothx_cutoff
+using Statistics
+
+####
+#### Local toy targets
+####
+#### These replace the `MCMCLogDensityProblems` test dependency, which transitively
+#### pulls in DistributionsAD/ReverseDiff — a stack that currently fails to precompile
+#### against recent Distributions releases. We only need two densities and their
+#### (value, gradient) closures, so we define them here directly, matching the original
+#### definitions (standard Gaussian; Neal's funnel with v ~ N(0,3), xᵢ ~ N(0, exp(v/2))).
+
+struct HighDimGaussian
+    dim::Int
+end
+
+struct Funnel
+    dim::Int
+end
+Funnel() = Funnel(2)
+
+dim(target::Union{HighDimGaussian,Funnel}) = target.dim
+
+_logpdf_normal_std(x, m, s) = -(log(2π) + 2 * log(s) + ((x - m) / s)^2) / 2
+
+target_logpdf(::HighDimGaussian, θ) = sum(x -> _logpdf_normal_std(x, 0, 1), θ)
+function target_logpdf(::Funnel, θ)
+    v = θ[1]
+    lp = _logpdf_normal_std(v, 0, 3)
+    s = exp(v / 2)
+    @inbounds for i in 2:length(θ)
+        lp += _logpdf_normal_std(θ[i], 0, s)
+    end
+    return lp
+end
+
+gen_logpdf(target) = θ -> target_logpdf(target, θ)
+function gen_logpdf_grad(target, _)
+    f = gen_logpdf(target)
+    return θ -> (f(θ), ForwardDiff.gradient(f, θ))
+end
+
+####
+#### Test utilities
+####
+
+function gen_hess_fwd(func, x::AbstractVector)
+    cfg = ForwardDiff.HessianConfig(func, x)
+    H = Matrix{eltype(x)}(undef, length(x), length(x))
+
+    function hess(x::AbstractVector)
+        ForwardDiff.hessian!(H, func, x, cfg)
+        return H
+    end
+    return hess
+end
+
+function gen_∂G∂θ_fwd(Vfunc, x; f=identity)
+    chunk = ForwardDiff.Chunk(x)
+    tag = ForwardDiff.Tag(Vfunc, eltype(x))
+    jac_cfg = ForwardDiff.JacobianConfig(Vfunc, x, chunk, tag)
+    hess_cfg = ForwardDiff.HessianConfig(Vfunc, jac_cfg.duals, chunk, tag)
+
+    d = length(x)
+    out = zeros(eltype(x), d^2, d)
+
+    function ∂G∂θ_fwd(y)
+        hess = z -> ForwardDiff.hessian(Vfunc, z, hess_cfg, Val{false}())
+        ForwardDiff.jacobian!(out, hess, y, jac_cfg, Val{false}())
+        return out
+    end
+
+    return ∂G∂θ_fwd
+end
+
+function reshape_∂G∂θ(H)
+    d = size(H, 2)
+    return cat((H[((i - 1) * d + 1):(i * d), :] for i in 1:d)...; dims=3)
+end
+
+function prepare_sample(ℓπ, initial_θ, λ)
+    Vfunc = x -> -ℓπ(x)
+    _Hfunc = gen_hess_fwd(Vfunc, initial_θ)
+    Hfunc = x -> copy.(_Hfunc(x))
+
+    fstabilize = H -> H + λ * I
+    Gfunc = x -> begin
+        H = fstabilize(Hfunc(x))
+        all(isfinite, H) ? H : diagm(ones(length(x)))
+    end
+    _∂G∂θfunc = gen_∂G∂θ_fwd(x -> -ℓπ(x), initial_θ; f=fstabilize)
+    ∂G∂θfunc = x -> reshape_∂G∂θ(_∂G∂θfunc(x))
+
+    return Vfunc, Hfunc, Gfunc, ∂G∂θfunc
+end
 
-# Taken from https://github.com/JuliaDiff/FiniteDiff.jl/blob/master/test/finitedifftests.jl
 δ(a, b) = maximum(abs.(a - b))
 
-@testset "Riemannian" begin
-    hps = (; λ=1e-2, α=20.0, ϵ=0.1, n=6, L=8)
+####
+#### Tests for SoftAbs numerical-stability helpers (P0.1)
+####
+
+@testset "SoftAbs stability helpers" begin
+    @testset "no NaN at x = 0 (previously Inf − Inf)" begin
+        # Limits: x·coth(x) → 1, its derivative → 0 as x → 0.
+        @test _xcothx(0.0) == 1.0
+        @test _xcothx_deriv(0.0) == 0.0
+    end
+
+    @testset "Taylor and exact branches agree at the switch threshold" begin
+        # Compare the two branch implementations directly at the threshold.
+        # Tolerances are conservative — the actual disagreement is bounded by
+        # the larger of Taylor truncation (~eps at threshold for both) and exact-
+        # form cancellation (~6e-14 for _xcothx_deriv).
+        thresh = _xcothx_cutoff(Float64)
+        @test _xcothx_taylor(thresh) ≈ _xcothx_exact(thresh) atol = 1e-12
+        @test _xcothx_deriv_taylor(thresh) ≈ _xcothx_deriv_exact(thresh) atol = 1e-10
+    end
+
+    @testset "make_J: no NaN at λ = 0 and matches exact formula" begin
+        α = 20.0
+        # Mixed: one zero, one well-separated, one negative — the previously-broken case.
+        λ = [0.0, 1.7, -0.5]
+        J = make_J(λ, α)
+        @test all(isfinite, J)
+        @test J ≈ J'
+        # Off-diagonal entries not touching λ=0 match the analytic divided difference of
+        # softabs(λ) = λ·coth(αλ).
+        @test J[2, 3] ≈ (λ[2] * coth(α * λ[2]) - λ[3] * coth(α * λ[3])) / (λ[2] - λ[3]) rtol =
+            1e-10
+        # Diagonal at λ=0 is the well-defined limit softabs'(0) = 0.
+        @test J[1, 1] == 0.0
+        # Diagonal at λ≠0 matches softabs'(λ) = coth(αλ) − αλ·csch²(αλ).
+        @test J[2, 2] ≈ coth(α * λ[2]) - α * λ[2] * csch(α * λ[2])^2 rtol = 1e-10
+    end
+end
+
+####
+#### Tests for SoftAbsRiemannianMetric canonicalize option
+####
+
+@testset "SoftAbsRiemannianMetric canonicalize" begin
+    α = 10.0
+    # Fixed Hessian with known structure: mixed positive/negative eigenvalues.
+    H_fixed = [2.0 1.0; 1.0 -1.0]
+    calc_H = _ -> H_fixed
+    calc_∂H∂θ = _ -> zeros(2, 2, 2)
+    θ = [0.0, 0.0]
+
+    @testset "factorization is valid (G unchanged by canonicalization)" begin
+        G_raw = metric_eval(SoftAbsRiemannianMetric((2,), calc_H, calc_∂H∂θ, α), θ)
+        G_can = metric_eval(
+            SoftAbsRiemannianMetric((2,), calc_H, calc_∂H∂θ, α; canonicalize=true), θ
+        )
+        # Canonicalization must not change the metric matrix G = Q diag(softabsλ) Q'.
+        G_mat_raw = G_raw.Q * Diagonal(G_raw.softabsλ) * G_raw.Q'
+        G_mat_can = G_can.Q * Diagonal(G_can.softabsλ) * G_can.Q'
+        @test G_mat_raw ≈ G_mat_can
+        # The sign convention: largest-magnitude element of each column is positive.
+        for j in axes(G_can.Q, 2)
+            col = G_can.Q[:, j]
+            @test col[argmax(abs.(col))] > 0
+        end
+    end
+
+    @testset "sampler runs with canonicalize=true" begin
+        rng = MersenneTwister(1)
+        target = HighDimGaussian(2)
+        ℓπ = gen_logpdf(target)
+        ∂ℓπ∂θ = gen_logpdf_grad(target, zeros(2))
+        _, _, G, ∂G∂θ = prepare_sample(ℓπ, zeros(2), 1e-2)
+        metric = SoftAbsRiemannianMetric((2,), G, ∂G∂θ, 10.0; canonicalize=true)
+        hamiltonian = Hamiltonian(metric, GaussianKinetic(), ℓπ, ∂ℓπ∂θ)
+        integrator = GeneralizedLeapfrog(0.01, 3)
+        kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
+        @test_nowarn sample(
+            rng, hamiltonian, kernel, zeros(2), 20; verbose=false, progress=false
+        )
+    end
+end
+
+####
+#### Tests for unified API (RiemannianMetric, SoftAbsRiemannianMetric)
+####
 
+@testset "New Riemannian API" begin
     @testset "$(nameof(typeof(target)))" for target in [HighDimGaussian(2), Funnel()]
         rng = MersenneTwister(1110)
+        λ = 1e-2
 
         θ₀ = rand(rng, dim(target))
 
-        ℓπ = MCMCLogDensityProblems.gen_logpdf(target)
-        ∂ℓπ∂θ = MCMCLogDensityProblems.gen_logpdf_grad(target, θ₀)
+        ℓπ = gen_logpdf(target)
+        ∂ℓπ∂θ = gen_logpdf_grad(target, θ₀)
 
-        Vfunc, Hfunc, Gfunc, ∂G∂θfunc = prepare_sample_target(hps, θ₀, ℓπ)
+        _, _, Gfunc, ∂G∂θfunc = prepare_sample(ℓπ, θ₀, λ)
 
-        D = dim(target) # ==2 for this test
-        x = zeros(D) # randn(rng, D)
+        D = dim(target)
+        x = zeros(D)
         r = randn(rng, D)
 
-        @testset "Autodiff" begin
-            @test δ(finite_difference_gradient(ℓπ, x), ∂ℓπ∂θ(x)[end]) < 1e-4
-            @test δ(finite_difference_hessian(Vfunc, x), Hfunc(x)[end]) < 1e-4
-            # finite_difference_jacobian returns shape of (4, 2), reshape_∂G∂θ turns it into (2, 2, 2)
-            @test δ(reshape_∂G∂θ(finite_difference_jacobian(Gfunc, x)), ∂G∂θfunc(x)) < 1e-4
+        @testset "RiemannianMetric (PDMat-style)" begin
+            metric = RiemannianMetric((D,), Gfunc, ∂G∂θfunc)
+            kinetic = GaussianKinetic()
+            hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
+
+            # Test metric_eval returns a matrix
+            G_eval = metric_eval(metric, x)
+            @test G_eval isa AbstractMatrix
+            @test size(G_eval) == (D, D)
+
+            # Test metric_sensitivity
+            ∂G = metric_sensitivity(metric, x)
+            @test size(∂G) == (D, D, D)
+
+            # Test gradient matrices
+            M_logdet = logdet_grad_matrix(G_eval)
+            @test size(M_logdet) == (D, D)
+
+            M_kinetic = kinetic_grad_matrix(G_eval, r)
+            @test size(M_kinetic) == (D, D)
+
+            # Test ∂H∂θ against finite differences
+            Hamifunc = (x, r) -> energy(hamiltonian, r, x) + energy(hamiltonian, x)
+            Hamifuncx = x -> Hamifunc(x, r)
+            @test δ(
+                finite_difference_gradient(Hamifuncx, x), ∂H∂θ(hamiltonian, x, r).gradient
+            ) < 1e-4
+
+            # Test ∂H∂r against finite differences
+            Hamifuncr = r -> Hamifunc(x, r)
+            @test δ(finite_difference_gradient(Hamifuncr, r), ∂H∂r(hamiltonian, x, r)) <
+                1e-4
         end
 
-        @testset "$(nameof(typeof(hessmap)))" for hessmap in
-                                                  [IdentityMap(), SoftAbsMap(hps.α)]
-            metric = DenseRiemannianMetric((D,), Gfunc, ∂G∂θfunc, hessmap)
+        @testset "SoftAbsRiemannianMetric" begin
+            α = 20.0
+            metric = SoftAbsRiemannianMetric((D,), Gfunc, ∂G∂θfunc, α)
             kinetic = GaussianKinetic()
             hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
 
-            if hessmap isa SoftAbsMap || # only test kinetic energy for SoftAbsMap as that of IdentityMap can be non-PD
-                all(iszero, x) # or for x==0 that I know it's PD
-                @testset "Kinetic energy" begin
-                    Σ = hamiltonian.metric.map(hamiltonian.metric.G(x))
-                    @test neg_energy(hamiltonian, r, x) ≈ logpdf(MvNormal(zeros(D), Σ), r)
-                end
-            end
+            # Test metric_eval returns SoftAbsEval
+            G_eval = metric_eval(metric, x)
+            @test G_eval isa SoftAbsEval
+            @test size(G_eval.Q) == (D, D)
+            @test length(G_eval.softabsλ) == D
+            @test size(G_eval.J) == (D, D)
+            @test size(G_eval.M_logdet) == (D, D)
+
+            # Test standard operations on SoftAbsEval
+            v = randn(rng, D)
+            @test length(G_eval \ v) == D
+            @test logdet(G_eval) isa Real
+
+            # Test gradient matrices
+            M_logdet = logdet_grad_matrix(G_eval)
+            @test M_logdet === G_eval.M_logdet  # Should be cached
 
+            M_kinetic = kinetic_grad_matrix(G_eval, r)
+            @test size(M_kinetic) == (D, D)
+
+            # Test kinetic energy matches MvNormal logpdf
+            G_matrix = G_eval.Q * Diagonal(G_eval.softabsλ) * G_eval.Q'
+            @test neg_energy(hamiltonian, r, x) ≈
+                logpdf(MvNormal(zeros(D), Symmetric(G_matrix)), r)
+
+            # Test ∂H∂θ against finite differences
             Hamifunc = (x, r) -> energy(hamiltonian, r, x) + energy(hamiltonian, x)
             Hamifuncx = x -> Hamifunc(x, r)
+            @test δ(
+                finite_difference_gradient(Hamifuncx, x), ∂H∂θ(hamiltonian, x, r).gradient
+            ) < 1e-4
+
+            # Test ∂H∂r against finite differences
             Hamifuncr = r -> Hamifunc(x, r)
+            @test δ(finite_difference_gradient(Hamifuncr, r), ∂H∂r(hamiltonian, x, r)) <
+                1e-4
+        end
+    end
+end
+
+####
+#### Deprecation forwards (DenseRiemannianMetric → {RiemannianMetric, SoftAbsRiemannianMetric})
+####
+
+@testset "Deprecated DenseRiemannianMetric forwards" begin
+    rng = MersenneTwister(1110)
+    target = HighDimGaussian(2)
+    D = dim(target)
+    θ₀ = rand(rng, D)
+    λ = 1e-2
+
+    ℓπ = gen_logpdf(target)
+    _, _, Gfunc, ∂G∂θfunc = prepare_sample(ℓπ, θ₀, λ)
+
+    # 3-arg form → RiemannianMetric
+    m1 = @test_deprecated DenseRiemannianMetric((D,), Gfunc, ∂G∂θfunc)
+    @test m1 isa RiemannianMetric
+
+    # 4-arg form with IdentityMap → RiemannianMetric
+    m2 = @test_deprecated DenseRiemannianMetric((D,), Gfunc, ∂G∂θfunc, IdentityMap())
+    @test m2 isa RiemannianMetric
+
+    # 4-arg form with SoftAbsMap → SoftAbsRiemannianMetric (α threaded through)
+    m3 = @test_deprecated DenseRiemannianMetric((D,), Gfunc, ∂G∂θfunc, SoftAbsMap(20.0))
+    @test m3 isa SoftAbsRiemannianMetric
+    @test m3.α == 20.0
+end
+
+####
+#### Integration tests with sampling
+####
 
-            @testset "∂H∂θ" begin
-                @test δ(
-                    finite_difference_gradient(Hamifuncx, x),
-                    ∂H∂θ(hamiltonian, x, r).gradient,
-                ) < 1e-4
+@testset "Sampling with unified RiemannianMetric" begin
+    n_samples = 100
+    rng = MersenneTwister(1110)
+    initial_θ = rand(rng, D)
+    λ = 1e-2
+    _, _, G, ∂G∂θ = prepare_sample(ℓπ, initial_θ, λ)
+
+    metric = RiemannianMetric((D,), G, ∂G∂θ)
+    kinetic = GaussianKinetic()
+    hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
+
+    initial_ϵ = 0.01
+    integrator = GeneralizedLeapfrog(initial_ϵ, 6)
+    kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(8)))
+
+    samples, stats = sample(rng, hamiltonian, kernel, initial_θ, n_samples; progress=false)
+    @test length(samples) == n_samples
+    @test length(stats) == n_samples
+end
+
+@testset "Sampling with SoftAbsRiemannianMetric" begin
+    n_samples = 100
+    rng = MersenneTwister(1110)
+    initial_θ = rand(rng, D)
+    λ = 1e-2
+    _, _, G, ∂G∂θ = prepare_sample(ℓπ, initial_θ, λ)
+
+    metric = SoftAbsRiemannianMetric((D,), G, ∂G∂θ, 20.0)
+    kinetic = GaussianKinetic()
+    hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
+
+    initial_ϵ = 0.01
+    integrator = GeneralizedLeapfrog(initial_ϵ, 6)
+    kernel = HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(8)))
+
+    samples, stats = sample(rng, hamiltonian, kernel, initial_θ, n_samples; progress=false)
+    @test length(samples) == n_samples
+    @test length(stats) == n_samples
+end
+
+####
+#### Energy conservation tests
+####
+
+@testset "Energy conservation" begin
+    rng = MersenneTwister(42)
+    D_test = 2
+    target = HighDimGaussian(D_test)
+    θ₀ = rand(rng, D_test)
+    λ = 1e-2
+
+    ℓπ = gen_logpdf(target)
+    ∂ℓπ∂θ = gen_logpdf_grad(target, θ₀)
+    _, _, G, ∂G∂θ = prepare_sample(ℓπ, θ₀, λ)
+
+    @testset "SoftAbsRiemannianMetric energy conservation" begin
+        metric = SoftAbsRiemannianMetric((D_test,), G, ∂G∂θ, 20.0)
+        kinetic = GaussianKinetic()
+        hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
+
+        # Small step size for better energy conservation
+        integrator = GeneralizedLeapfrog(0.001, 10)
+
+        # Create initial phase point
+        θ_init = zeros(D_test)
+        r_init = randn(rng, D_test)
+        z0 = AdvancedHMC.phasepoint(hamiltonian, θ_init, r_init)
+        H0 = -AdvancedHMC.neg_energy(z0)
+
+        # Take 10 leapfrog steps
+        z1 = AdvancedHMC.step(integrator, hamiltonian, z0, 10)
+        H1 = -AdvancedHMC.neg_energy(z1)
+
+        # Energy should be approximately conserved
+        @test abs(H1 - H0) < 0.1
+    end
+end
+
+####
+#### Validation tests
+####
+
+@testset "Validation testing" begin
+    # 1D Wasserstein-1 distance
+    function w1(a::AbstractVector, b::AbstractVector)
+        sa = sort(a)
+        sb = sort(b)
+        return mean(abs.(sa .- sb))
+    end
+
+    @testset "Validation testing (Gaussian)" begin
+
+        # 1D normal Wasserstein-1 distance tolerance estimator
+        function w1_tol_normal_1d(;
+            n::Int, reps::Int=200, q::Float64=0.999, rng=Random.default_rng()
+        )
+            poolN = max(50_000, 50n)
+            pool = randn(rng, poolN)
+
+            vals = Vector{Float64}(undef, reps)
+            for i in 1:reps
+                a = pool[rand(rng, 1:poolN, n)]
+                b = pool[rand(rng, 1:poolN, n)]
+                vals[i] = w1(a, b)
+            end
+            sort!(vals)
+            return vals[clamp(ceil(Int, q * reps), 1, reps)]
+        end
+
+        target = HighDimGaussian(2)
+        rng = MersenneTwister(125)
+        λ = 1e-2
+
+        initial_θ = rand(rng, dim(target))
+
+        ℓπ = gen_logpdf(target)
+        ∂ℓπ∂θ = gen_logpdf_grad(target, initial_θ)
+
+        _, _, G, ∂G∂θ = prepare_sample(ℓπ, initial_θ, λ)
+
+        D = dim(target)
+
+        n_samples = 100
+        n_adapts = 50
+
+        tol_w1 = w1_tol_normal_1d(; n=n_samples, rng=rng)
+
+        tol_w1 *= 1.5
+
+        x_true = randn(rng, n_samples)
+        y_true = randn(rng, n_samples)
+
+        @testset "RiemannianMetric (PDMat-style)" begin
+            metric = RiemannianMetric((D,), G, ∂G∂θ)
+            kinetic = GaussianKinetic()
+            hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
+
+            initial_ϵ = 0.01
+            integrator = GeneralizedLeapfrog(initial_ϵ, 15)
+            kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
+
+            acceptance_rate = 0.9
+            adaptor = StepSizeAdaptor(acceptance_rate, integrator)
+
+            samples, stats = sample(
+                rng,
+                hamiltonian,
+                kernel,
+                initial_θ,
+                n_samples,
+                adaptor,
+                n_adapts;
+                progress=false,
+            )
+            θ = reduce(vcat, (permutedims(s) for s in samples))
+            # 1st marginal
+            @test w1(θ[:, 1], x_true) < tol_w1
+            # 2nd marginal
+            @test w1(θ[:, 2], y_true) < tol_w1
+        end
+
+        @testset "SoftAbsRiemannianMetric" begin
+            # We do not need SoftAbs for Gaussian target, so using small α
+            metric = SoftAbsRiemannianMetric((D,), G, ∂G∂θ, 1.0)
+            kinetic = GaussianKinetic()
+            hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
+
+            initial_ϵ = 0.01
+            integrator = GeneralizedLeapfrog(initial_ϵ, 15)
+            kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
+
+            acceptance_rate = 0.9
+            adaptor = StepSizeAdaptor(acceptance_rate, integrator)
+
+            samples, stats = sample(
+                rng,
+                hamiltonian,
+                kernel,
+                initial_θ,
+                n_samples,
+                adaptor,
+                n_adapts;
+                progress=false,
+            )
+
+            θ = reduce(vcat, (permutedims(s) for s in samples))
+            # 1st marginal
+            @test w1(θ[:, 1], x_true) < tol_w1
+            # 2nd marginal
+            @test w1(θ[:, 2], y_true) < tol_w1
+        end
+    end
+
+    @testset "Validation testing (Funnel)" begin
+
+        # Funnel i.i.d. sampler
+        # θ layout: [v, x1]
+        function funnel_iid(rng::AbstractRNG, n::Int)
+            v = 3.0 .* randn(rng, n)
+            X = Matrix{Float64}(undef, n, 1)
+            for i in 1:n
+                s = exp(v[i] / 2)
+                @inbounds X[i, :] .= s .* randn(rng, 1)
             end
+            return v, X
+        end
+
+        # 1D Wasserstein-1 distance tolerances for Funnel marginals
+        function funnel_w1_tols(;
+            n::Int,
+            reps::Int=200,
+            q::Float64=0.999,
+            inflate::Float64=1.8,
+            rng::AbstractRNG=Random.default_rng(),
+        )
+            vals_v = Vector{Float64}(undef, reps)
+            vals_x1 = Vector{Float64}(undef, reps)
+
+            for i in 1:reps
+                vA, XA = funnel_iid(rng, n)
+                vB, XB = funnel_iid(rng, n)
 
-            @testset "∂H∂r" begin
-                @test δ(finite_difference_gradient(Hamifuncr, r), ∂H∂r(hamiltonian, x, r)) <
-                    1e-4
+                x1A = XA[:, 1]
+                x1B = XB[:, 1]
+
+                vals_v[i] = w1(vA, vB)
+                vals_x1[i] = w1(x1A, x1B)
             end
+
+            sort!(vals_v)
+            sort!(vals_x1)
+            idx = clamp(ceil(Int, q * reps), 1, reps)
+
+            return (tol_v=inflate * vals_v[idx], tol_x1=inflate * vals_x1[idx])
+        end
+
+        target = Funnel()
+        rng = MersenneTwister(234)
+        λ = 1e-2
+
+        initial_θ = rand(rng, dim(target))
+
+        ℓπ = gen_logpdf(target)
+        ∂ℓπ∂θ = gen_logpdf_grad(target, initial_θ)
+
+        _, _, G, ∂G∂θ = prepare_sample(ℓπ, initial_θ, λ)
+
+        D = dim(target)
+
+        n_samples = 1500
+        n_adapts = 500
+        n_post = n_samples - n_adapts
+
+        # True samples (sized to match post-warmup chain samples)
+        v_true, X_true = funnel_iid(rng, n_post)
+
+        # Wasserstein-1 distance tolerances
+        tols = funnel_w1_tols(; n=n_post, rng=rng)
+
+        @testset "SoftAbsRiemannianMetric" begin
+            metric = SoftAbsRiemannianMetric((D,), G, ∂G∂θ, 40.0)
+            kinetic = GaussianKinetic()
+            hamiltonian = Hamiltonian(metric, kinetic, ℓπ, ∂ℓπ∂θ)
+
+            initial_ϵ = 0.01
+            integrator = GeneralizedLeapfrog(initial_ϵ, 5)
+            kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
+
+            acceptance_rate = 0.7
+            adaptor = StepSizeAdaptor(acceptance_rate, integrator)
+
+            samples, stats = sample(
+                rng,
+                hamiltonian,
+                kernel,
+                initial_θ,
+                n_samples,
+                adaptor,
+                n_adapts;
+                drop_warmup=true,
+                progress=false,
+            )
+
+            θ = reduce(vcat, (permutedims(s) for s in samples))
+            # 1st marginal
+            @test w1(θ[:, 1], v_true) < tols.tol_v
+            # 2nd marginal
+            @test w1(θ[:, 2], X_true[:, 1]) < tols.tol_x1
         end
     end
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 0015529d..32af650b 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -31,6 +31,7 @@ if GROUP == "All" || GROUP == "AdvancedHMC"
     include("abstractmcmc.jl")
     include("mcmcchains.jl")
     include("constructors.jl")
+    include("riemannian.jl")
     retest(; dry=false, verbose=Inf)
 end