Allow structured parameter containers in ParamsWithStats

Project.toml

@@ -3,7 +3,7 @@ uuid = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 keywords = ["markov chain monte carlo", "probabilistic programming"]
 license = "MIT"
 desc = "A lightweight interface for common MCMC methods."
-version = "5.15.1"
+version = "5.16.0"
 
 [deps]
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"

docs/src/callbacks.md

@@ -250,11 +250,28 @@ function AbstractMCMC.getstats(state::MyState)
 end
 ```
 
-The `ParamsWithStats` constructors normalize all inputs to `NamedTuple`:
+The default `ParamsWithStats` extraction constructors normalize common inputs to
+`NamedTuple`:
 - `Vector{<:Real}` gets default `θ[i]` names
 - `Vector{Pair}` is converted to `NamedTuple` with the provided names
 - `NamedTuple` is used directly
 
+Packages with a structured parameter representation can also construct
+`ParamsWithStats(params, stats)` directly. The parameter container must implement `pairs`
+and `isempty`; this is not validated at construction (some Base types already define those
+methods and will not fail later). Keys should have a meaningful `string` form so
+name-based filtering and logging work. The two-argument constructor uses an empty
+`NamedTuple` for `extras`. Note the normalization above applies to `AbstractVector`
+subtypes: a vector-like container is converted to a `Symbol`-keyed `NamedTuple` rather
+than stored as given.
+
+!!! warning "Mixed key types"
+    Because the parameter container's keys need not be `Symbol`s, `Base.pairs(pws)` may
+    yield pairs with mixed key types (e.g., `VarName` keys from the parameters and
+    `Symbol` keys from the statistics). Consumers should iterate the pairs generically
+    (e.g., stringify keys) and must not assume `Symbol` keys, a concrete element type, or
+    that the pairs can be collected into a `NamedTuple`.
+
 !!! note "stats vs extras"
     Use `stats` for values that change once per MCMC iteration (e.g., log probability, acceptance rate).
     Use `extras` for values that are constant across iterations (e.g., preconditioning matrix, number of particles)

src/callbacks.jl

@@ -118,13 +118,22 @@ end
 """
     ParamsWithStats{P,S,E}
 
-A container for MCMC parameters, statistics, and extras.
+A container for MCMC parameters, statistics, and extras. The parameter container can be a
+structured type; statistics and extras are stored as `NamedTuple`s. Use `Base.pairs(pws)`
+to iterate over all `(name, value)` pairs.
 
-All fields are stored as `NamedTuple`s to ensure a tight, well-defined interface.
-Use `Base.pairs(pws)` to iterate over `(name, value)` pairs.
+The parameter container must implement `pairs` and `isempty`. This is not validated at
+construction; some Base types already define those methods and will not fail later. Keys
+should have a meaningful `string` form so that name-based filtering and logging callbacks
+work. Keys are not required to be `Symbol`s, so `pairs(pws)` may yield pairs with mixed
+key types; consumers should not assume `Symbol` keys or a concrete element type.
+
+Note that `AbstractVector{<:Real}` and `AbstractVector{<:Pair}` parameter inputs are not
+stored as given: the extraction constructors normalize them to `Symbol`-keyed
+`NamedTuple`s (see the constructor docs below).
 
 # Fields
-- `params::P`: Parameter values as a NamedTuple
+- `params::P`: Parameter values in a container implementing `pairs` and `isempty`
 - `stats::S`: Statistics as a NamedTuple (e.g., `(lp=...,)`)
 - `extras::E`: Extra diagnostics as a NamedTuple
 
@@ -139,12 +148,21 @@ end
 pws2 = ParamsWithStats(pws; params=true, stats=false)
 ```
 """
-struct ParamsWithStats{P<:NamedTuple,S<:NamedTuple,E<:NamedTuple}
+struct ParamsWithStats{P,S<:NamedTuple,E<:NamedTuple}
     params::P
     stats::S
     extras::E
 end
 
+"""
+    ParamsWithStats(params, stats::NamedTuple)
+
+Construct a `ParamsWithStats` with no extra diagnostics.
+"""
+function ParamsWithStats(params, stats::NamedTuple)
+    return ParamsWithStats(params, stats, NamedTuple())
+end
+
 # Constructor from Vector{<:Real} - adds default θ[i] names
 function ParamsWithStats(
     v::AbstractVector{<:Real}, stats::S, extras::E
@@ -222,6 +240,10 @@ end
 
 Return an iterator of `(name, value)` pairs for all selected data in `pws`.
 
+Parameter keys need not be `Symbol`s, so the iterator may yield mixed key types (e.g.
+structured parameter keys together with `Symbol` keys from `stats`/`extras`). Consumers
+should iterate generically and must not assume `Symbol` keys or a concrete element type.
+
 This is the canonical way to iterate over a `ParamsWithStats`:
 ```julia
 for (name, value) in Base.pairs(pws)

test/callbacks.jl

@@ -166,6 +166,20 @@ end
 ### ParamsWithStats   ###
 #########################
 
+# StructuredKey: non-Symbol key with string form (VarName stand-in). StructuredParams
+# wraps a Vector of Pairs so AbstractVector{<:Pair} normalization does not kick in.
+struct StructuredKey
+    name::String
+end
+Base.string(k::StructuredKey) = k.name
+Base.:(==)(a::StructuredKey, b::StructuredKey) = a.name == b.name
+
+struct StructuredParams
+    data::Vector{Pair{StructuredKey,Float64}}
+end
+Base.pairs(params::StructuredParams) = params.data
+Base.isempty(params::StructuredParams) = isempty(params.data)
+
 @testset "ParamsWithStats" begin
     @testset "Constructor from NamedTuple" begin
         pws = AbstractMCMC.ParamsWithStats((a=1.0, b=2.0), (lp=-10.0,), NamedTuple())
@@ -178,6 +192,11 @@ end
     @testset "Constructor from Vector{Real} - default names" begin
         pws = AbstractMCMC.ParamsWithStats([1.0, 2.0, 3.0], NamedTuple(), NamedTuple())
         @test pws.params == (var"θ[1]"=1.0, var"θ[2]"=2.0, var"θ[3]"=3.0)
+
+        pws_without_extras = AbstractMCMC.ParamsWithStats([1.0, 2.0], (lp=-1.0,))
+        @test pws_without_extras.params == (var"θ[1]"=1.0, var"θ[2]"=2.0)
+        @test pws_without_extras.stats == (lp=-1.0,)
+        @test isempty(pws_without_extras.extras)
     end
 
     @testset "Constructor from Vector{Pair} - named" begin
@@ -187,6 +206,35 @@ end
         @test pws.params == (μ=1.0, σ=2.0)
     end
 
+    @testset "Structured parameter container" begin
+        params = StructuredParams([StructuredKey("μ") => 1.0, StructuredKey("y") => 3.0])
+
+        # The two-argument constructor stores the container as-is and defaults extras.
+        pws = AbstractMCMC.ParamsWithStats(params, (lp=-10.0,))
+        @test pws.params === params
+        @test pws.stats == (lp=-10.0,)
+        @test pws.extras == NamedTuple()
+
+        # pairs flattens params/stats/extras and yields mixed key types.
+        pws = AbstractMCMC.ParamsWithStats(params, (lp=-10.0,), (step_size=0.1,))
+        pairs_list = collect(Base.pairs(pws))
+        @test pairs_list == [
+            StructuredKey("μ") => 1.0,
+            StructuredKey("y") => 3.0,
+            :lp => -10.0,
+            :step_size => 0.1,
+        ]
+        @test pairs_list[1][1] isa StructuredKey
+        @test pairs_list[3][1] isa Symbol
+        @test !isempty(pws)
+        @test isempty(AbstractMCMC.ParamsWithStats(StructuredParams([]), NamedTuple()))
+
+        # NameFilter stringifies keys, so non-Symbol parameter keys still filter.
+        f = AbstractMCMC.NameFilter(; include=["μ", "lp"])
+        filtered = filter(p -> f(first(p)), pairs_list)
+        @test filtered == [StructuredKey("μ") => 1.0, :lp => -10.0]
+    end
+
     @testset "Constructor from state" begin
         state = 5
         pws = AbstractMCMC.ParamsWithStats(
@@ -198,6 +246,23 @@ end
         @test pws.extras == NamedTuple()
     end
 
+    @testset "Structured container through extraction and re-selection" begin
+        # Let the container act as its own state so getparams/getstats drive extraction.
+        AbstractMCMC.getparams(params::StructuredParams) = params
+        AbstractMCMC.getstats(::StructuredParams) = (lp=-3.0,)
+
+        params = StructuredParams([StructuredKey("α") => 0.5])
+        pws = AbstractMCMC.ParamsWithStats(
+            MyModel(), MySampler(), nothing, params; params=true, stats=true
+        )
+        @test pws.params === params
+        @test collect(Base.pairs(pws)) == [StructuredKey("α") => 0.5, :lp => -3.0]
+
+        # Re-selection preserves the container; deselecting substitutes an empty NamedTuple.
+        @test AbstractMCMC.ParamsWithStats(pws; stats=false).params === params
+        @test AbstractMCMC.ParamsWithStats(pws; params=false).params == NamedTuple()
+    end
+
     @testset "Copy constructor with selection" begin
         pws = AbstractMCMC.ParamsWithStats((a=1.0,), (lp=-10.0,), NamedTuple())
 
@@ -229,12 +294,14 @@ end
         @test isempty(pws_empty)
     end
 
-    @testset "Illegal states are unrepresentable" begin
-        # Should not be able to construct with arbitrary types
+    @testset "Stats and extras are constrained" begin
+        # Statistics and extras must be NamedTuples. Params are not validated at
+        # construction (e.g. `1` still constructs; Base defines pairs/isempty for it).
         @test_throws MethodError AbstractMCMC.ParamsWithStats(1, 2, 3)
-        @test_throws MethodError AbstractMCMC.ParamsWithStats(
-            "bad", NamedTuple(), NamedTuple()
-        )
+        @test_throws MethodError AbstractMCMC.ParamsWithStats("params", 2, NamedTuple())
+        @test_throws MethodError AbstractMCMC.ParamsWithStats("params", NamedTuple(), 3)
+        @test AbstractMCMC.ParamsWithStats(1, NamedTuple(), NamedTuple()) isa
+            AbstractMCMC.ParamsWithStats
     end
 end