From 7a5ecd9020567addcb0c0f03f4ec4f90a965ce17 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Fri, 30 Oct 2026 08:54:14 -0400 Subject: [PATCH 1/2] Split interface.jl by concern Move the contents of src/interface.jl unchanged into model.jl, chains.jl, samplers/{mala,parallel_mala,adaptive_mala}.jl and sample.jl. --- README.md | 6 +- src/ParallelMCMC.jl | 9 +- src/chains.jl | 55 ++ src/interface.jl | 1563 ------------------------------ src/model.jl | 422 ++++++++ src/sample.jl | 28 + src/samplers/adaptive_mala.jl | 222 +++++ src/samplers/mala.jl | 156 +++ src/samplers/parallel_mala.jl | 668 +++++++++++++ test/test-Reactant-Validation.jl | 2 +- 10 files changed, 1562 insertions(+), 1569 deletions(-) create mode 100644 src/chains.jl delete mode 100644 src/interface.jl create mode 100644 src/model.jl create mode 100644 src/sample.jl create mode 100644 src/samplers/adaptive_mala.jl create mode 100644 src/samplers/mala.jl create mode 100644 src/samplers/parallel_mala.jl diff --git a/README.md b/README.md index b443e61..5616afb 100644 --- a/README.md +++ b/README.md @@ -37,9 +37,9 @@ The approach and its scaling tricks (stochastic Hutchinson Jacobian estimators, | Sampler | Role | |---|---| -| [`ParallelMALASampler`](src/interface.jl) | **Primary** — parallel-across-sequence MALA via DEER; $O(\log T)$ per solve | -| [`MALASampler`](src/interface.jl) | Baseline — sequential MALA with a fixed step size | -| [`AdaptiveMALASampler`](src/interface.jl) | Baseline — sequential MALA with dual-averaging step-size adaptation | +| [`ParallelMALASampler`](src/samplers/parallel_mala.jl) | **Primary** — parallel-across-sequence MALA via DEER; $O(\log T)$ per solve | +| [`MALASampler`](src/samplers/mala.jl) | Baseline — sequential MALA with a fixed step size | +| [`AdaptiveMALASampler`](src/samplers/adaptive_mala.jl) | Baseline — sequential MALA with dual-averaging step-size adaptation | All samplers implement the [AbstractMCMC](https://github.com/TuringLang/AbstractMCMC.jl) interface and return [`FlexiChains`](https://pysm.dev/FlexiChains.jl) objects, so they slot into existing AbstractMCMC workflows. diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index 69a5873..c44b793 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -61,13 +61,18 @@ finished, so the caller may release the source. _copy_from_device_pointer!(::AbstractArray, ::Ptr{Cvoid}, ::AbstractString) = false #= Lives here rather than in `DEER` because both DEER's `ReactantHVP` fallbacks -and `interface.jl`'s gradient hooks report it. =# +and `model.jl`'s gradient hooks report it. =# const _REACTANT_LOAD_HINT = "AutoReactant requires Reactant.jl: add `using Reactant` to load ParallelMCMC's ReactantExt." include("MALA/MALA.jl") include("DEER/DEERScan.jl") include("DEER/DEER.jl") -include("interface.jl") +include("model.jl") +include("chains.jl") +include("samplers/mala.jl") +include("samplers/parallel_mala.jl") +include("samplers/adaptive_mala.jl") +include("sample.jl") export DensityModel export MALASampler, MALATransition, MALAState diff --git a/src/chains.jl b/src/chains.jl new file mode 100644 index 0000000..168a01c --- /dev/null +++ b/src/chains.jl @@ -0,0 +1,55 @@ +function _construct_flexichain( + ::Type{TKey}, + vals::AbstractMatrix{<:Real}, + internals::NamedTuple, + param_names::Any, + model::DensityModel, +) where {TKey} + # Wrap user-supplied names in `Parameter`. This allows people to specify, e.g., + # `param_names=(:x, :y, :z=>(2,))` without faffing with `Parameter` themselves. Also + # 'upgrade' symbol parameter names to VarNames if the user requested a VNChain. + to_parameter(vn::VarName) = FlexiChains.Parameter(vn) + to_parameter(s::Symbol) = FlexiChains.Parameter(TKey <: VarName ? VarName{s}() : s) + + N, D = size(vals) + + if param_names === nothing + param_names = model.param_names + end + wrapped_param_names = if param_names === nothing + # Wasn't defined either as `param_names` or `model.param_names`, so make some up + (to_parameter(:x) => (D,),) + else + map(param_names) do n + if n isa Pair + to_parameter(n.first) => n.second + elseif n isa TKey || n isa Symbol + to_parameter(n) + else + throw( + ArgumentError( + "param_names must be a collection of Pairs, Symbols, or $TKey, got $(typeof(n))", + ), + ) + end + end + end + + arr = reshape(vals, N, 1, D) + param_chain = FlexiChains.FlexiChain{TKey}(arr, Tuple(wrapped_param_names)) + + isempty(internals) && return param_chain + + #= + The internals have mixed element types (e.g. `Bool` for `accepted`/`is_warmup`, + `FP` for `logp`), so they cannot share a single stacked array without widening. + Store each in its own column via the dict-of-arrays constructor, which preserves + the natural element type (see issue #49), then merge with the parameters. + =# + extras = OrderedDict{FlexiChains.ParameterOrExtra{<:TKey},AbstractArray}( + FlexiChains.Extra(name) => internals[name] for name in keys(internals) + ) + extra_chain = FlexiChains.FlexiChain{TKey}(N, 1, extras) + + return merge(param_chain, extra_chain) +end diff --git a/src/interface.jl b/src/interface.jl deleted file mode 100644 index ddaf7fe..0000000 --- a/src/interface.jl +++ /dev/null @@ -1,1563 +0,0 @@ -#= -AbstractMCMC interface for ParallelMCMC samplers. - -Defines model/sampler/state/transition types and implements -`AbstractMCMC.step` so that `sample(model, sampler, N)` works out of the box. -=# - -""" - DensityModel(logdensity, grad_logdensity, dim; param_names, logdensity_batch, grad_logdensity_batch, hvp, hvp_batch) - -Wrap a log-density and its derivatives for ParallelMCMC samplers. - -Each derivative slot (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, -`hvp_batch`) takes a callable or an `ADTypes.AbstractADType`, so a model can be -built from the log-density alone: - - DensityModel(logp, AutoForwardDiff(), dim) - -Backends are prepared when sampling starts. An HVP backend differentiates a -callable gradient once. With an AD-derived gradient it forms -`DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` on the -log-density. An explicit `SecondOrder` also bypasses the gradient slot. - -- `logdensity(x::AbstractVector) -> Real` -- `grad_logdensity` — callable `x -> AbstractVector`, or an AD backend. -- `hvp` — optional callable `(x, v) -> AbstractVector`, or a backend. If - `nothing`, DEER builds the HVP from the sampler's `backend`. -- `logdensity_batch(X::AbstractMatrix) -> AbstractVector` — optional batched - log-density over independent columns. -- `grad_logdensity_batch` — optional callable `X -> AbstractMatrix`, or a - backend. It is derived from `logdensity_batch` when `grad_logdensity` is a - backend. -- `hvp_batch` — optional callable `(X, V) -> AbstractMatrix`, or a backend, - resolved like `hvp`. -- `dim::Int` — dimensionality of the parameter space -- `param_names` — optional collection of parameter names used in `FlexiChains` output. If - `nothing` (the default), uses a single vector-valued parameter `:x` with shape `(dim,)`. - See the [`Parameter names`](@ref parameter-names) section of the docs for more - information. - -Both batched derivative slots require `logdensity_batch`. `AutoReactant()` has -additional pairing and tracing constraints; see the GPU guide. -""" -struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel - logdensity::F - grad_logdensity::G - hvp::H - logdensity_batch::FB - grad_logdensity_batch::GB - hvp_batch::HB - dim::Int - param_names::PN -end - -function DensityModel( - logdensity, - grad_logdensity, - dim::Int; - param_names=nothing, - hvp=nothing, - logdensity_batch=nothing, - grad_logdensity_batch=nothing, - hvp_batch=nothing, -) - logdensity isa AbstractADType && - throw(ArgumentError("logdensity must be callable, not an AD backend")) - logdensity_batch isa AbstractADType && - throw(ArgumentError("logdensity_batch must be callable, not an AD backend")) - grad_logdensity === nothing && throw( - ArgumentError( - "grad_logdensity must be a callable or an ADTypes.AbstractADType backend" - ), - ) - _batch_needs_logp(name) = throw(ArgumentError("$name requires logdensity_batch")) - if logdensity_batch === nothing - grad_logdensity_batch === nothing || _batch_needs_logp("grad_logdensity_batch") - hvp_batch === nothing || _batch_needs_logp("hvp_batch") - end - return DensityModel( - logdensity, - grad_logdensity, - hvp, - logdensity_batch, - grad_logdensity_batch, - hvp_batch, - dim, - param_names, - ) -end - -""" -A [`DensityModel`](@ref) whose AD backends have been resolved to callables. -`source` prevents a cached preparation from being reused with another model. -""" -struct PreppedDensityModel{F,G,H,FB,GB,HB,PN,SM<:DensityModel} - logdensity::F - grad_logdensity::G - hvp::H - logdensity_batch::FB - grad_logdensity_batch::GB - hvp_batch::HB - dim::Int - param_names::PN - source::SM -end - -# Identity is required: reusing another model's preparation would sample the wrong target. -_prepped_for(prepped::PreppedDensityModel, model::DensityModel) = prepped.source === model - -# DI keys preparations on function identity. Other input types use an unprepared call. -struct _ADGradient{F,B<:AbstractADType,P,TX} - logdensity::F - backend::B - prep::P -end -function (g::_ADGradient{F,B,P,TX})(x) where {F,B,P,TX} - if x isa TX - return DI.gradient(g.logdensity, g.prep, g.backend, x) - else - return DI.gradient(g.logdensity, g.backend, x) - end -end - -# Columns of X are independent samples, so ∇_X sum(logdensity_batch(X)) -# stacks the per-column gradients. -struct _BatchLogdensitySum{F} - logdensity_batch::F -end -(c::_BatchLogdensitySum)(X) = sum(c.logdensity_batch(X)) - -struct _ADGradientBatch{C<:_BatchLogdensitySum,B<:AbstractADType,P,TX} - closure::C - backend::B - prep::P -end -function (g::_ADGradientBatch{C,B,P,TX})(X) where {C,B,P,TX} - if X isa TX - return DI.gradient(g.closure, g.prep, g.backend, X) - else - return DI.gradient(g.closure, g.backend, X) - end -end - -function _resolve_gradient(logdensity, backend::AbstractADType, x_template::AbstractVector) - prep = DI.prepare_gradient(logdensity, backend, x_template) - return _ADGradient{typeof(logdensity),typeof(backend),typeof(prep),typeof(x_template)}( - logdensity, backend, prep - ) -end - -function _resolve_gradient_batch( - logdensity_batch, backend::AbstractADType, X_template::AbstractMatrix -) - closure = _BatchLogdensitySum(logdensity_batch) - prep = DI.prepare_gradient(closure, backend, X_template) - return _ADGradientBatch{typeof(closure),typeof(backend),typeof(prep),typeof(X_template)}( - closure, backend, prep - ) -end - -function _resolve_gradient( - logdensity, backend::ADTypes.AutoReactant, x_template::AbstractVector -) - return _reactant_resolve_gradient(logdensity, backend, x_template) -end -function _reactant_resolve_gradient(logdensity, backend, x_template) - return error(_REACTANT_LOAD_HINT) -end - -function _resolve_gradient_batch( - logdensity_batch, backend::ADTypes.AutoReactant, X_template::AbstractMatrix -) - return _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) -end -function _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) - return error(_REACTANT_LOAD_HINT) -end - -# An explicit SecondOrder differentiates logdensity and bypasses the gradient slot. -function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) - if hvp_backend isa DI.SecondOrder - return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) - elseif grad_backend !== nothing - return DEER._make_hvp_fn_second_order( - logdensity, _second_order(hvp_backend, grad_backend), x_template - ) - else - return DEER._make_hvp_fn( - DEER._hvp_strategy(hvp_backend), grad, hvp_backend, x_template - ) - end -end - -function _resolve_hvp_batch( - logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template -) - if hvp_backend isa DI.SecondOrder - return DEER._make_hvp_batch_fn_second_order( - _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template - ) - elseif grad_batch_backend !== nothing - return DEER._make_hvp_batch_fn_second_order( - _BatchLogdensitySum(logdensity_batch), - _second_order(hvp_backend, grad_batch_backend), - X_template, - ) - else - return DEER._make_hvp_batch_fn( - DEER._hvp_strategy(hvp_backend), grad_batch, hvp_backend, X_template - ) - end -end - -# With no HVP source, run forward mode over an AD-derived gradient. -# AutoReactant gradient maps to itself, so the Reactant pairing still holds. -_default_hvp_backend(::Nothing) = nothing -_default_hvp_backend(grad_backend::AbstractADType) = DI.forward_counterpart(grad_backend) - -# Reactant implements the two passes itself; DI composes all other backends. -_second_order(hvp_backend, grad_backend) = DI.SecondOrder(hvp_backend, grad_backend) -function _second_order(hvp_backend::ADTypes.AutoReactant, ::ADTypes.AutoReactant) - return hvp_backend -end - -_check_reactant_pair(grad_backend, hvp_backend) = nothing -_check_reactant_pair(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = nothing -_check_reactant_pair(::Nothing, ::ADTypes.AutoReactant) = nothing - -function _check_reactant_pair(grad_backend::ADTypes.AutoReactant, hvp_backend) - return throw( - ArgumentError( - "an AutoReactant gradient requires an AutoReactant HVP backend; " * - "got $(hvp_backend)", - ), - ) -end - -function _check_reactant_pair(grad_backend, hvp_backend::ADTypes.AutoReactant) - return throw( - ArgumentError( - "an AutoReactant HVP requires an AutoReactant or callable gradient; " * - "got backend $(grad_backend)", - ), - ) -end - -function _check_reactant_pair(::ADTypes.AutoReactant, hvp_backend::DI.SecondOrder) - return throw( - ArgumentError( - "an AutoReactant gradient requires a bare AutoReactant HVP backend; " * - "got $(hvp_backend)", - ), - ) -end - -function _check_reactant_pair(grad_backend, hvp_backend::DI.SecondOrder) - _second_order_has_reactant(hvp_backend) || return nothing - return throw( - ArgumentError( - "AutoReactant cannot be used inside DifferentiationInterface.SecondOrder; " * - "use a bare AutoReactant() HVP backend", - ), - ) -end - -function _second_order_has_reactant(so::DI.SecondOrder) - return DI.outer(so) isa ADTypes.AutoReactant || DI.inner(so) isa ADTypes.AutoReactant -end - -_check_reactant_hvp_source(grad, hvp_backend) = nothing - -""" - _prepare_model(model, x_template) -> PreppedDensityModel - _prepare_model(model, x_template, T::Int, backend) -> PreppedDensityModel - -Resolve AD backends into prepared callables at `x_template`. The two-argument -form prepares only the gradient. The four-argument form also prepares DEER's -HVP and batched slots at shape `(model.dim, T)`. -""" -function _prepare_model(model::DensityModel, x_template::AbstractVector) - grad = if model.grad_logdensity isa AbstractADType - _resolve_gradient(model.logdensity, model.grad_logdensity, x_template) - else - model.grad_logdensity - end - return PreppedDensityModel( - model.logdensity, - grad, - nothing, - model.logdensity_batch, - nothing, - nothing, - model.dim, - model.param_names, - model, - ) -end - -function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, backend) - grad_backend = - model.grad_logdensity isa AbstractADType ? model.grad_logdensity : nothing - - # Validate the pair before an AutoReactant gradient triggers compilation. - needs_hvp = model.hvp === nothing || model.hvp isa AbstractADType - hvp_backend = if needs_hvp - hb = if model.hvp !== nothing - model.hvp - elseif backend !== nothing - backend - else - _default_hvp_backend(grad_backend) - end - hb === nothing && throw( - ArgumentError( - "ParallelMALASampler needs a Hessian-vector product: supply `hvp` " * - "on the DensityModel (callable or AD backend), pass `backend=` " * - "to ParallelMALASampler, or give `grad_logdensity` as an AD backend", - ), - ) - _check_reactant_pair(grad_backend, hb) - _check_reactant_hvp_source(model.grad_logdensity, hb) - hb - else - nothing - end - - grad = if grad_backend !== nothing - _resolve_gradient(model.logdensity, grad_backend, x_template) - else - model.grad_logdensity - end - - hvp = if needs_hvp - _resolve_hvp(model.logdensity, grad, grad_backend, hvp_backend, x_template) - else - model.hvp - end - - grad_batch = model.grad_logdensity_batch - if grad_batch === nothing && model.logdensity_batch !== nothing - grad_batch = grad_backend - end - grad_batch_backend = grad_batch isa AbstractADType ? grad_batch : nothing - - hvp_batch = model.hvp_batch - batch_active = model.logdensity_batch !== nothing && grad_batch !== nothing - - if batch_active - # Prepare on x0 in every column; zeros need not be in the support. - X_template = similar(x_template, length(x_template), T) - X_template .= x_template - - needs_hvp_batch = hvp_batch === nothing || hvp_batch isa AbstractADType - hvp_batch_backend = if needs_hvp_batch - hbb = if hvp_batch !== nothing - hvp_batch - elseif model.hvp isa AbstractADType - model.hvp - elseif backend !== nothing - backend - else - _default_hvp_backend(grad_batch_backend) - end - hbb === nothing && throw( - ArgumentError("batched DEER requires `hvp_batch` or a sampler `backend`"), - ) - _check_reactant_pair(grad_batch_backend, hbb) - hbb - else - nothing - end - - if grad_batch_backend !== nothing - grad_batch = _resolve_gradient_batch( - model.logdensity_batch, grad_batch_backend, X_template - ) - end - - if needs_hvp_batch - hvp_batch = _resolve_hvp_batch( - model.logdensity_batch, - grad_batch, - grad_batch_backend, - hvp_batch_backend, - X_template, - ) - end - elseif hvp_batch !== nothing - # Do not silently ignore an explicitly supplied HVP. - throw( - ArgumentError( - "hvp_batch has no batched gradient; supply grad_logdensity_batch or an AD-backed grad_logdensity", - ), - ) - end - - return PreppedDensityModel( - model.logdensity, - grad, - hvp, - model.logdensity_batch, - grad_batch, - hvp_batch, - model.dim, - model.param_names, - model, - ) -end - -#= Wrappers for LogDensityProblems objects. They are defined here rather than in -LogDensityProblemsExt so `_check_reactant_hvp_source` can dispatch on them. The call -methods live in the extension. =# -struct LogDensityProblemPrimal{L} - ld::L -end -struct LogDensityProblemGradient{L} - ld::L -end - -function _check_reactant_hvp_source( - ::LogDensityProblemGradient, hvp_backend::ADTypes.AutoReactant -) - return throw( - ArgumentError( - "AutoReactant cannot differentiate a LogDensityProblems-derived gradient; " * - "supply a Reactant-traceable gradient callable instead", - ), - ) -end - -""" - MALASampler(epsilon; cholM=nothing) - -Metropolis-Adjusted Langevin Algorithm sampler with step size `epsilon`. - -Optionally pass `cholM = cholesky(M)` to use a mass matrix `M` as a -preconditioner. The proposal becomes `y = x + ε M ∇logp(x) + √(2ε) L ξ` -where `L` is the Cholesky factor of `M`. -""" -struct MALASampler{FP<:AbstractFloat,CM} <: AbstractMCMC.AbstractSampler - epsilon::FP - cholM::CM -end - -function MALASampler(epsilon::Real; cholM=nothing) - epsilon > 0 || throw(ArgumentError("epsilon must be > 0, got $epsilon")) - eps_f = float(epsilon) - return MALASampler{typeof(eps_f),typeof(cholM)}(eps_f, cholM) -end - -""" -State for a `MALASampler` chain. Holds the prepped model so AD preparation -happens once per chain. -""" -struct MALAState{V<:AbstractVector,L<:Real,W,NV<:AbstractVector,H,DM<:PreppedDensityModel} - x::V - logp::L - workspace::W - noise::NV - noise_host::H - model::DM -end - -""" -One `MALASampler` sample: parameter vector `x`, its log-density `logp`, and an -accept/reject flag. -""" -struct MALATransition{V<:AbstractVector,L<:Real} - x::V - logp::L - accepted::Bool -end - -function _make_noise_buffer(x::AbstractVector, ::Type{FP}, D::Int) where {FP} - ξ = similar(x, FP, D) - host = needs_host_staging(ξ) ? Vector{FP}(undef, D) : nothing - return ξ, host -end - -function _randn_like!( - rng::Random.AbstractRNG, ξ::AbstractVector{FP}, host::Union{Nothing,AbstractVector{FP}} -) where {FP} - if needs_host_staging(ξ) - host === nothing && - error("normal noise for $(typeof(ξ)) requires a reusable host buffer") - randn!(rng, host) - copyto!(ξ, host) - else - randn!(rng, ξ) - end - return ξ -end - -function AbstractMCMC.step( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::MALASampler{FP}; - initial_params=nothing, - kwargs..., -) where {FP} - x = if initial_params !== nothing - copy(initial_params) - else - randn(rng, FP, model.dim) - end - model = _prepare_model(model, x) - logp_val = model.logdensity(x) - ws = MALA.MALAWorkspace(x) - noise, noise_host = _make_noise_buffer(x, FP, model.dim) - t = MALATransition(x, logp_val, true) - s = MALAState(x, logp_val, ws, noise, noise_host, model) - return t, s -end - -function AbstractMCMC.step( - rng::Random.AbstractRNG, - user_model::DensityModel, - sampler::MALASampler, - state::MALAState; - kwargs..., -) - #= Reuse the state's preparation, unless the state came from a different - model (via `initial_state`), in which case the model we were handed is the - one to sample. =# - model = if _prepped_for(state.model, user_model) - state.model - else - _prepare_model(user_model, state.x) - end - x = state.x - ϵ = sampler.epsilon - D = model.dim - - ξ = _randn_like!(rng, state.noise, state.noise_host) - u = rand(rng) - - x_next = similar(x) - x_next, accepted, _ = MALA.mala_step_with_logα!( - x_next, - state.workspace, - model.logdensity, - model.grad_logdensity, - x, - ϵ, - ξ, - u; - cholM=sampler.cholM, - ) - - logp_val = accepted ? model.logdensity(x_next) : state.logp - t = MALATransition(x_next, logp_val, accepted) - s = MALAState(x_next, logp_val, state.workspace, state.noise, state.noise_host, model) - return t, s -end - -for TKey in (Symbol, VarName) - @eval function AbstractMCMC.bundle_samples( - samples::Vector{<:MALATransition}, - model::DensityModel, - sampler::MALASampler, - state::MALAState, - ::Type{FlexiChains.FlexiChain{$TKey}}; - param_names=nothing, - kwargs..., - ) - N = length(samples) - D = model.dim - # Follow the sampler's working precision instead of pinning `Float64`, so a - # Float32 (e.g. GPU) run yields a Float32 chain rather than a silently widened one. - FP = typeof(sampler.epsilon) - - vals = Matrix{FP}(undef, N, D) - logp = Vector{FP}(undef, N) - accepted = Vector{Bool}(undef, N) - - for i in 1:N - s = samples[i] - vals[i, :] .= s.x - logp[i] = s.logp - accepted[i] = s.accepted - end - - internals = (logp=logp, accepted=accepted) - return _construct_flexichain($TKey, vals, internals, param_names, model) - end -end - -""" - MALATapeElement(ξ, u) - -One element of the MALA noise tape: a noise vector `ξ ~ N(0,I)` and a uniform -scalar `u ~ Uniform(0,1)`. Stored with a concrete vector type `V` for type -stability inside `DEER.TapedRecursion`. -""" -struct MALATapeElement{FP<:AbstractFloat,V<:AbstractVector{FP}} - ξ::V - u::FP -end - -""" - ParallelMALASampler(epsilon; T, maxiter, tol_abs, tol_rel, jacobian, damping, probes, cholM, backend) - -DEER-parallelized MALA sampler. - -Supported Jacobian modes are `:stoch_diag` (the default Hutchinson diagonal -estimator) and `:diag` (exact diagonal via `D` JVPs). - -`backend` supplies Hessian-vector products when the `DensityModel` brings no -`hvp` / `hvp_batch` of its own. It does not supply gradients. Leave it out when -the model carries its own HVPs. When it is left out and `grad_logdensity` is an -AD backend, the HVP runs the forward-mode counterpart of that backend over the -gradient (`DifferentiationInterface.forward_counterpart`). - -With `backend=ADTypes.AutoReactant()`, `grad_logdensity` must be `AutoReactant()` -or a callable. - -`callback` and `initial_state` route through the generic step-wise loop instead. -""" -struct ParallelMALASampler{FP<:AbstractFloat,CM,AD} <: AbstractMCMC.AbstractSampler - epsilon::FP - T::Int - maxiter::Int - tol_abs::FP - tol_rel::FP - jacobian::Symbol - damping::FP - probes::Int - cholM::CM - backend::AD -end - -function ParallelMALASampler( - epsilon::Real; - T::Int=64, - maxiter::Int=200, - tol_abs::Real=1e-6, - tol_rel::Real=1e-5, - jacobian::Symbol=:stoch_diag, - damping::Real=0.5, - probes::Int=1, - cholM=nothing, - backend=nothing, -) - epsilon > 0 || throw(ArgumentError("epsilon must be > 0, got $epsilon")) - (jacobian === :stoch_diag || jacobian === :diag) || - throw(ArgumentError("jacobian must be :stoch_diag or :diag")) - eps_f = float(epsilon) - FP = typeof(eps_f) - return ParallelMALASampler{FP,typeof(cholM),typeof(backend)}( - eps_f, - T, - maxiter, - FP(tol_abs), - FP(tol_rel), - jacobian, - FP(damping), - probes, - cholM, - backend, - ) -end - -""" -State for a `ParallelMALASampler` chain. Holds the prepped model so AD -preparation happens once per chain. -""" -struct ParallelMALAState{ - V<:AbstractVector,L<:Real,M<:AbstractMatrix,LV<:AbstractVector,W,DM<:PreppedDensityModel -} - x::V - logp::L - trajectory::M - logps::LV - workspace::W - tape::Vector{<:MALATapeElement} - t::Int - model::DM -end - -""" -One `ParallelMALASampler` sample: parameter vector `x` and its log-density `logp`. -""" -struct ParallelMALATransition{V<:AbstractVector,L<:Real} - x::V - logp::L -end - -#= Kept samples of a chain, one block of columns per DEER solve. `ends[b]` is the -number of kept samples up to and including block `b`; with thinning the blocks -are not all the same width. =# -struct ParallelMALABlockSamples{B<:AbstractVector,L<:AbstractVector} <: - AbstractVector{ParallelMALATransition} - blocks::B - logps::L - ends::Vector{Int} - n::Int -end - -Base.size(samples::ParallelMALABlockSamples) = (samples.n,) -Base.length(samples::ParallelMALABlockSamples) = samples.n -Base.IndexStyle(::Type{<:ParallelMALABlockSamples}) = IndexLinear() - -function Base.getindex(samples::ParallelMALABlockSamples, i::Int) - 1 <= i <= samples.n || throw(BoundsError(samples, i)) - b = searchsortedfirst(samples.ends, i) - t = i - (b == 1 ? 0 : samples.ends[b - 1]) - return ParallelMALATransition(samples.blocks[b][:, t], samples.logps[b][t]) -end - -function _make_mala_tape_block( - rng::Random.AbstractRNG, x0::AbstractVector, ::Type{FP}, D::Int, T::Int -) where {FP} - Xi = similar(x0, FP, D, T) - U_host = Vector{FP}(undef, T) - - if needs_host_staging(Xi) - Xi_host = Matrix{FP}(undef, D, T) - for t in 1:T - randn!(rng, view(Xi_host, :, t)) - U_host[t] = FP(rand(rng)) - end - copyto!(Xi, Xi_host) - else - for t in 1:T - randn!(rng, view(Xi, :, t)) - U_host[t] = FP(rand(rng)) - end - end - - U = similar(x0, FP, T) - copyto!(U, U_host) - tape = [MALATapeElement(view(Xi, :, t), U_host[t]) for t in 1:T] - return tape, Xi, U -end - -function _build_mala_deer_rec( - model::PreppedDensityModel, - ε::Real, - tape::Vector{<:MALATapeElement}, - x0_like::AbstractVector; - cholM=nothing, - tape_noise=nothing, - tape_uniforms=nothing, -) - logp = model.logdensity - gradlogp = model.grad_logdensity - - hvp_fn = model.hvp - - # Exact forward step. - step_fwd = - (x, te) -> MALA.mala_step_taped(logp, gradlogp, x, ε, te.ξ, te.u; cholM=cholM) - - # Explicit analytical JVP of the surrogate step. - jvp = - (x, te, v) -> MALA.mala_step_surrogate_sigmoid_jvp( - logp, gradlogp, x, ε, te.ξ, te.u, v, hvp_fn; cholM=cholM - ) - - # Fused forward step + JVP. - fwd_and_jvp = - (x, te, v) -> MALA.mala_step_taped_and_jvp( - logp, gradlogp, x, ε, te.ξ, te.u, v, hvp_fn; cholM=cholM - ) - - fwd_and_jvp_batch = - if model.logdensity_batch !== nothing && model.grad_logdensity_batch !== nothing - D = length(x0_like) - T = length(tape) - - (tape_noise === nothing) == (tape_uniforms === nothing) || throw( - ArgumentError("tape_noise and tape_uniforms must be provided together") - ) - Xi, U = if tape_noise === nothing - Xi_local = similar(x0_like, D, T) - U_host = Vector{typeof(tape[1].u)}(undef, T) - U_local = similar(x0_like, typeof(tape[1].u), T) - - for t in 1:T - copyto!(Xi_local, (t - 1) * D + 1, tape[t].ξ, 1, D) - U_host[t] = tape[t].u - end - copyto!(U_local, U_host) - Xi_local, U_local - else - size(tape_noise) == (D, T) || - throw(DimensionMismatch("tape_noise must have size (D, T)")) - length(tape_uniforms) == T || - throw(DimensionMismatch("tape_uniforms must have length T")) - tape_noise, tape_uniforms - end - - X_template = similar(x0_like, D, T) - fill!(X_template, zero(eltype(X_template))) - batch_ws = MALA.MALABatchedWorkspace(X_template) - FT = similar(X_template) - Jt = similar(X_template) - hvp_batch = model.hvp_batch - - (Xbar, Z) -> MALA.mala_step_batched_fwd_and_jvp!( - FT, - Jt, - batch_ws, - model.logdensity_batch, - model.grad_logdensity_batch, - hvp_batch, - Xbar, - ε, - Xi, - U, - Z; - cholM=cholM, - ) - else - nothing - end - - return DEER.TapedRecursion( - step_fwd, jvp, tape; fwd_and_jvp=fwd_and_jvp, fwd_and_jvp_batch=fwd_and_jvp_batch - ) -end - -function _deer_solve_new_tape( - rng::Random.AbstractRNG, - model::PreppedDensityModel, - sampler::ParallelMALASampler, - x0::AbstractVector; - workspace=nothing, -) - D = model.dim - T = sampler.T - FP = typeof(sampler.epsilon) - - tape, Xi, U = _make_mala_tape_block(rng, x0, FP, D, T) - - rec = _build_mala_deer_rec( - model, - sampler.epsilon, - tape, - x0; - cholM=sampler.cholM, - tape_noise=Xi, - tape_uniforms=U, - ) - - ws = workspace === nothing ? DEER.DEERWorkspace(x0, T) : workspace - - S = DEER.solve( - rec, - x0; - tol_abs=sampler.tol_abs, - tol_rel=sampler.tol_rel, - maxiter=sampler.maxiter, - jacobian=sampler.jacobian, - damping=sampler.damping, - probes=sampler.probes, - rng=rng, - workspace=ws, - copy_result=false, - ) - return S, tape, ws -end - -function _trajectory_logps(model::PreppedDensityModel, S::AbstractMatrix) - if model.logdensity_batch !== nothing - return Array(model.logdensity_batch(S)) - end - - T = size(S, 2) - return [model.logdensity(S[:, t]) for t in 1:T] -end - -function _parallel_mala_initial_x( - rng::Random.AbstractRNG, model::DensityModel, ::ParallelMALASampler{FP}, initial_params -) where {FP} - return initial_params !== nothing ? copy(initial_params) : randn(rng, FP, model.dim) -end - -function _parallel_mala_progress(progress, progressname) - progress === true && return AbstractMCMC.CreateNewProgressBar(progressname) - progress === false && return AbstractMCMC.NoLogging() - return progress -end - -function _parallel_mala_update_progress!( - progress, nsteps::Int, Ntotal::Int, next_update::Real, threshold::Real -) - if nsteps >= next_update && next_update <= Ntotal - AbstractMCMC.update_progress!(progress, min(nsteps / Ntotal, 1)) - next_update += threshold - end - return next_update -end - -""" -Which steps of a `ParallelMALASampler` run are kept. Global step `g` is kept iff -`g > discard_initial` and `(g - discard_initial - 1) % thinning == 0`, i.e., the -steps the `AbstractMCMC.mcmcsample` loop saves, so the batched path and the -step-wise fallback return the same chain for the same RNG. -""" -struct ParallelMALASchedule - discard_initial::Int - thinning::Int - Ntotal::Int -end - -function ParallelMALASchedule( - N::Int; discard_initial::Int=0, thinning::Int=1, num_warmup::Int=0 -) - N > 0 || error("the number of samples must be ≥ 1") - discard_initial >= 0 || - throw(ArgumentError("number of discarded samples must be non-negative")) - thinning >= 1 || throw(ArgumentError("thinning must be ≥ 1")) - num_warmup >= 0 || - throw(ArgumentError("number of warm-up samples must be non-negative")) - Ntotal = thinning * (N - 1) + discard_initial + 1 - Ntotal >= num_warmup || throw( - ArgumentError("number of warm-up samples exceeds the total number of samples") - ) - return ParallelMALASchedule(discard_initial, thinning, Ntotal) -end - -#= Columns of a block covering global steps `offset+1:offset+ncols` that the schedule -keeps, and the chain row the first of them lands in. =# -function _kept_columns(sched::ParallelMALASchedule, offset::Int, ncols::Int) - k0 = max(0, cld(offset - sched.discard_initial, sched.thinning)) - first_col = sched.discard_initial + 1 + k0 * sched.thinning - offset - return first_col:(sched.thinning):ncols, k0 + 1 -end - -#= Runs DEER blocks until `sched.Ntotal` steps are consumed, passing kept columns -to `sink!(S, logps, cols, first_row)`. Uses full-length blocks (`T`) to match the -fallback's RNG stream. Returns the last block; `S` aliases the workspace. =# -function _parallel_mala_run!( - sink!, - rng::Random.AbstractRNG, - model::DensityModel, - sampler::ParallelMALASampler, - sched::ParallelMALASchedule; - initial_params=nothing, - progress=AbstractMCMC.PROGRESS[], - progressname="Sampling", -) - progress = _parallel_mala_progress(progress, progressname) - x0 = _parallel_mala_initial_x(rng, model, sampler, initial_params) - prepped = _prepare_model(model, x0, sampler.T, sampler.backend) - ws = nothing - nsteps = 0 - Ntotal = sched.Ntotal - next_update = Ntotal / AbstractMCMC.get_n_updates(progress) - threshold = next_update - last = nothing - - AbstractMCMC.@maybewithricherlogger begin - AbstractMCMC.init_progress!(progress) - try - while nsteps < Ntotal - S, tape, ws = _deer_solve_new_tape(rng, prepped, sampler, x0; workspace=ws) - logps = _trajectory_logps(prepped, S) - nconsume = min(sampler.T, Ntotal - nsteps) - cols, first_row = _kept_columns(sched, nsteps, nconsume) - sink!(S, logps, cols, first_row) - - x0 = copy(view(S, :, nconsume)) - nsteps += nconsume - last = ( - S=S, logps=logps, workspace=ws, tape=tape, t=nconsume, model=prepped - ) - next_update = _parallel_mala_update_progress!( - progress, nsteps, Ntotal, next_update, threshold - ) - end - finally - AbstractMCMC.finish_progress!(progress) - end - end - return last -end - -function _copy_trajectory_rows!( - vals::AbstractMatrix{<:Real}, first_row::Int, S::AbstractMatrix, cols::AbstractRange -) - rows = first_row:(first_row + length(cols) - 1) - S_host = Array(S[:, cols]) - vals[rows, :] .= transpose(S_host) - return vals -end - -function _sample_parallel_mala_chain( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::ParallelMALASampler, - N::Int, - ::Type{FlexiChains.FlexiChain{TKey}}, - sched::ParallelMALASchedule; - initial_params=nothing, - param_names=nothing, - progress=AbstractMCMC.PROGRESS[], - progressname="Sampling", -) where {TKey} - D = model.dim - FP = typeof(sampler.epsilon) - - vals = Matrix{FP}(undef, N, D) - logp = Vector{FP}(undef, N) - - _parallel_mala_run!( - rng, model, sampler, sched; initial_params, progress, progressname - ) do S, logps, cols, first_row - _copy_trajectory_rows!(vals, first_row, S, cols) - logp[first_row:(first_row + length(cols) - 1)] .= view(logps, cols) - return nothing - end - - internals = (logp=logp,) - return _construct_flexichain(TKey, vals, internals, param_names, model) -end - -function _construct_flexichain( - ::Type{TKey}, - vals::AbstractMatrix{<:Real}, - internals::NamedTuple, - param_names::Any, - model::DensityModel, -) where {TKey} - # Wrap user-supplied names in `Parameter`. This allows people to specify, e.g., - # `param_names=(:x, :y, :z=>(2,))` without faffing with `Parameter` themselves. Also - # 'upgrade' symbol parameter names to VarNames if the user requested a VNChain. - to_parameter(vn::VarName) = FlexiChains.Parameter(vn) - to_parameter(s::Symbol) = FlexiChains.Parameter(TKey <: VarName ? VarName{s}() : s) - - N, D = size(vals) - - if param_names === nothing - param_names = model.param_names - end - wrapped_param_names = if param_names === nothing - # Wasn't defined either as `param_names` or `model.param_names`, so make some up - (to_parameter(:x) => (D,),) - else - map(param_names) do n - if n isa Pair - to_parameter(n.first) => n.second - elseif n isa TKey || n isa Symbol - to_parameter(n) - else - throw( - ArgumentError( - "param_names must be a collection of Pairs, Symbols, or $TKey, got $(typeof(n))", - ), - ) - end - end - end - - arr = reshape(vals, N, 1, D) - param_chain = FlexiChains.FlexiChain{TKey}(arr, Tuple(wrapped_param_names)) - - isempty(internals) && return param_chain - - #= - The internals have mixed element types (e.g. `Bool` for `accepted`/`is_warmup`, - `FP` for `logp`), so they cannot share a single stacked array without widening. - Store each in its own column via the dict-of-arrays constructor, which preserves - the natural element type (see issue #49), then merge with the parameters. - =# - extras = OrderedDict{FlexiChains.ParameterOrExtra{<:TKey},AbstractArray}( - FlexiChains.Extra(name) => internals[name] for name in keys(internals) - ) - extra_chain = FlexiChains.FlexiChain{TKey}(N, 1, extras) - - return merge(param_chain, extra_chain) -end - -function _sample_parallel_mala_blocks( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::ParallelMALASampler, - N::Int, - sched::ParallelMALASchedule; - initial_params=nothing, - progress=AbstractMCMC.PROGRESS[], - progressname="Sampling", -) - nblocks = cld(sched.Ntotal, sampler.T) - blocks = Vector{AbstractMatrix}(undef, 0) - logp_blocks = Vector{AbstractVector}(undef, 0) - ends = Int[] - sizehint!(blocks, nblocks) - sizehint!(logp_blocks, nblocks) - sizehint!(ends, nblocks) - - last = _parallel_mala_run!( - rng, model, sampler, sched; initial_params, progress, progressname - ) do S, logps, cols, first_row - isempty(cols) && return nothing - push!(blocks, S[:, cols]) - push!(logp_blocks, logps[cols]) - push!(ends, first_row + length(cols) - 1) - return nothing - end - - S = copy(last.S) - final_state = ParallelMALAState( - S[:, last.t], - last.logps[last.t], - S, - last.logps, - last.workspace, - last.tape, - last.t, - last.model, - ) - return ParallelMALABlockSamples(blocks, logp_blocks, ends, N), final_state -end - -function _default_parallel_mala_mcmcsample( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::ParallelMALASampler, - N::Integer; - kwargs..., -) - return invoke( - AbstractMCMC.mcmcsample, - Tuple{ - Random.AbstractRNG, - AbstractMCMC.AbstractModel, - AbstractMCMC.AbstractSampler, - Integer, - }, - rng, - model, - sampler, - N; - kwargs..., - ) -end - -function AbstractMCMC.mcmcsample( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::ParallelMALASampler, - N::Integer; - progress=AbstractMCMC.PROGRESS[], - progressname="Sampling", - callback=nothing, - num_warmup::Int=0, - discard_initial::Int=num_warmup, - thinning=1, - chain_type::Type=Any, - initial_state=nothing, - initial_params=nothing, - param_names=nothing, - kwargs..., -) - #= `callback` wants a transition per step and `initial_state` a trajectory to - replay, so both go through the generic step loop. =# - if callback !== nothing || initial_state !== nothing - return _default_parallel_mala_mcmcsample( - rng, - model, - sampler, - N; - progress=progress, - progressname=progressname, - callback=callback, - num_warmup=num_warmup, - discard_initial=discard_initial, - thinning=thinning, - chain_type=chain_type, - initial_state=initial_state, - initial_params=initial_params, - param_names=param_names, - kwargs..., - ) - end - - N_int = Int(N) - sched = ParallelMALASchedule(N_int; discard_initial, thinning, num_warmup) - - if chain_type <: FlexiChains.FlexiChain - return _sample_parallel_mala_chain( - rng, - model, - sampler, - N_int, - chain_type, - sched; - initial_params=initial_params, - param_names=param_names, - progress=progress, - progressname=progressname, - ) - end - - samples, state = _sample_parallel_mala_blocks( - rng, - model, - sampler, - N_int, - sched; - initial_params=initial_params, - progress=progress, - progressname=progressname, - ) - chain_type === Any && return samples - - sample_vec = [samples[i] for i in 1:length(samples)] - return AbstractMCMC.bundle_samples( - sample_vec, model, sampler, state, chain_type; param_names=param_names, kwargs... - ) -end - -function AbstractMCMC.step( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::ParallelMALASampler{FP}; - initial_params=nothing, - kwargs..., -) where {FP} - x0 = if initial_params !== nothing - copy(initial_params) - else - randn(rng, FP, model.dim) - end - model = _prepare_model(model, x0, sampler.T, sampler.backend) - - S, tape, ws = _deer_solve_new_tape(rng, model, sampler, x0) - logps = _trajectory_logps(model, S) - x1 = S[:, 1] - logp1 = logps[1] - trans = ParallelMALATransition(x1, logp1) - state = ParallelMALAState(x1, logp1, S, logps, ws, tape, 1, model) - return trans, state -end - -function AbstractMCMC.step( - rng::Random.AbstractRNG, - user_model::DensityModel, - sampler::ParallelMALASampler, - state::ParallelMALAState; - kwargs..., -) - #= A state from a different model (via `initial_state`) brings a trajectory - solved under that model, so the rest of it can't be replayed either: pick up - from the current position with a fresh tape under the model we were handed. =# - reuse = _prepped_for(state.model, user_model) - model = if reuse - state.model - else - _prepare_model(user_model, state.x, sampler.T, sampler.backend) - end - T = sampler.T - t_next = state.t + 1 - - if reuse && t_next <= T - x_new = state.trajectory[:, t_next] - logp_new = state.logps[t_next] - trans = ParallelMALATransition(x_new, logp_new) - new_state = ParallelMALAState( - x_new, - logp_new, - state.trajectory, - state.logps, - state.workspace, - state.tape, - t_next, - model, - ) - return trans, new_state - else - x0 = reuse ? state.trajectory[:, T] : copy(state.x) - S_new, tape, ws = _deer_solve_new_tape( - rng, model, sampler, x0; workspace=state.workspace - ) - logps = _trajectory_logps(model, S_new) - x_new = S_new[:, 1] - logp_new = logps[1] - trans = ParallelMALATransition(x_new, logp_new) - new_state = ParallelMALAState(x_new, logp_new, S_new, logps, ws, tape, 1, model) - return trans, new_state - end -end - -for TKey in (Symbol, VarName) - @eval function AbstractMCMC.bundle_samples( - samples::Vector{<:ParallelMALATransition}, - model::DensityModel, - sampler::ParallelMALASampler, - state::ParallelMALAState, - ::Type{FlexiChains.FlexiChain{$TKey}}; - param_names=nothing, - kwargs..., - ) - N = length(samples) - D = model.dim - FP = typeof(sampler.epsilon) - - vals = Matrix{FP}(undef, N, D) - logp = Vector{FP}(undef, N) - - for i in 1:N - vals[i, :] .= samples[i].x - logp[i] = samples[i].logp - end - - internals = (logp=logp,) - return _construct_flexichain($TKey, vals, internals, param_names, model) - end -end - -""" - AdaptiveMALASampler(epsilon_init; n_warmup, target_accept, gamma, t0, kappa, cholM) -""" -struct AdaptiveMALASampler{FP<:AbstractFloat,CM} <: AbstractMCMC.AbstractSampler - epsilon_init::FP - n_warmup::Int - target_accept::FP - gamma::FP - t0::FP - kappa::FP - cholM::CM -end - -function AdaptiveMALASampler( - epsilon_init::Real; - n_warmup::Int=1000, - target_accept::Real=0.574, - gamma::Real=0.05, - t0::Real=10.0, - kappa::Real=0.75, - cholM=nothing, -) - epsilon_init > 0 || throw(ArgumentError("epsilon_init must be > 0, got $epsilon_init")) - 0 < target_accept < 1 || - throw(ArgumentError("target_accept must be in (0,1), got $target_accept")) - gamma > 0 || throw(ArgumentError("gamma must be > 0, got $gamma")) - t0 > 0 || throw(ArgumentError("t0 must be > 0, got $t0")) - 0.5 < kappa <= 1.0 || throw(ArgumentError("kappa must be in (0.5, 1], got $kappa")) - - eps_f = float(epsilon_init) - FP = typeof(eps_f) - return AdaptiveMALASampler{FP,typeof(cholM)}( - eps_f, n_warmup, FP(target_accept), FP(gamma), FP(t0), FP(kappa), cholM - ) -end - -""" -State for an `AdaptiveMALASampler` chain, including dual-averaging adaptation -statistics. Holds the prepped model so AD preparation happens once per chain. -""" -struct AdaptiveMALAState{ - V<:AbstractVector,FP<:AbstractFloat,W,NV<:AbstractVector,H,DM<:PreppedDensityModel -} - x::V - logp::FP - epsilon::FP - epsilon_bar::FP - H_bar::FP - step::Int - workspace::W - noise::NV - noise_host::H - model::DM -end - -""" -One `AdaptiveMALASampler` sample: parameter vector `x`, its log-density `logp`, -the step size used for that transition, and whether the sample came from warmup. -""" -struct AdaptiveMALATransition{V<:AbstractVector,FP<:AbstractFloat} - x::V - logp::FP - accepted::Bool - step_size::FP - is_warmup::Bool -end - -function _dual_average_update( - epsilon_init::FP, - epsilon_bar::FP, - H_bar::FP, - m::Int, - logα::FP, - sampler::AdaptiveMALASampler{FP}, -) where {FP<:AbstractFloat} - α = min(one(FP), exp(logα)) - δ = sampler.target_accept - γ = sampler.gamma - t0 = sampler.t0 - κ = sampler.kappa - μ = log(10 * epsilon_init) - - inv_mt0 = one(FP) / (FP(m) + t0) - H_bar_new = (one(FP) - inv_mt0) * H_bar + inv_mt0 * (δ - α) - log_ε = μ - sqrt(FP(m)) / γ * H_bar_new - mk = FP(m)^(-κ) - log_ε_bar_new = mk * log_ε + (one(FP) - mk) * log(epsilon_bar) - - return exp(log_ε), exp(log_ε_bar_new), H_bar_new -end - -function AbstractMCMC.step( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::AdaptiveMALASampler{FP}; - initial_params=nothing, - kwargs..., -) where {FP} - x = if initial_params !== nothing - copy(initial_params) - else - randn(rng, FP, model.dim) - end - model = _prepare_model(model, x) - logp_val = FP(model.logdensity(x)) - ws = MALA.MALAWorkspace(x) - noise, noise_host = _make_noise_buffer(x, FP, model.dim) - trans = AdaptiveMALATransition(x, logp_val, true, sampler.epsilon_init, true) - state = AdaptiveMALAState( - x, - logp_val, - sampler.epsilon_init, - sampler.epsilon_init, - zero(FP), - 0, - ws, - noise, - noise_host, - model, - ) - return trans, state -end - -function AbstractMCMC.step( - rng::Random.AbstractRNG, - user_model::DensityModel, - sampler::AdaptiveMALASampler{FP}, - state::AdaptiveMALAState; - kwargs..., -) where {FP} - model = if _prepped_for(state.model, user_model) - state.model - else - _prepare_model(user_model, state.x) - end - D = model.dim - in_warmup = state.step < sampler.n_warmup - ε = in_warmup ? state.epsilon : state.epsilon_bar - - ξ = _randn_like!(rng, state.noise, state.noise_host) - u = rand(rng) - - x_next = similar(state.x) - x_next, accepted, logα = MALA.mala_step_with_logα!( - x_next, - state.workspace, - model.logdensity, - model.grad_logdensity, - state.x, - ε, - ξ, - u; - cholM=sampler.cholM, - ) - - logp_next = accepted ? FP(model.logdensity(x_next)) : state.logp - - m_new = state.step + 1 - ε_new, ε_bar_new, H_bar_new = if in_warmup - _dual_average_update( - sampler.epsilon_init, - state.epsilon_bar, - state.H_bar, - m_new, - FP(logα), - sampler, - ) - else - state.epsilon, state.epsilon_bar, state.H_bar - end - - trans = AdaptiveMALATransition(x_next, logp_next, accepted, ε, in_warmup) - new_state = AdaptiveMALAState( - x_next, - logp_next, - ε_new, - ε_bar_new, - H_bar_new, - m_new, - state.workspace, - state.noise, - state.noise_host, - model, - ) - return trans, new_state -end - -for TKey in (Symbol, VarName) - @eval function AbstractMCMC.bundle_samples( - samples::Vector{<:AdaptiveMALATransition}, - model::DensityModel, - sampler::AdaptiveMALASampler, - state::AdaptiveMALAState, - ::Type{FlexiChains.FlexiChain{$TKey}}; - param_names=nothing, - discard_warmup=false, - kwargs..., - ) - filtered = discard_warmup ? filter(s -> !s.is_warmup, samples) : samples - N = length(filtered) - D = model.dim - FP = typeof(sampler.epsilon_init) - - vals = Matrix{FP}(undef, N, D) - logp = Vector{FP}(undef, N) - accepted = Vector{Bool}(undef, N) - step_size = Vector{FP}(undef, N) - is_warmup = Vector{Bool}(undef, N) - - for i in 1:N - s = filtered[i] - vals[i, :] .= s.x - logp[i] = s.logp - accepted[i] = s.accepted - step_size[i] = s.step_size - is_warmup[i] = s.is_warmup - end - - internals = (logp=logp, accepted=accepted, step_size=step_size, is_warmup=is_warmup) - return _construct_flexichain($TKey, vals, internals, param_names, model) - end -end - -#= `sample` on a `DensityModel` returns a `VNChain` unless `chain_type` says -otherwise. AbstractMCMC's own default is `Any`, which hands back the raw -transition vector; pass `chain_type=Any` to get that. =# -function AbstractMCMC.sample( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::Union{MALASampler,AdaptiveMALASampler,ParallelMALASampler}, - N_or_isdone; - chain_type::Type=VNChain, - kwargs..., -) - return AbstractMCMC.mcmcsample(rng, model, sampler, N_or_isdone; chain_type, kwargs...) -end - -function AbstractMCMC.sample( - rng::Random.AbstractRNG, - model::DensityModel, - sampler::Union{MALASampler,AdaptiveMALASampler,ParallelMALASampler}, - parallel::AbstractMCMC.AbstractMCMCEnsemble, - N::Integer, - nchains::Integer; - chain_type::Type=VNChain, - kwargs..., -) - return AbstractMCMC.mcmcsample( - rng, model, sampler, parallel, N, nchains; chain_type, kwargs... - ) -end diff --git a/src/model.jl b/src/model.jl new file mode 100644 index 0000000..6d13d2e --- /dev/null +++ b/src/model.jl @@ -0,0 +1,422 @@ +""" + DensityModel(logdensity, grad_logdensity, dim; param_names, logdensity_batch, grad_logdensity_batch, hvp, hvp_batch) + +Wrap a log-density and its derivatives for ParallelMCMC samplers. + +Each derivative slot (`grad_logdensity`, `hvp`, `grad_logdensity_batch`, +`hvp_batch`) takes a callable or an `ADTypes.AbstractADType`, so a model can be +built from the log-density alone: + + DensityModel(logp, AutoForwardDiff(), dim) + +Backends are prepared when sampling starts. An HVP backend differentiates a +callable gradient once. With an AD-derived gradient it forms +`DifferentiationInterface.SecondOrder(hvp_backend, grad_backend)` on the +log-density. An explicit `SecondOrder` also bypasses the gradient slot. + +- `logdensity(x::AbstractVector) -> Real` +- `grad_logdensity` — callable `x -> AbstractVector`, or an AD backend. +- `hvp` — optional callable `(x, v) -> AbstractVector`, or a backend. If + `nothing`, DEER builds the HVP from the sampler's `backend`. +- `logdensity_batch(X::AbstractMatrix) -> AbstractVector` — optional batched + log-density over independent columns. +- `grad_logdensity_batch` — optional callable `X -> AbstractMatrix`, or a + backend. It is derived from `logdensity_batch` when `grad_logdensity` is a + backend. +- `hvp_batch` — optional callable `(X, V) -> AbstractMatrix`, or a backend, + resolved like `hvp`. +- `dim::Int` — dimensionality of the parameter space +- `param_names` — optional collection of parameter names used in `FlexiChains` output. If + `nothing` (the default), uses a single vector-valued parameter `:x` with shape `(dim,)`. + See the [`Parameter names`](@ref parameter-names) section of the docs for more + information. + +Both batched derivative slots require `logdensity_batch`. `AutoReactant()` has +additional pairing and tracing constraints; see the GPU guide. +""" +struct DensityModel{F,G,H,FB,GB,HB,PN} <: AbstractMCMC.AbstractModel + logdensity::F + grad_logdensity::G + hvp::H + logdensity_batch::FB + grad_logdensity_batch::GB + hvp_batch::HB + dim::Int + param_names::PN +end + +function DensityModel( + logdensity, + grad_logdensity, + dim::Int; + param_names=nothing, + hvp=nothing, + logdensity_batch=nothing, + grad_logdensity_batch=nothing, + hvp_batch=nothing, +) + logdensity isa AbstractADType && + throw(ArgumentError("logdensity must be callable, not an AD backend")) + logdensity_batch isa AbstractADType && + throw(ArgumentError("logdensity_batch must be callable, not an AD backend")) + grad_logdensity === nothing && throw( + ArgumentError( + "grad_logdensity must be a callable or an ADTypes.AbstractADType backend" + ), + ) + _batch_needs_logp(name) = throw(ArgumentError("$name requires logdensity_batch")) + if logdensity_batch === nothing + grad_logdensity_batch === nothing || _batch_needs_logp("grad_logdensity_batch") + hvp_batch === nothing || _batch_needs_logp("hvp_batch") + end + return DensityModel( + logdensity, + grad_logdensity, + hvp, + logdensity_batch, + grad_logdensity_batch, + hvp_batch, + dim, + param_names, + ) +end + +""" +A [`DensityModel`](@ref) whose AD backends have been resolved to callables. +`source` prevents a cached preparation from being reused with another model. +""" +struct PreppedDensityModel{F,G,H,FB,GB,HB,PN,SM<:DensityModel} + logdensity::F + grad_logdensity::G + hvp::H + logdensity_batch::FB + grad_logdensity_batch::GB + hvp_batch::HB + dim::Int + param_names::PN + source::SM +end + +# Identity is required: reusing another model's preparation would sample the wrong target. +_prepped_for(prepped::PreppedDensityModel, model::DensityModel) = prepped.source === model + +# DI keys preparations on function identity. Other input types use an unprepared call. +struct _ADGradient{F,B<:AbstractADType,P,TX} + logdensity::F + backend::B + prep::P +end +function (g::_ADGradient{F,B,P,TX})(x) where {F,B,P,TX} + if x isa TX + return DI.gradient(g.logdensity, g.prep, g.backend, x) + else + return DI.gradient(g.logdensity, g.backend, x) + end +end + +# Columns of X are independent samples, so ∇_X sum(logdensity_batch(X)) +# stacks the per-column gradients. +struct _BatchLogdensitySum{F} + logdensity_batch::F +end +(c::_BatchLogdensitySum)(X) = sum(c.logdensity_batch(X)) + +struct _ADGradientBatch{C<:_BatchLogdensitySum,B<:AbstractADType,P,TX} + closure::C + backend::B + prep::P +end +function (g::_ADGradientBatch{C,B,P,TX})(X) where {C,B,P,TX} + if X isa TX + return DI.gradient(g.closure, g.prep, g.backend, X) + else + return DI.gradient(g.closure, g.backend, X) + end +end + +function _resolve_gradient(logdensity, backend::AbstractADType, x_template::AbstractVector) + prep = DI.prepare_gradient(logdensity, backend, x_template) + return _ADGradient{typeof(logdensity),typeof(backend),typeof(prep),typeof(x_template)}( + logdensity, backend, prep + ) +end + +function _resolve_gradient_batch( + logdensity_batch, backend::AbstractADType, X_template::AbstractMatrix +) + closure = _BatchLogdensitySum(logdensity_batch) + prep = DI.prepare_gradient(closure, backend, X_template) + return _ADGradientBatch{typeof(closure),typeof(backend),typeof(prep),typeof(X_template)}( + closure, backend, prep + ) +end + +function _resolve_gradient( + logdensity, backend::ADTypes.AutoReactant, x_template::AbstractVector +) + return _reactant_resolve_gradient(logdensity, backend, x_template) +end +function _reactant_resolve_gradient(logdensity, backend, x_template) + return error(_REACTANT_LOAD_HINT) +end + +function _resolve_gradient_batch( + logdensity_batch, backend::ADTypes.AutoReactant, X_template::AbstractMatrix +) + return _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) +end +function _reactant_resolve_gradient_batch(logdensity_batch, backend, X_template) + return error(_REACTANT_LOAD_HINT) +end + +# An explicit SecondOrder differentiates logdensity and bypasses the gradient slot. +function _resolve_hvp(logdensity, grad, grad_backend, hvp_backend, x_template) + if hvp_backend isa DI.SecondOrder + return DEER._make_hvp_fn_second_order(logdensity, hvp_backend, x_template) + elseif grad_backend !== nothing + return DEER._make_hvp_fn_second_order( + logdensity, _second_order(hvp_backend, grad_backend), x_template + ) + else + return DEER._make_hvp_fn( + DEER._hvp_strategy(hvp_backend), grad, hvp_backend, x_template + ) + end +end + +function _resolve_hvp_batch( + logdensity_batch, grad_batch, grad_batch_backend, hvp_backend, X_template +) + if hvp_backend isa DI.SecondOrder + return DEER._make_hvp_batch_fn_second_order( + _BatchLogdensitySum(logdensity_batch), hvp_backend, X_template + ) + elseif grad_batch_backend !== nothing + return DEER._make_hvp_batch_fn_second_order( + _BatchLogdensitySum(logdensity_batch), + _second_order(hvp_backend, grad_batch_backend), + X_template, + ) + else + return DEER._make_hvp_batch_fn( + DEER._hvp_strategy(hvp_backend), grad_batch, hvp_backend, X_template + ) + end +end + +# With no HVP source, run forward mode over an AD-derived gradient. +# AutoReactant gradient maps to itself, so the Reactant pairing still holds. +_default_hvp_backend(::Nothing) = nothing +_default_hvp_backend(grad_backend::AbstractADType) = DI.forward_counterpart(grad_backend) + +# Reactant implements the two passes itself; DI composes all other backends. +_second_order(hvp_backend, grad_backend) = DI.SecondOrder(hvp_backend, grad_backend) +function _second_order(hvp_backend::ADTypes.AutoReactant, ::ADTypes.AutoReactant) + return hvp_backend +end + +_check_reactant_pair(grad_backend, hvp_backend) = nothing +_check_reactant_pair(::ADTypes.AutoReactant, ::ADTypes.AutoReactant) = nothing +_check_reactant_pair(::Nothing, ::ADTypes.AutoReactant) = nothing + +function _check_reactant_pair(grad_backend::ADTypes.AutoReactant, hvp_backend) + return throw( + ArgumentError( + "an AutoReactant gradient requires an AutoReactant HVP backend; " * + "got $(hvp_backend)", + ), + ) +end + +function _check_reactant_pair(grad_backend, hvp_backend::ADTypes.AutoReactant) + return throw( + ArgumentError( + "an AutoReactant HVP requires an AutoReactant or callable gradient; " * + "got backend $(grad_backend)", + ), + ) +end + +function _check_reactant_pair(::ADTypes.AutoReactant, hvp_backend::DI.SecondOrder) + return throw( + ArgumentError( + "an AutoReactant gradient requires a bare AutoReactant HVP backend; " * + "got $(hvp_backend)", + ), + ) +end + +function _check_reactant_pair(grad_backend, hvp_backend::DI.SecondOrder) + _second_order_has_reactant(hvp_backend) || return nothing + return throw( + ArgumentError( + "AutoReactant cannot be used inside DifferentiationInterface.SecondOrder; " * + "use a bare AutoReactant() HVP backend", + ), + ) +end + +function _second_order_has_reactant(so::DI.SecondOrder) + return DI.outer(so) isa ADTypes.AutoReactant || DI.inner(so) isa ADTypes.AutoReactant +end + +_check_reactant_hvp_source(grad, hvp_backend) = nothing + +""" + _prepare_model(model, x_template) -> PreppedDensityModel + _prepare_model(model, x_template, T::Int, backend) -> PreppedDensityModel + +Resolve AD backends into prepared callables at `x_template`. The two-argument +form prepares only the gradient. The four-argument form also prepares DEER's +HVP and batched slots at shape `(model.dim, T)`. +""" +function _prepare_model(model::DensityModel, x_template::AbstractVector) + grad = if model.grad_logdensity isa AbstractADType + _resolve_gradient(model.logdensity, model.grad_logdensity, x_template) + else + model.grad_logdensity + end + return PreppedDensityModel( + model.logdensity, + grad, + nothing, + model.logdensity_batch, + nothing, + nothing, + model.dim, + model.param_names, + model, + ) +end + +function _prepare_model(model::DensityModel, x_template::AbstractVector, T::Int, backend) + grad_backend = + model.grad_logdensity isa AbstractADType ? model.grad_logdensity : nothing + + # Validate the pair before an AutoReactant gradient triggers compilation. + needs_hvp = model.hvp === nothing || model.hvp isa AbstractADType + hvp_backend = if needs_hvp + hb = if model.hvp !== nothing + model.hvp + elseif backend !== nothing + backend + else + _default_hvp_backend(grad_backend) + end + hb === nothing && throw( + ArgumentError( + "ParallelMALASampler needs a Hessian-vector product: supply `hvp` " * + "on the DensityModel (callable or AD backend), pass `backend=` " * + "to ParallelMALASampler, or give `grad_logdensity` as an AD backend", + ), + ) + _check_reactant_pair(grad_backend, hb) + _check_reactant_hvp_source(model.grad_logdensity, hb) + hb + else + nothing + end + + grad = if grad_backend !== nothing + _resolve_gradient(model.logdensity, grad_backend, x_template) + else + model.grad_logdensity + end + + hvp = if needs_hvp + _resolve_hvp(model.logdensity, grad, grad_backend, hvp_backend, x_template) + else + model.hvp + end + + grad_batch = model.grad_logdensity_batch + if grad_batch === nothing && model.logdensity_batch !== nothing + grad_batch = grad_backend + end + grad_batch_backend = grad_batch isa AbstractADType ? grad_batch : nothing + + hvp_batch = model.hvp_batch + batch_active = model.logdensity_batch !== nothing && grad_batch !== nothing + + if batch_active + # Prepare on x0 in every column; zeros need not be in the support. + X_template = similar(x_template, length(x_template), T) + X_template .= x_template + + needs_hvp_batch = hvp_batch === nothing || hvp_batch isa AbstractADType + hvp_batch_backend = if needs_hvp_batch + hbb = if hvp_batch !== nothing + hvp_batch + elseif model.hvp isa AbstractADType + model.hvp + elseif backend !== nothing + backend + else + _default_hvp_backend(grad_batch_backend) + end + hbb === nothing && throw( + ArgumentError("batched DEER requires `hvp_batch` or a sampler `backend`"), + ) + _check_reactant_pair(grad_batch_backend, hbb) + hbb + else + nothing + end + + if grad_batch_backend !== nothing + grad_batch = _resolve_gradient_batch( + model.logdensity_batch, grad_batch_backend, X_template + ) + end + + if needs_hvp_batch + hvp_batch = _resolve_hvp_batch( + model.logdensity_batch, + grad_batch, + grad_batch_backend, + hvp_batch_backend, + X_template, + ) + end + elseif hvp_batch !== nothing + # Do not silently ignore an explicitly supplied HVP. + throw( + ArgumentError( + "hvp_batch has no batched gradient; supply grad_logdensity_batch or an AD-backed grad_logdensity", + ), + ) + end + + return PreppedDensityModel( + model.logdensity, + grad, + hvp, + model.logdensity_batch, + grad_batch, + hvp_batch, + model.dim, + model.param_names, + model, + ) +end + +#= Wrappers for LogDensityProblems objects. They are defined here rather than in +LogDensityProblemsExt so `_check_reactant_hvp_source` can dispatch on them. The call +methods live in the extension. =# +struct LogDensityProblemPrimal{L} + ld::L +end +struct LogDensityProblemGradient{L} + ld::L +end + +function _check_reactant_hvp_source( + ::LogDensityProblemGradient, hvp_backend::ADTypes.AutoReactant +) + return throw( + ArgumentError( + "AutoReactant cannot differentiate a LogDensityProblems-derived gradient; " * + "supply a Reactant-traceable gradient callable instead", + ), + ) +end diff --git a/src/sample.jl b/src/sample.jl new file mode 100644 index 0000000..177925f --- /dev/null +++ b/src/sample.jl @@ -0,0 +1,28 @@ +#= `sample` on a `DensityModel` returns a `VNChain` unless `chain_type` says +otherwise. AbstractMCMC's own default is `Any`, which hands back the raw +transition vector; pass `chain_type=Any` to get that. =# +function AbstractMCMC.sample( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::Union{MALASampler,AdaptiveMALASampler,ParallelMALASampler}, + N_or_isdone; + chain_type::Type=VNChain, + kwargs..., +) + return AbstractMCMC.mcmcsample(rng, model, sampler, N_or_isdone; chain_type, kwargs...) +end + +function AbstractMCMC.sample( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::Union{MALASampler,AdaptiveMALASampler,ParallelMALASampler}, + parallel::AbstractMCMC.AbstractMCMCEnsemble, + N::Integer, + nchains::Integer; + chain_type::Type=VNChain, + kwargs..., +) + return AbstractMCMC.mcmcsample( + rng, model, sampler, parallel, N, nchains; chain_type, kwargs... + ) +end diff --git a/src/samplers/adaptive_mala.jl b/src/samplers/adaptive_mala.jl new file mode 100644 index 0000000..5157c77 --- /dev/null +++ b/src/samplers/adaptive_mala.jl @@ -0,0 +1,222 @@ +""" + AdaptiveMALASampler(epsilon_init; n_warmup, target_accept, gamma, t0, kappa, cholM) +""" +struct AdaptiveMALASampler{FP<:AbstractFloat,CM} <: AbstractMCMC.AbstractSampler + epsilon_init::FP + n_warmup::Int + target_accept::FP + gamma::FP + t0::FP + kappa::FP + cholM::CM +end + +function AdaptiveMALASampler( + epsilon_init::Real; + n_warmup::Int=1000, + target_accept::Real=0.574, + gamma::Real=0.05, + t0::Real=10.0, + kappa::Real=0.75, + cholM=nothing, +) + epsilon_init > 0 || throw(ArgumentError("epsilon_init must be > 0, got $epsilon_init")) + 0 < target_accept < 1 || + throw(ArgumentError("target_accept must be in (0,1), got $target_accept")) + gamma > 0 || throw(ArgumentError("gamma must be > 0, got $gamma")) + t0 > 0 || throw(ArgumentError("t0 must be > 0, got $t0")) + 0.5 < kappa <= 1.0 || throw(ArgumentError("kappa must be in (0.5, 1], got $kappa")) + + eps_f = float(epsilon_init) + FP = typeof(eps_f) + return AdaptiveMALASampler{FP,typeof(cholM)}( + eps_f, n_warmup, FP(target_accept), FP(gamma), FP(t0), FP(kappa), cholM + ) +end + +""" +State for an `AdaptiveMALASampler` chain, including dual-averaging adaptation +statistics. Holds the prepped model so AD preparation happens once per chain. +""" +struct AdaptiveMALAState{ + V<:AbstractVector,FP<:AbstractFloat,W,NV<:AbstractVector,H,DM<:PreppedDensityModel +} + x::V + logp::FP + epsilon::FP + epsilon_bar::FP + H_bar::FP + step::Int + workspace::W + noise::NV + noise_host::H + model::DM +end + +""" +One `AdaptiveMALASampler` sample: parameter vector `x`, its log-density `logp`, +the step size used for that transition, and whether the sample came from warmup. +""" +struct AdaptiveMALATransition{V<:AbstractVector,FP<:AbstractFloat} + x::V + logp::FP + accepted::Bool + step_size::FP + is_warmup::Bool +end + +function _dual_average_update( + epsilon_init::FP, + epsilon_bar::FP, + H_bar::FP, + m::Int, + logα::FP, + sampler::AdaptiveMALASampler{FP}, +) where {FP<:AbstractFloat} + α = min(one(FP), exp(logα)) + δ = sampler.target_accept + γ = sampler.gamma + t0 = sampler.t0 + κ = sampler.kappa + μ = log(10 * epsilon_init) + + inv_mt0 = one(FP) / (FP(m) + t0) + H_bar_new = (one(FP) - inv_mt0) * H_bar + inv_mt0 * (δ - α) + log_ε = μ - sqrt(FP(m)) / γ * H_bar_new + mk = FP(m)^(-κ) + log_ε_bar_new = mk * log_ε + (one(FP) - mk) * log(epsilon_bar) + + return exp(log_ε), exp(log_ε_bar_new), H_bar_new +end + +function AbstractMCMC.step( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::AdaptiveMALASampler{FP}; + initial_params=nothing, + kwargs..., +) where {FP} + x = if initial_params !== nothing + copy(initial_params) + else + randn(rng, FP, model.dim) + end + model = _prepare_model(model, x) + logp_val = FP(model.logdensity(x)) + ws = MALA.MALAWorkspace(x) + noise, noise_host = _make_noise_buffer(x, FP, model.dim) + trans = AdaptiveMALATransition(x, logp_val, true, sampler.epsilon_init, true) + state = AdaptiveMALAState( + x, + logp_val, + sampler.epsilon_init, + sampler.epsilon_init, + zero(FP), + 0, + ws, + noise, + noise_host, + model, + ) + return trans, state +end + +function AbstractMCMC.step( + rng::Random.AbstractRNG, + user_model::DensityModel, + sampler::AdaptiveMALASampler{FP}, + state::AdaptiveMALAState; + kwargs..., +) where {FP} + model = if _prepped_for(state.model, user_model) + state.model + else + _prepare_model(user_model, state.x) + end + D = model.dim + in_warmup = state.step < sampler.n_warmup + ε = in_warmup ? state.epsilon : state.epsilon_bar + + ξ = _randn_like!(rng, state.noise, state.noise_host) + u = rand(rng) + + x_next = similar(state.x) + x_next, accepted, logα = MALA.mala_step_with_logα!( + x_next, + state.workspace, + model.logdensity, + model.grad_logdensity, + state.x, + ε, + ξ, + u; + cholM=sampler.cholM, + ) + + logp_next = accepted ? FP(model.logdensity(x_next)) : state.logp + + m_new = state.step + 1 + ε_new, ε_bar_new, H_bar_new = if in_warmup + _dual_average_update( + sampler.epsilon_init, + state.epsilon_bar, + state.H_bar, + m_new, + FP(logα), + sampler, + ) + else + state.epsilon, state.epsilon_bar, state.H_bar + end + + trans = AdaptiveMALATransition(x_next, logp_next, accepted, ε, in_warmup) + new_state = AdaptiveMALAState( + x_next, + logp_next, + ε_new, + ε_bar_new, + H_bar_new, + m_new, + state.workspace, + state.noise, + state.noise_host, + model, + ) + return trans, new_state +end + +for TKey in (Symbol, VarName) + @eval function AbstractMCMC.bundle_samples( + samples::Vector{<:AdaptiveMALATransition}, + model::DensityModel, + sampler::AdaptiveMALASampler, + state::AdaptiveMALAState, + ::Type{FlexiChains.FlexiChain{$TKey}}; + param_names=nothing, + discard_warmup=false, + kwargs..., + ) + filtered = discard_warmup ? filter(s -> !s.is_warmup, samples) : samples + N = length(filtered) + D = model.dim + FP = typeof(sampler.epsilon_init) + + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) + accepted = Vector{Bool}(undef, N) + step_size = Vector{FP}(undef, N) + is_warmup = Vector{Bool}(undef, N) + + for i in 1:N + s = filtered[i] + vals[i, :] .= s.x + logp[i] = s.logp + accepted[i] = s.accepted + step_size[i] = s.step_size + is_warmup[i] = s.is_warmup + end + + internals = (logp=logp, accepted=accepted, step_size=step_size, is_warmup=is_warmup) + return _construct_flexichain($TKey, vals, internals, param_names, model) + end +end diff --git a/src/samplers/mala.jl b/src/samplers/mala.jl new file mode 100644 index 0000000..93fbe85 --- /dev/null +++ b/src/samplers/mala.jl @@ -0,0 +1,156 @@ +""" + MALASampler(epsilon; cholM=nothing) + +Metropolis-Adjusted Langevin Algorithm sampler with step size `epsilon`. + +Optionally pass `cholM = cholesky(M)` to use a mass matrix `M` as a +preconditioner. The proposal becomes `y = x + ε M ∇logp(x) + √(2ε) L ξ` +where `L` is the Cholesky factor of `M`. +""" +struct MALASampler{FP<:AbstractFloat,CM} <: AbstractMCMC.AbstractSampler + epsilon::FP + cholM::CM +end + +function MALASampler(epsilon::Real; cholM=nothing) + epsilon > 0 || throw(ArgumentError("epsilon must be > 0, got $epsilon")) + eps_f = float(epsilon) + return MALASampler{typeof(eps_f),typeof(cholM)}(eps_f, cholM) +end + +""" +State for a `MALASampler` chain. Holds the prepped model so AD preparation +happens once per chain. +""" +struct MALAState{V<:AbstractVector,L<:Real,W,NV<:AbstractVector,H,DM<:PreppedDensityModel} + x::V + logp::L + workspace::W + noise::NV + noise_host::H + model::DM +end + +""" +One `MALASampler` sample: parameter vector `x`, its log-density `logp`, and an +accept/reject flag. +""" +struct MALATransition{V<:AbstractVector,L<:Real} + x::V + logp::L + accepted::Bool +end + +function _make_noise_buffer(x::AbstractVector, ::Type{FP}, D::Int) where {FP} + ξ = similar(x, FP, D) + host = needs_host_staging(ξ) ? Vector{FP}(undef, D) : nothing + return ξ, host +end + +function _randn_like!( + rng::Random.AbstractRNG, ξ::AbstractVector{FP}, host::Union{Nothing,AbstractVector{FP}} +) where {FP} + if needs_host_staging(ξ) + host === nothing && + error("normal noise for $(typeof(ξ)) requires a reusable host buffer") + randn!(rng, host) + copyto!(ξ, host) + else + randn!(rng, ξ) + end + return ξ +end + +function AbstractMCMC.step( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::MALASampler{FP}; + initial_params=nothing, + kwargs..., +) where {FP} + x = if initial_params !== nothing + copy(initial_params) + else + randn(rng, FP, model.dim) + end + model = _prepare_model(model, x) + logp_val = model.logdensity(x) + ws = MALA.MALAWorkspace(x) + noise, noise_host = _make_noise_buffer(x, FP, model.dim) + t = MALATransition(x, logp_val, true) + s = MALAState(x, logp_val, ws, noise, noise_host, model) + return t, s +end + +function AbstractMCMC.step( + rng::Random.AbstractRNG, + user_model::DensityModel, + sampler::MALASampler, + state::MALAState; + kwargs..., +) + #= Reuse the state's preparation, unless the state came from a different + model (via `initial_state`), in which case the model we were handed is the + one to sample. =# + model = if _prepped_for(state.model, user_model) + state.model + else + _prepare_model(user_model, state.x) + end + x = state.x + ϵ = sampler.epsilon + D = model.dim + + ξ = _randn_like!(rng, state.noise, state.noise_host) + u = rand(rng) + + x_next = similar(x) + x_next, accepted, _ = MALA.mala_step_with_logα!( + x_next, + state.workspace, + model.logdensity, + model.grad_logdensity, + x, + ϵ, + ξ, + u; + cholM=sampler.cholM, + ) + + logp_val = accepted ? model.logdensity(x_next) : state.logp + t = MALATransition(x_next, logp_val, accepted) + s = MALAState(x_next, logp_val, state.workspace, state.noise, state.noise_host, model) + return t, s +end + +for TKey in (Symbol, VarName) + @eval function AbstractMCMC.bundle_samples( + samples::Vector{<:MALATransition}, + model::DensityModel, + sampler::MALASampler, + state::MALAState, + ::Type{FlexiChains.FlexiChain{$TKey}}; + param_names=nothing, + kwargs..., + ) + N = length(samples) + D = model.dim + # Follow the sampler's working precision instead of pinning `Float64`, so a + # Float32 (e.g. GPU) run yields a Float32 chain rather than a silently widened one. + FP = typeof(sampler.epsilon) + + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) + accepted = Vector{Bool}(undef, N) + + for i in 1:N + s = samples[i] + vals[i, :] .= s.x + logp[i] = s.logp + accepted[i] = s.accepted + end + + internals = (logp=logp, accepted=accepted) + return _construct_flexichain($TKey, vals, internals, param_names, model) + end +end diff --git a/src/samplers/parallel_mala.jl b/src/samplers/parallel_mala.jl new file mode 100644 index 0000000..b430004 --- /dev/null +++ b/src/samplers/parallel_mala.jl @@ -0,0 +1,668 @@ +""" + MALATapeElement(ξ, u) + +One element of the MALA noise tape: a noise vector `ξ ~ N(0,I)` and a uniform +scalar `u ~ Uniform(0,1)`. Stored with a concrete vector type `V` for type +stability inside `DEER.TapedRecursion`. +""" +struct MALATapeElement{FP<:AbstractFloat,V<:AbstractVector{FP}} + ξ::V + u::FP +end + +""" + ParallelMALASampler(epsilon; T, maxiter, tol_abs, tol_rel, jacobian, damping, probes, cholM, backend) + +DEER-parallelized MALA sampler. + +Supported Jacobian modes are `:stoch_diag` (the default Hutchinson diagonal +estimator) and `:diag` (exact diagonal via `D` JVPs). + +`backend` supplies Hessian-vector products when the `DensityModel` brings no +`hvp` / `hvp_batch` of its own. It does not supply gradients. Leave it out when +the model carries its own HVPs. When it is left out and `grad_logdensity` is an +AD backend, the HVP runs the forward-mode counterpart of that backend over the +gradient (`DifferentiationInterface.forward_counterpart`). + +With `backend=ADTypes.AutoReactant()`, `grad_logdensity` must be `AutoReactant()` +or a callable. + +`callback` and `initial_state` route through the generic step-wise loop instead. +""" +struct ParallelMALASampler{FP<:AbstractFloat,CM,AD} <: AbstractMCMC.AbstractSampler + epsilon::FP + T::Int + maxiter::Int + tol_abs::FP + tol_rel::FP + jacobian::Symbol + damping::FP + probes::Int + cholM::CM + backend::AD +end + +function ParallelMALASampler( + epsilon::Real; + T::Int=64, + maxiter::Int=200, + tol_abs::Real=1e-6, + tol_rel::Real=1e-5, + jacobian::Symbol=:stoch_diag, + damping::Real=0.5, + probes::Int=1, + cholM=nothing, + backend=nothing, +) + epsilon > 0 || throw(ArgumentError("epsilon must be > 0, got $epsilon")) + (jacobian === :stoch_diag || jacobian === :diag) || + throw(ArgumentError("jacobian must be :stoch_diag or :diag")) + eps_f = float(epsilon) + FP = typeof(eps_f) + return ParallelMALASampler{FP,typeof(cholM),typeof(backend)}( + eps_f, + T, + maxiter, + FP(tol_abs), + FP(tol_rel), + jacobian, + FP(damping), + probes, + cholM, + backend, + ) +end + +""" +State for a `ParallelMALASampler` chain. Holds the prepped model so AD +preparation happens once per chain. +""" +struct ParallelMALAState{ + V<:AbstractVector,L<:Real,M<:AbstractMatrix,LV<:AbstractVector,W,DM<:PreppedDensityModel +} + x::V + logp::L + trajectory::M + logps::LV + workspace::W + tape::Vector{<:MALATapeElement} + t::Int + model::DM +end + +""" +One `ParallelMALASampler` sample: parameter vector `x` and its log-density `logp`. +""" +struct ParallelMALATransition{V<:AbstractVector,L<:Real} + x::V + logp::L +end + +#= Kept samples of a chain, one block of columns per DEER solve. `ends[b]` is the +number of kept samples up to and including block `b`; with thinning the blocks +are not all the same width. =# +struct ParallelMALABlockSamples{B<:AbstractVector,L<:AbstractVector} <: + AbstractVector{ParallelMALATransition} + blocks::B + logps::L + ends::Vector{Int} + n::Int +end + +Base.size(samples::ParallelMALABlockSamples) = (samples.n,) +Base.length(samples::ParallelMALABlockSamples) = samples.n +Base.IndexStyle(::Type{<:ParallelMALABlockSamples}) = IndexLinear() + +function Base.getindex(samples::ParallelMALABlockSamples, i::Int) + 1 <= i <= samples.n || throw(BoundsError(samples, i)) + b = searchsortedfirst(samples.ends, i) + t = i - (b == 1 ? 0 : samples.ends[b - 1]) + return ParallelMALATransition(samples.blocks[b][:, t], samples.logps[b][t]) +end + +function _make_mala_tape_block( + rng::Random.AbstractRNG, x0::AbstractVector, ::Type{FP}, D::Int, T::Int +) where {FP} + Xi = similar(x0, FP, D, T) + U_host = Vector{FP}(undef, T) + + if needs_host_staging(Xi) + Xi_host = Matrix{FP}(undef, D, T) + for t in 1:T + randn!(rng, view(Xi_host, :, t)) + U_host[t] = FP(rand(rng)) + end + copyto!(Xi, Xi_host) + else + for t in 1:T + randn!(rng, view(Xi, :, t)) + U_host[t] = FP(rand(rng)) + end + end + + U = similar(x0, FP, T) + copyto!(U, U_host) + tape = [MALATapeElement(view(Xi, :, t), U_host[t]) for t in 1:T] + return tape, Xi, U +end + +function _build_mala_deer_rec( + model::PreppedDensityModel, + ε::Real, + tape::Vector{<:MALATapeElement}, + x0_like::AbstractVector; + cholM=nothing, + tape_noise=nothing, + tape_uniforms=nothing, +) + logp = model.logdensity + gradlogp = model.grad_logdensity + + hvp_fn = model.hvp + + # Exact forward step. + step_fwd = + (x, te) -> MALA.mala_step_taped(logp, gradlogp, x, ε, te.ξ, te.u; cholM=cholM) + + # Explicit analytical JVP of the surrogate step. + jvp = + (x, te, v) -> MALA.mala_step_surrogate_sigmoid_jvp( + logp, gradlogp, x, ε, te.ξ, te.u, v, hvp_fn; cholM=cholM + ) + + # Fused forward step + JVP. + fwd_and_jvp = + (x, te, v) -> MALA.mala_step_taped_and_jvp( + logp, gradlogp, x, ε, te.ξ, te.u, v, hvp_fn; cholM=cholM + ) + + fwd_and_jvp_batch = + if model.logdensity_batch !== nothing && model.grad_logdensity_batch !== nothing + D = length(x0_like) + T = length(tape) + + (tape_noise === nothing) == (tape_uniforms === nothing) || throw( + ArgumentError("tape_noise and tape_uniforms must be provided together") + ) + Xi, U = if tape_noise === nothing + Xi_local = similar(x0_like, D, T) + U_host = Vector{typeof(tape[1].u)}(undef, T) + U_local = similar(x0_like, typeof(tape[1].u), T) + + for t in 1:T + copyto!(Xi_local, (t - 1) * D + 1, tape[t].ξ, 1, D) + U_host[t] = tape[t].u + end + copyto!(U_local, U_host) + Xi_local, U_local + else + size(tape_noise) == (D, T) || + throw(DimensionMismatch("tape_noise must have size (D, T)")) + length(tape_uniforms) == T || + throw(DimensionMismatch("tape_uniforms must have length T")) + tape_noise, tape_uniforms + end + + X_template = similar(x0_like, D, T) + fill!(X_template, zero(eltype(X_template))) + batch_ws = MALA.MALABatchedWorkspace(X_template) + FT = similar(X_template) + Jt = similar(X_template) + hvp_batch = model.hvp_batch + + (Xbar, Z) -> MALA.mala_step_batched_fwd_and_jvp!( + FT, + Jt, + batch_ws, + model.logdensity_batch, + model.grad_logdensity_batch, + hvp_batch, + Xbar, + ε, + Xi, + U, + Z; + cholM=cholM, + ) + else + nothing + end + + return DEER.TapedRecursion( + step_fwd, jvp, tape; fwd_and_jvp=fwd_and_jvp, fwd_and_jvp_batch=fwd_and_jvp_batch + ) +end + +function _deer_solve_new_tape( + rng::Random.AbstractRNG, + model::PreppedDensityModel, + sampler::ParallelMALASampler, + x0::AbstractVector; + workspace=nothing, +) + D = model.dim + T = sampler.T + FP = typeof(sampler.epsilon) + + tape, Xi, U = _make_mala_tape_block(rng, x0, FP, D, T) + + rec = _build_mala_deer_rec( + model, + sampler.epsilon, + tape, + x0; + cholM=sampler.cholM, + tape_noise=Xi, + tape_uniforms=U, + ) + + ws = workspace === nothing ? DEER.DEERWorkspace(x0, T) : workspace + + S = DEER.solve( + rec, + x0; + tol_abs=sampler.tol_abs, + tol_rel=sampler.tol_rel, + maxiter=sampler.maxiter, + jacobian=sampler.jacobian, + damping=sampler.damping, + probes=sampler.probes, + rng=rng, + workspace=ws, + copy_result=false, + ) + return S, tape, ws +end + +function _trajectory_logps(model::PreppedDensityModel, S::AbstractMatrix) + if model.logdensity_batch !== nothing + return Array(model.logdensity_batch(S)) + end + + T = size(S, 2) + return [model.logdensity(S[:, t]) for t in 1:T] +end + +function _parallel_mala_initial_x( + rng::Random.AbstractRNG, model::DensityModel, ::ParallelMALASampler{FP}, initial_params +) where {FP} + return initial_params !== nothing ? copy(initial_params) : randn(rng, FP, model.dim) +end + +function _parallel_mala_progress(progress, progressname) + progress === true && return AbstractMCMC.CreateNewProgressBar(progressname) + progress === false && return AbstractMCMC.NoLogging() + return progress +end + +function _parallel_mala_update_progress!( + progress, nsteps::Int, Ntotal::Int, next_update::Real, threshold::Real +) + if nsteps >= next_update && next_update <= Ntotal + AbstractMCMC.update_progress!(progress, min(nsteps / Ntotal, 1)) + next_update += threshold + end + return next_update +end + +""" +Which steps of a `ParallelMALASampler` run are kept. Global step `g` is kept iff +`g > discard_initial` and `(g - discard_initial - 1) % thinning == 0`, i.e., the +steps the `AbstractMCMC.mcmcsample` loop saves, so the batched path and the +step-wise fallback return the same chain for the same RNG. +""" +struct ParallelMALASchedule + discard_initial::Int + thinning::Int + Ntotal::Int +end + +function ParallelMALASchedule( + N::Int; discard_initial::Int=0, thinning::Int=1, num_warmup::Int=0 +) + N > 0 || error("the number of samples must be ≥ 1") + discard_initial >= 0 || + throw(ArgumentError("number of discarded samples must be non-negative")) + thinning >= 1 || throw(ArgumentError("thinning must be ≥ 1")) + num_warmup >= 0 || + throw(ArgumentError("number of warm-up samples must be non-negative")) + Ntotal = thinning * (N - 1) + discard_initial + 1 + Ntotal >= num_warmup || throw( + ArgumentError("number of warm-up samples exceeds the total number of samples") + ) + return ParallelMALASchedule(discard_initial, thinning, Ntotal) +end + +#= Columns of a block covering global steps `offset+1:offset+ncols` that the schedule +keeps, and the chain row the first of them lands in. =# +function _kept_columns(sched::ParallelMALASchedule, offset::Int, ncols::Int) + k0 = max(0, cld(offset - sched.discard_initial, sched.thinning)) + first_col = sched.discard_initial + 1 + k0 * sched.thinning - offset + return first_col:(sched.thinning):ncols, k0 + 1 +end + +#= Runs DEER blocks until `sched.Ntotal` steps are consumed, passing kept columns +to `sink!(S, logps, cols, first_row)`. Uses full-length blocks (`T`) to match the +fallback's RNG stream. Returns the last block; `S` aliases the workspace. =# +function _parallel_mala_run!( + sink!, + rng::Random.AbstractRNG, + model::DensityModel, + sampler::ParallelMALASampler, + sched::ParallelMALASchedule; + initial_params=nothing, + progress=AbstractMCMC.PROGRESS[], + progressname="Sampling", +) + progress = _parallel_mala_progress(progress, progressname) + x0 = _parallel_mala_initial_x(rng, model, sampler, initial_params) + prepped = _prepare_model(model, x0, sampler.T, sampler.backend) + ws = nothing + nsteps = 0 + Ntotal = sched.Ntotal + next_update = Ntotal / AbstractMCMC.get_n_updates(progress) + threshold = next_update + last = nothing + + AbstractMCMC.@maybewithricherlogger begin + AbstractMCMC.init_progress!(progress) + try + while nsteps < Ntotal + S, tape, ws = _deer_solve_new_tape(rng, prepped, sampler, x0; workspace=ws) + logps = _trajectory_logps(prepped, S) + nconsume = min(sampler.T, Ntotal - nsteps) + cols, first_row = _kept_columns(sched, nsteps, nconsume) + sink!(S, logps, cols, first_row) + + x0 = copy(view(S, :, nconsume)) + nsteps += nconsume + last = ( + S=S, logps=logps, workspace=ws, tape=tape, t=nconsume, model=prepped + ) + next_update = _parallel_mala_update_progress!( + progress, nsteps, Ntotal, next_update, threshold + ) + end + finally + AbstractMCMC.finish_progress!(progress) + end + end + return last +end + +function _copy_trajectory_rows!( + vals::AbstractMatrix{<:Real}, first_row::Int, S::AbstractMatrix, cols::AbstractRange +) + rows = first_row:(first_row + length(cols) - 1) + S_host = Array(S[:, cols]) + vals[rows, :] .= transpose(S_host) + return vals +end + +function _sample_parallel_mala_chain( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::ParallelMALASampler, + N::Int, + ::Type{FlexiChains.FlexiChain{TKey}}, + sched::ParallelMALASchedule; + initial_params=nothing, + param_names=nothing, + progress=AbstractMCMC.PROGRESS[], + progressname="Sampling", +) where {TKey} + D = model.dim + FP = typeof(sampler.epsilon) + + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) + + _parallel_mala_run!( + rng, model, sampler, sched; initial_params, progress, progressname + ) do S, logps, cols, first_row + _copy_trajectory_rows!(vals, first_row, S, cols) + logp[first_row:(first_row + length(cols) - 1)] .= view(logps, cols) + return nothing + end + + internals = (logp=logp,) + return _construct_flexichain(TKey, vals, internals, param_names, model) +end + +function _sample_parallel_mala_blocks( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::ParallelMALASampler, + N::Int, + sched::ParallelMALASchedule; + initial_params=nothing, + progress=AbstractMCMC.PROGRESS[], + progressname="Sampling", +) + nblocks = cld(sched.Ntotal, sampler.T) + blocks = Vector{AbstractMatrix}(undef, 0) + logp_blocks = Vector{AbstractVector}(undef, 0) + ends = Int[] + sizehint!(blocks, nblocks) + sizehint!(logp_blocks, nblocks) + sizehint!(ends, nblocks) + + last = _parallel_mala_run!( + rng, model, sampler, sched; initial_params, progress, progressname + ) do S, logps, cols, first_row + isempty(cols) && return nothing + push!(blocks, S[:, cols]) + push!(logp_blocks, logps[cols]) + push!(ends, first_row + length(cols) - 1) + return nothing + end + + S = copy(last.S) + final_state = ParallelMALAState( + S[:, last.t], + last.logps[last.t], + S, + last.logps, + last.workspace, + last.tape, + last.t, + last.model, + ) + return ParallelMALABlockSamples(blocks, logp_blocks, ends, N), final_state +end + +function _default_parallel_mala_mcmcsample( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::ParallelMALASampler, + N::Integer; + kwargs..., +) + return invoke( + AbstractMCMC.mcmcsample, + Tuple{ + Random.AbstractRNG, + AbstractMCMC.AbstractModel, + AbstractMCMC.AbstractSampler, + Integer, + }, + rng, + model, + sampler, + N; + kwargs..., + ) +end + +function AbstractMCMC.mcmcsample( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::ParallelMALASampler, + N::Integer; + progress=AbstractMCMC.PROGRESS[], + progressname="Sampling", + callback=nothing, + num_warmup::Int=0, + discard_initial::Int=num_warmup, + thinning=1, + chain_type::Type=Any, + initial_state=nothing, + initial_params=nothing, + param_names=nothing, + kwargs..., +) + #= `callback` wants a transition per step and `initial_state` a trajectory to + replay, so both go through the generic step loop. =# + if callback !== nothing || initial_state !== nothing + return _default_parallel_mala_mcmcsample( + rng, + model, + sampler, + N; + progress=progress, + progressname=progressname, + callback=callback, + num_warmup=num_warmup, + discard_initial=discard_initial, + thinning=thinning, + chain_type=chain_type, + initial_state=initial_state, + initial_params=initial_params, + param_names=param_names, + kwargs..., + ) + end + + N_int = Int(N) + sched = ParallelMALASchedule(N_int; discard_initial, thinning, num_warmup) + + if chain_type <: FlexiChains.FlexiChain + return _sample_parallel_mala_chain( + rng, + model, + sampler, + N_int, + chain_type, + sched; + initial_params=initial_params, + param_names=param_names, + progress=progress, + progressname=progressname, + ) + end + + samples, state = _sample_parallel_mala_blocks( + rng, + model, + sampler, + N_int, + sched; + initial_params=initial_params, + progress=progress, + progressname=progressname, + ) + chain_type === Any && return samples + + sample_vec = [samples[i] for i in 1:length(samples)] + return AbstractMCMC.bundle_samples( + sample_vec, model, sampler, state, chain_type; param_names=param_names, kwargs... + ) +end + +function AbstractMCMC.step( + rng::Random.AbstractRNG, + model::DensityModel, + sampler::ParallelMALASampler{FP}; + initial_params=nothing, + kwargs..., +) where {FP} + x0 = if initial_params !== nothing + copy(initial_params) + else + randn(rng, FP, model.dim) + end + model = _prepare_model(model, x0, sampler.T, sampler.backend) + + S, tape, ws = _deer_solve_new_tape(rng, model, sampler, x0) + logps = _trajectory_logps(model, S) + x1 = S[:, 1] + logp1 = logps[1] + trans = ParallelMALATransition(x1, logp1) + state = ParallelMALAState(x1, logp1, S, logps, ws, tape, 1, model) + return trans, state +end + +function AbstractMCMC.step( + rng::Random.AbstractRNG, + user_model::DensityModel, + sampler::ParallelMALASampler, + state::ParallelMALAState; + kwargs..., +) + #= A state from a different model (via `initial_state`) brings a trajectory + solved under that model, so the rest of it can't be replayed either: pick up + from the current position with a fresh tape under the model we were handed. =# + reuse = _prepped_for(state.model, user_model) + model = if reuse + state.model + else + _prepare_model(user_model, state.x, sampler.T, sampler.backend) + end + T = sampler.T + t_next = state.t + 1 + + if reuse && t_next <= T + x_new = state.trajectory[:, t_next] + logp_new = state.logps[t_next] + trans = ParallelMALATransition(x_new, logp_new) + new_state = ParallelMALAState( + x_new, + logp_new, + state.trajectory, + state.logps, + state.workspace, + state.tape, + t_next, + model, + ) + return trans, new_state + else + x0 = reuse ? state.trajectory[:, T] : copy(state.x) + S_new, tape, ws = _deer_solve_new_tape( + rng, model, sampler, x0; workspace=state.workspace + ) + logps = _trajectory_logps(model, S_new) + x_new = S_new[:, 1] + logp_new = logps[1] + trans = ParallelMALATransition(x_new, logp_new) + new_state = ParallelMALAState(x_new, logp_new, S_new, logps, ws, tape, 1, model) + return trans, new_state + end +end + +for TKey in (Symbol, VarName) + @eval function AbstractMCMC.bundle_samples( + samples::Vector{<:ParallelMALATransition}, + model::DensityModel, + sampler::ParallelMALASampler, + state::ParallelMALAState, + ::Type{FlexiChains.FlexiChain{$TKey}}; + param_names=nothing, + kwargs..., + ) + N = length(samples) + D = model.dim + FP = typeof(sampler.epsilon) + + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) + + for i in 1:N + vals[i, :] .= samples[i].x + logp[i] = samples[i].logp + end + + internals = (logp=logp,) + return _construct_flexichain($TKey, vals, internals, param_names, model) + end +end diff --git a/test/test-Reactant-Validation.jl b/test/test-Reactant-Validation.jl index f0af59d..41839c3 100644 --- a/test/test-Reactant-Validation.jl +++ b/test/test-Reactant-Validation.jl @@ -4,7 +4,7 @@ using ADTypes using ForwardDiff: ForwardDiff using LogDensityProblems: LogDensityProblems -#= The AutoReactant validation in src/interface.jl needs no Reactant, so it runs +#= The AutoReactant validation in src/model.jl needs no Reactant, so it runs in every test invocation. The Reactant-backed tests are opt-in in test-Reactant-HVP.jl. =# From f18e14fafe34ed3b86c997c17e2cdafe486fa2ad Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Fri, 30 Oct 2026 09:02:55 -0400 Subject: [PATCH 2/2] Move test-only MALA references into test/utils.jl run_mala_sequential_taped, mala_accept_indicator and mala_step_full had no callers in src or ext. They now live in a TestUtils module that the tests include, and the docs GIF script builds its sequential path from mala_step_taped. --- CHANGELOG.md | 4 ++ docs/src/assets/make_julia_deer_gif.jl | 8 +-- src/MALA/MALA.jl | 56 --------------------- test/test-Adaptive-MALA.jl | 4 +- test/test-DEER-Logistic.jl | 6 +-- test/test-Deer-vs-MALA.jl | 16 +++--- test/test-GPU-DEER.jl | 4 +- test/test-GPU-MALA.jl | 4 +- test/test-MALA-Kernel.jl | 14 +++--- test/utils.jl | 67 ++++++++++++++++++++++++++ 10 files changed, 101 insertions(+), 82 deletions(-) create mode 100644 test/utils.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index cf4aa93..ec492ab 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -70,6 +70,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 `DensityModel(::DynamicPPL.Model)` constructor are gone, since DynamicPPL models cannot run on the GPU. Use ParallelMCMC 0.2 for Turing models, or wrap the log-density in a `LogDensityProblems` object or a plain `DensityModel`. +- `MALA.run_mala_sequential_taped`, `MALA.mala_accept_indicator` and + `MALA.mala_step_full`. Nothing in the package called them; they now live in + the test suite as references. Use `MALA.mala_step_taped` or + `MALA.mala_step_with_logα` instead. ## [0.2.0] - 2026-06-29 diff --git a/docs/src/assets/make_julia_deer_gif.jl b/docs/src/assets/make_julia_deer_gif.jl index c96d504..1e9d6c6 100644 --- a/docs/src/assets/make_julia_deer_gif.jl +++ b/docs/src/assets/make_julia_deer_gif.jl @@ -354,10 +354,10 @@ function main() noise = [step.xi for step in tape] uniforms = [step.u for step in tape] - sequential = MALA.run_mala_sequential_taped( - logposterior, gradposterior, x0, epsilon, noise, uniforms - ) - final_trajectory = reduce(hcat, sequential[2:end]) + sequential = accumulate(zip(noise, uniforms); init=x0) do x, (xi, u) + return MALA.mala_step_taped(logposterior, gradposterior, x, epsilon, xi, u) + end + final_trajectory = reduce(hcat, sequential) output = isempty(ARGS) ? joinpath(@__DIR__, "julia_deer_posterior.gif") : ARGS[1] mkpath(dirname(output)) diff --git a/src/MALA/MALA.jl b/src/MALA/MALA.jl index 465434f..65805ce 100644 --- a/src/MALA/MALA.jl +++ b/src/MALA/MALA.jl @@ -8,13 +8,10 @@ export MALAWorkspace, logq_mala!, mala_step_taped, mala_step_taped!, - run_mala_sequential_taped, mala_proposal, mala_proposal!, mala_logα, mala_logα!, - mala_accept_indicator, - mala_step_full, mala_step_with_logα, mala_step_with_logα!, mala_step_surrogate_sigmoid, @@ -313,59 +310,6 @@ function mala_step_taped( return x_next end -""" -Run sequential taped MALA for T steps. - -Returns: -- xs: Vector of states length T+1, xs[1]=x0, xs[t+1]=x_t -""" -function run_mala_sequential_taped( - logp, - gradlogp, - x0::AbstractVector, - ϵ::Real, - ξs::Vector{<:AbstractVector}, - us::AbstractVector; - cholM=nothing, -) - T = length(us) - length(ξs) == T || throw(DimensionMismatch("ξs and us must have the same length")) - xs = Vector{typeof(x0)}(undef, T + 1) - xs[1] = copy(x0) - - ws = MALAWorkspace(x0) - x = copy(x0) - x_next = similar(x0) - - for t in 1:T - mala_step_taped!(x_next, ws, logp, gradlogp, x, ϵ, ξs[t], us[t]; cholM=cholM) - xs[t + 1] = copy(x_next) - x, x_next = x_next, x - end - return xs -end - -""" -Primal accept indicator for a taped MALA step. -Returns a float in {0, 1} matching the precision of `u`. -""" -function mala_accept_indicator( - logp, gradlogp, x::AbstractVector, ϵ::Real, ξ::AbstractVector, u::Real; cholM=nothing -) - ws = MALAWorkspace(x) - mala_proposal!(ws.y, ws, logp, gradlogp, x, ϵ, ξ; cholM=cholM) - logα = mala_logα!(ws, logp, gradlogp, x, ws.y, ϵ; cholM=cholM) - FP = typeof(float(u)) - return (log(u) < logα) ? one(FP) : zero(FP) -end - -function mala_step_full( - logp, gradlogp, x::AbstractVector, ϵ::Real, ξ::AbstractVector, u::Real; cholM=nothing -) - x_next, accepted, _ = mala_step_with_logα(logp, gradlogp, x, ϵ, ξ, u; cholM=cholM) - return x_next, accepted -end - function mala_step_with_logα( logp, gradlogp, x::AbstractVector, ϵ::Real, ξ::AbstractVector, u::Real; cholM=nothing ) diff --git a/test/test-Adaptive-MALA.jl b/test/test-Adaptive-MALA.jl index 85adccf..588ab5e 100644 --- a/test/test-Adaptive-MALA.jl +++ b/test/test-Adaptive-MALA.jl @@ -5,6 +5,8 @@ using Statistics using FlexiChains using ParallelMCMC +isdefined(@__MODULE__, :TestUtils) || include("utils.jl") +using .TestUtils const MALA = ParallelMCMC.MALA logp_adapt(x) = -0.5 * dot(x, x) @@ -18,7 +20,7 @@ gradlogp_adapt(x) = -x ξ = randn(rng, D) u = rand(rng) * 0.999 + 1e-15 - x1, a1 = MALA.mala_step_full(logp_adapt, gradlogp_adapt, x, 0.1, ξ, u) + x1, a1 = mala_step_full(logp_adapt, gradlogp_adapt, x, 0.1, ξ, u) x2, a2, logα = MALA.mala_step_with_logα(logp_adapt, gradlogp_adapt, x, 0.1, ξ, u) @test x1 == x2 diff --git a/test/test-DEER-Logistic.jl b/test/test-DEER-Logistic.jl index 730ab5e..95e5562 100644 --- a/test/test-DEER-Logistic.jl +++ b/test/test-DEER-Logistic.jl @@ -5,6 +5,8 @@ using Statistics using FlexiChains using ParallelMCMC +isdefined(@__MODULE__, :TestUtils) || include("utils.jl") +using .TestUtils using ADTypes using Enzyme @@ -124,9 +126,7 @@ else logp_cpu = β -> _logp_lr(β, _X_f32, _y_f32) gradlogp_cpu = β -> _gradlogp_lr(β, _X_f32, _y_f32) - xs_seq = ParallelMCMC.MALA.run_mala_sequential_taped( - logp_cpu, gradlogp_cpu, x0_cpu, ε, ξs_cpu, us - ) + xs_seq = run_mala_sequential_taped(logp_cpu, gradlogp_cpu, x0_cpu, ε, ξs_cpu, us) X_gpu = CUDA.CuMatrix(_X_f32) y_gpu = CUDA.CuVector(_y_f32) diff --git a/test/test-Deer-vs-MALA.jl b/test/test-Deer-vs-MALA.jl index 09281cb..b0602aa 100644 --- a/test/test-Deer-vs-MALA.jl +++ b/test/test-Deer-vs-MALA.jl @@ -5,6 +5,8 @@ using ADTypes using Enzyme using ParallelMCMC +isdefined(@__MODULE__, :TestUtils) || include("utils.jl") +using .TestUtils const MALA = ParallelMCMC.MALA const DEER = ParallelMCMC.DEER const _AD = ADTypes.AutoEnzyme() @@ -20,7 +22,7 @@ function count_mala_rejections(logp, gradlogp, x0, ϵ, ξs, us) T = length(ξs) @assert length(us) == T for t in 1:T - a = MALA.mala_accept_indicator(logp, gradlogp, x, ϵ, ξs[t], us[t]) # 1.0 accept, 0.0 reject + a = mala_accept_indicator(logp, gradlogp, x, ϵ, ξs[t], us[t]) # 1.0 accept, 0.0 reject if a == 0.0 rejects += 1 # x unchanged on rejection @@ -59,9 +61,7 @@ end us = rand(rng, T) # Ground truth: sequential taped MALA - xs_seq = MALA.run_mala_sequential_taped( - logp_stdnormal, gradlogp_stdnormal, x0, ϵ, ξs, us - ) + xs_seq = run_mala_sequential_taped(logp_stdnormal, gradlogp_stdnormal, x0, ϵ, ξs, us) @test length(xs_seq) == T + 1 # Ensure we actually exercised the rejection (gate a=0) case at least once. rejects = count_mala_rejections(logp_stdnormal, gradlogp_stdnormal, x0, ϵ, ξs, us) @@ -80,9 +80,7 @@ end us = rand(rng, T) # Ground truth: sequential taped MALA - xs_seq = MALA.run_mala_sequential_taped( - logp_stdnormal, gradlogp_stdnormal, x0, ϵ, ξs, us - ) + xs_seq = run_mala_sequential_taped(logp_stdnormal, gradlogp_stdnormal, x0, ϵ, ξs, us) @test length(xs_seq) == T + 1 tape = [(ξ=ξs[t], u=us[t]) for t in 1:T] @@ -170,9 +168,7 @@ end us = rand(rng, T) # Ground truth: sequential taped MALA - xs_seq = MALA.run_mala_sequential_taped( - logp_stdnormal, gradlogp_stdnormal, x0, ϵ, ξs, us - ) + xs_seq = run_mala_sequential_taped(logp_stdnormal, gradlogp_stdnormal, x0, ϵ, ξs, us) @test length(xs_seq) == T + 1 # Ensure at least one rejection occurred (so gating is exercised) diff --git a/test/test-GPU-DEER.jl b/test/test-GPU-DEER.jl index eb8e2fd..66f302c 100644 --- a/test/test-GPU-DEER.jl +++ b/test/test-GPU-DEER.jl @@ -4,6 +4,8 @@ using LinearAlgebra using Statistics using ParallelMCMC +isdefined(@__MODULE__, :TestUtils) || include("utils.jl") +using .TestUtils import ParallelMCMC: DEER, MALA using ADTypes: ADTypes using ForwardDiff: ForwardDiff @@ -91,7 +93,7 @@ else us = Float32.(rand(rng, T)) x0_cpu = randn(rng, Float32, D) - xs_seq = MALA.run_mala_sequential_taped(_logp, _gradlogp, x0_cpu, ε, ξs_cpu, us) + xs_seq = run_mala_sequential_taped(_logp, _gradlogp, x0_cpu, ε, ξs_cpu, us) tape = [(ξ=CUDA.CuArray(ξs_cpu[t]), u=us[t]) for t in 1:T] rec = _make_gpu_rec(tape, ε, backend) diff --git a/test/test-GPU-MALA.jl b/test/test-GPU-MALA.jl index 8aee951..7d523e2 100644 --- a/test/test-GPU-MALA.jl +++ b/test/test-GPU-MALA.jl @@ -4,6 +4,8 @@ using LinearAlgebra using Statistics using ParallelMCMC +isdefined(@__MODULE__, :TestUtils) || include("utils.jl") +using .TestUtils const MALA = ParallelMCMC.MALA using CUDA: CUDA @@ -104,7 +106,7 @@ else # Scalar step on CPU logp_scalar = x -> -0.5f0 * dot(x, x) grad_scalar = x -> -x - x_scalar, acc_scalar = MALA.mala_step_full(logp_scalar, grad_scalar, x, 0.1f0, ξ, u) + x_scalar, acc_scalar = mala_step_full(logp_scalar, grad_scalar, x, 0.1f0, ξ, u) # Batched N=1 on GPU X_gpu = CUDA.CuArray(reshape(copy(x), D, 1)) diff --git a/test/test-MALA-Kernel.jl b/test/test-MALA-Kernel.jl index 3da55ed..990d1df 100644 --- a/test/test-MALA-Kernel.jl +++ b/test/test-MALA-Kernel.jl @@ -3,6 +3,8 @@ using Random using LinearAlgebra using ParallelMCMC +isdefined(@__MODULE__, :TestUtils) || include("utils.jl") +using .TestUtils const MALA = ParallelMCMC.MALA # Standard normal target in R^d @@ -52,10 +54,10 @@ end # Fixed tape ξs, us = make_tape(rng, D, T) - xs1 = MALA.run_mala_sequential_taped( + xs1 = run_mala_sequential_taped( logp_stdnormal_kernel, gradlogp_stdnormal_kernel, x0, ϵ, ξs, us ) - xs2 = MALA.run_mala_sequential_taped( + xs2 = run_mala_sequential_taped( logp_stdnormal_kernel, gradlogp_stdnormal_kernel, x0, ϵ, ξs, us ) @@ -91,7 +93,7 @@ end x0 = randn(rng, D) ξs, us = make_tape(rng, D, T) - xs = MALA.run_mala_sequential_taped( + xs = run_mala_sequential_taped( logp_stdnormal_kernel, gradlogp_stdnormal_kernel, x0, ϵ, ξs, us ) @@ -156,10 +158,10 @@ end @test x_rej == x # Also sanity-check the accept indicator - a_acc = MALA.mala_accept_indicator( + a_acc = mala_accept_indicator( logp_stdnormal_kernel, gradlogp_stdnormal_kernel, x, ϵ, ξ, u_accept ) - a_rej = MALA.mala_accept_indicator( + a_rej = mala_accept_indicator( logp_stdnormal_kernel, gradlogp_stdnormal_kernel, x, ϵ, ξ, u_reject ) @test a_acc == 1.0 @@ -228,7 +230,7 @@ end for j in 1:N xj = copy(view(X, :, j)) ξj = copy(view(Ξ, :, j)) - x_ref, a_ref = MALA.mala_step_full( + x_ref, a_ref = mala_step_full( logp_stdnormal_kernel, gradlogp_stdnormal_kernel, xj, diff --git a/test/utils.jl b/test/utils.jl new file mode 100644 index 0000000..dc6a9e4 --- /dev/null +++ b/test/utils.jl @@ -0,0 +1,67 @@ +#= Sequential MALA references the tests compare DEER and the batched kernels +against. They are built from the kernel functions in `ParallelMCMC.MALA`; the +package itself never calls them. Include with + + isdefined(@__MODULE__, :TestUtils) || include("utils.jl") + using .TestUtils +=# +module TestUtils + +using ParallelMCMC: MALA + +export run_mala_sequential_taped, mala_accept_indicator, mala_step_full + +""" +Run sequential taped MALA for T steps. + +Returns: +- xs: Vector of states length T+1, xs[1]=x0, xs[t+1]=x_t +""" +function run_mala_sequential_taped( + logp, + gradlogp, + x0::AbstractVector, + ϵ::Real, + ξs::Vector{<:AbstractVector}, + us::AbstractVector; + cholM=nothing, +) + T = length(us) + length(ξs) == T || throw(DimensionMismatch("ξs and us must have the same length")) + xs = Vector{typeof(x0)}(undef, T + 1) + xs[1] = copy(x0) + + ws = MALA.MALAWorkspace(x0) + x = copy(x0) + x_next = similar(x0) + + for t in 1:T + MALA.mala_step_taped!(x_next, ws, logp, gradlogp, x, ϵ, ξs[t], us[t]; cholM=cholM) + xs[t + 1] = copy(x_next) + x, x_next = x_next, x + end + return xs +end + +""" +Primal accept indicator for a taped MALA step. +Returns a float in {0, 1} matching the precision of `u`. +""" +function mala_accept_indicator( + logp, gradlogp, x::AbstractVector, ϵ::Real, ξ::AbstractVector, u::Real; cholM=nothing +) + ws = MALA.MALAWorkspace(x) + MALA.mala_proposal!(ws.y, ws, logp, gradlogp, x, ϵ, ξ; cholM=cholM) + logα = MALA.mala_logα!(ws, logp, gradlogp, x, ws.y, ϵ; cholM=cholM) + FP = typeof(float(u)) + return (log(u) < logα) ? one(FP) : zero(FP) +end + +function mala_step_full( + logp, gradlogp, x::AbstractVector, ϵ::Real, ξ::AbstractVector, u::Real; cholM=nothing +) + x_next, accepted, _ = MALA.mala_step_with_logα(logp, gradlogp, x, ϵ, ξ, u; cholM=cholM) + return x_next, accepted +end + +end # module