reorganize interface

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
 

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.
 

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))

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
 )

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

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