diff --git a/Project.toml b/Project.toml index 7a667fa3..37d9aa42 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md index 0156f238..3ca93714 100644 --- a/docs/src/callbacks.md +++ b/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) diff --git a/src/callbacks.jl b/src/callbacks.jl index 9c43b17a..836c3c5e 100644 --- a/src/callbacks.jl +++ b/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) diff --git a/test/callbacks.jl b/test/callbacks.jl index 24903b55..8f1f2821 100644 --- a/test/callbacks.jl +++ b/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