diff --git a/Project.toml b/Project.toml index 13a1bc47..007b104b 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.11.0" +version = "5.12.0" [deps] BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md index 7745e480..bbf7783e 100644 --- a/docs/src/callbacks.md +++ b/docs/src/callbacks.md @@ -122,8 +122,8 @@ cb = mcmc_callback( name_filter=( include=["mu", "sigma"], # Only log these parameters exclude=["_internal"], # Exclude matching names - extras=true, # Include extra stats (log density, etc.) - hyperparams=true, # Include hyperparameters (logged once) + stats=true, # Include step-level statistics + extras=true, # Include extra diagnostics ), ) ``` @@ -176,13 +176,12 @@ Navigate to `localhost:6006` in your browser to see the dashboard. You'll see re ## API Reference -### Main Functions - ```@docs -mcmc_callback +AbstractMCMC.mcmc_callback +AbstractMCMC.ParamsWithStats ``` -## Default Values +### Default Values ### stats_options defaults @@ -198,8 +197,8 @@ mcmc_callback |--------------|------------|----------------------------------| | `include` | `String[]` | Only log these (empty=all) | | `exclude` | `String[]` | Don't log these | -| `extras` | `false` | Include extra stats | -| `hyperparams`| `false` | Include hyperparameters | +| `stats` | `false` | Include step-level statistics | +| `extras` | `false` | Include extra diagnostics | ## Implementing Custom Callbacks @@ -211,6 +210,67 @@ function my_callback(rng, model, sampler, transition, state, iteration; kwargs.. end ``` +## ParamsWithStats + +`ParamsWithStats` is a container for extracting and iterating over MCMC parameters, statistics, and extras. + +### Basic Usage + +```julia +# Extract params and stats from state +pws = ParamsWithStats(model, sampler, transition, state; params=true, stats=true) + +# Iterate using Base.pairs +for (name, value) in Base.pairs(pws) + @info name value +end + +# Re-select to get only params +pws_params = ParamsWithStats(pws; params=true, stats=false, extras=false) +``` + +### Overriding for Your Package + +To provide meaningful variable names, override the extraction hooks: + +```julia +# Option 1: Return Vector{<:Real} - default names (θ[1], θ[2], ...) will be used +function AbstractMCMC.getparams(state::MyState) + return [state.mu, state.sigma] +end + +# Option 2: Return named pairs - will be converted to NamedTuple +function AbstractMCMC.getparams(state::MyState) + return ["μ" => state.mu, "σ" => state.sigma] +end + +# Override getstats to return step-level statistics as NamedTuple +function AbstractMCMC.getstats(state::MyState) + return (lp=state.logp, acceptance_rate=state.accept_rate) +end +``` + +The `ParamsWithStats` constructors normalize all inputs to `NamedTuple`: +- `Vector{<:Real}` gets default `θ[i]` names +- `Vector{Pair}` is converted to `NamedTuple` with the provided names +- `NamedTuple` is used directly + +!!! 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) + or that change multiple times within a single iteration (e.g., leapfrog phase points). + +### Usage in TensorBoard Callback + +The TensorBoard callback uses `ParamsWithStats` with `Base.pairs`: + +```julia +pws = ParamsWithStats(model, sampler, t, state; params=true, stats=true) +for (k, val) in Base.pairs(pws) + @info "$k" val +end +``` + ## Internals !!! note @@ -235,20 +295,3 @@ When using statistics, AbstractMCMC provides wrappers that modify how samples ar | `WindowStat(n, stat)` | Use a rolling window of `n` observations | These are applied automatically via `stats_options`, but can also be used directly if needed. - -### Internal Functions - -The unified `_names_and_values` function extracts all relevant data from a sampler state: - -```julia -for (name, value) in AbstractMCMC._names_and_values( - model, sampler, transition, state; - params=true, - hyperparams=false, - extra=false, -) - println("$name = $value") -end -``` - -Samplers can override `AbstractMCMC.getparams(state)` and `AbstractMCMC.getstats(state)` to provide custom information extraction. diff --git a/ext/AbstractMCMCOnlineStatsExt.jl b/ext/AbstractMCMCOnlineStatsExt.jl index c9020df1..ab7e7b5e 100644 --- a/ext/AbstractMCMCOnlineStatsExt.jl +++ b/ext/AbstractMCMCOnlineStatsExt.jl @@ -144,21 +144,39 @@ end Update and log statistics. Called from TensorBoard callback. """ -function log_stat_impl!(stats::AbstractDict, prototype, key, val, prefix) +function log_stat_impl!(stats::AbstractDict, prototype, key, val::Real, prefix) + float_val = try + Float64(val) + catch + return nothing + end + + str_key = string(key) + stat = if prototype !== nothing - get!(stats, key) do + get!(stats, str_key) do deepcopy(prototype) end else - get(stats, key, nothing) + get(stats, str_key, nothing) end if stat !== nothing - fit!(stat, val) - @info "$(prefix)$key" stat + fit!(stat, float_val) + @info "$(prefix)$str_key" stat end end +function log_stat_impl!( + stats::AbstractDict, prototype, key, val::Pair{<:Any,T}, prefix +) where {T<:Real} + return log_stat_impl!(stats, prototype, key, last(val), prefix) +end + +function log_stat_impl!(stats::AbstractDict, prototype, key, val, prefix) + return nothing +end + log_stat_impl!(::Nothing, prototype, key, val, prefix) = nothing # tb_name helpers for formatting stat names in TensorBoard diff --git a/ext/AbstractMCMCTensorBoardLoggerExt.jl b/ext/AbstractMCMCTensorBoardLoggerExt.jl index ab60e6d2..0839bf72 100644 --- a/ext/AbstractMCMCTensorBoardLoggerExt.jl +++ b/ext/AbstractMCMCTensorBoardLoggerExt.jl @@ -4,8 +4,7 @@ using AbstractMCMC using AbstractMCMC: MultiCallback, NameFilter, - _names_and_values, - hyperparam_metrics, + ParamsWithStats, merge_with_defaults, create_stats_with_options, DEFAULT_STATS_OPTIONS, @@ -14,6 +13,10 @@ using TensorBoardLogger using TensorBoardLogger: TBLogger using Logging: AbstractLogger, with_logger, @info +########################### +### TensorBoardCallback ### +########################### + """ TensorBoardCallback @@ -25,8 +28,8 @@ struct TensorBoardCallback{L,S,P,F} stats::S stat_prototype::P variable_filter::F + include_stats::Bool include_extras::Bool - include_hyperparams::Bool param_prefix::String extras_prefix::String end @@ -43,7 +46,7 @@ Create a TensorBoard logging callback. - `true` or `:default`: Use default statistics (Mean, Variance, KHist) - requires OnlineStats - An OnlineStat or tuple of OnlineStats - requires OnlineStats - `stats_options`: NamedTuple with `thin`, `skip`, `window` -- `name_filter`: NamedTuple with `include`, `exclude`, `extras`, `hyperparams` +- `name_filter`: NamedTuple with `include`, `exclude`, `stats`, `extras` - `num_bins`: Number of histogram bins (default: 100) # Examples @@ -85,8 +88,8 @@ function AbstractMCMC.mcmc_callback(; stats_dict, prototype, variable_filter, + merged_name_filter.stats, merged_name_filter.extras, - merged_name_filter.hyperparams, "", "extras/", ) @@ -105,38 +108,19 @@ function (cb::TensorBoardCallback)( lg = cb.logger filter_fn = Base.Fix1(filter_name_and_value, cb) - if iteration == 1 && cb.include_hyperparams - hp_iter = _names_and_values( - model, - sampler, - transition, - state; - params=false, - hyperparams=true, - extra=false, - kwargs..., - ) - hparams = Dict(hp_iter) - if !isempty(hparams) - TensorBoardLogger.write_hparams!( - lg, hparams, AbstractMCMC.hyperparam_metrics(model, sampler) - ) - end - end - with_logger(lg) do - all_values = _names_and_values( + # Use ParamsWithStats container with Base.pairs iteration + pws = ParamsWithStats( model, sampler, transition, state; params=true, - hyperparams=false, - extra=cb.include_extras, - kwargs..., + stats=cb.include_stats, + extras=(iteration == 1 && cb.include_extras), ) - for (k, val) in Iterators.filter(filter_fn, all_values) + for (k, val) in Iterators.filter(filter_fn, Base.pairs(pws)) @info "$(cb.param_prefix)$k" val if stats !== nothing diff --git a/src/AbstractMCMC.jl b/src/AbstractMCMC.jl index 79398d24..ade29651 100644 --- a/src/AbstractMCMC.jl +++ b/src/AbstractMCMC.jl @@ -23,7 +23,7 @@ export sample export MCMCThreads, MCMCDistributed, MCMCSerial # Callback API -export mcmc_callback +export mcmc_callback, ParamsWithStats """ AbstractChains diff --git a/src/callbacks.jl b/src/callbacks.jl index f7138385..393f54a6 100644 --- a/src/callbacks.jl +++ b/src/callbacks.jl @@ -52,7 +52,8 @@ end (f::NameFilter)(name, value) = f(name) function (f::NameFilter)(name) - return name ∉ f.exclude && (isempty(f.include) || name ∈ f.include) + str_name = string(name) + return str_name ∉ f.exclude && (isempty(f.include) || str_name ∈ f.include) end ############################## @@ -62,7 +63,7 @@ end const DEFAULT_STATS_OPTIONS = (; thin=0, skip=0, window=typemax(Int)) const DEFAULT_NAME_FILTER = (; - include=String[], exclude=String[], extras=false, hyperparams=false + include=String[], exclude=String[], stats=false, extras=false ) """ @@ -115,90 +116,125 @@ end ################################ """ - default_param_names_for_values(x) + ParamsWithStats{P,S,E} -Return an iterator of `θ[i]` for each element in `x`. -""" -default_param_names_for_values(x) = ("θ[$i]" for i in 1:length(x)) +A container for MCMC parameters, statistics, and extras. + +All fields are stored as `NamedTuple`s to ensure a tight, well-defined interface. +Use `Base.pairs(pws)` to iterate over `(name, value)` pairs. + +# Fields +- `params::P`: Parameter values as a NamedTuple +- `stats::S`: Statistics as a NamedTuple (e.g., `(lp=...,)`) +- `extras::E`: Extra diagnostics as a NamedTuple + +# Example +```julia +pws = ParamsWithStats(model, sampler, transition, state; params=true, stats=true) +for (name, value) in Base.pairs(pws) + println("\$name: \$value") +end +# Re-select to exclude stats: +pws2 = ParamsWithStats(pws; params=true, stats=false) +``` """ - _names_and_values( - model, - sampler, - transition, - state; - params::Bool = true, - hyperparams::Bool = false, - extra::Bool = false, - kwargs... - ) +struct ParamsWithStats{P<:NamedTuple,S<:NamedTuple,E<:NamedTuple} + params::P + stats::S + extras::E +end -Return an iterator over parameter names and values. +# Constructor from Vector{<:Real} - adds default θ[i] names +function ParamsWithStats( + v::AbstractVector{<:Real}, stats::S, extras::E +) where {S<:NamedTuple,E<:NamedTuple} + names = ntuple(i -> Symbol("θ[$i]"), length(v)) + params = NamedTuple{names}(Tuple(v)) + return ParamsWithStats(params, stats, extras) +end -This function is not part of the public API and may change or break at any time. +# Constructor from Vector{Pair} - converts to NamedTuple +function ParamsWithStats( + v::AbstractVector{<:Pair}, stats::S, extras::E +) where {S<:NamedTuple,E<:NamedTuple} + names = Tuple(Symbol(first(p)) for p in v) + values = Tuple(last(p) for p in v) + params = NamedTuple{names}(values) + return ParamsWithStats(params, stats, extras) +end + +# Constructor for nothing params (when params=false) +function ParamsWithStats(::Nothing, stats::S, extras::E) where {S<:NamedTuple,E<:NamedTuple} + return ParamsWithStats(NamedTuple(), stats, extras) +end -## Keywords -- `params`: include model parameters. -- `hyperparams`: include sampler hyperparameters -- `extra`: include additional statistics. -- `kwargs...`: reserved for internal extensibility. """ -function _names_and_values( + ParamsWithStats(model, sampler, transition, state; params=true, stats=false, extras=false) + +Construct a `ParamsWithStats` by extracting values from the MCMC state. + +# Arguments +- `params=true`: Include model parameters via `getparams(state)`. +- `stats=true`: Include step-level statistics via `getstats(state)`. These are values that + change once per MCMC iteration (e.g., log probability, acceptance rate). +- `extras=true`: Include extra diagnostics. These are values that remain constant across + MCMC iterations (e.g., preconditioning matrix, number of particles) or change multiple + times within a single iteration (e.g., leapfrog phase points in HMC). +""" +function ParamsWithStats( model, sampler, transition, state; params::Bool=true, - hyperparams::Bool=false, - extra::Bool=false, - kwargs..., + stats::Bool=false, + extras::Bool=false, ) - iters = [] - - if params - try - p = getparams(state) - push!(iters, zip(default_param_names_for_values(p), p)) - catch - # No params available - end - end - - if hyperparams - hp = _hyperparams_impl(model, sampler, state; kwargs...) - if !isempty(hp) - push!(iters, hp) - end - end + p = params ? getparams(state) : nothing + s = stats ? getstats(state) : NamedTuple() + e = extras ? NamedTuple() : NamedTuple() + return ParamsWithStats(p, s, e) +end - if extra - try - stats = getstats(state) - if stats isa NamedTuple - push!(iters, pairs(stats)) - end - catch - # No extras available - end - end +""" + ParamsWithStats(pws::ParamsWithStats; params=true, stats=true, extras=true) - return Iterators.flatten(iters) -end +Create a new `ParamsWithStats` by selecting subsets of an existing one. -# Internal helper for hyperparams extraction -function _hyperparams_impl(model, sampler, state; kwargs...) - return Pair{String,Any}[] +This enables filtering without re-extracting from state: +```julia +pws = ParamsWithStats(model, sampler, transition, state; params=true, stats=true) +pws_params_only = ParamsWithStats(pws; params=true, stats=false, extras=false) +``` +""" +function ParamsWithStats( + pws::ParamsWithStats; params::Bool=true, stats::Bool=true, extras::Bool=true +) + p = params ? pws.params : NamedTuple() + s = stats ? pws.stats : NamedTuple() + e = extras ? pws.extras : NamedTuple() + return ParamsWithStats(p, s, e) end """ - hyperparam_metrics(model, sampler[, state]; kwargs...) + Base.pairs(pws::ParamsWithStats) -Return a Vector{String} of metrics for hyperparameters. -Override this to specify which logged values should be used as hyperparam metrics in TensorBoard. +Return an iterator of `(name, value)` pairs for all selected data in `pws`. + +This is the canonical way to iterate over a `ParamsWithStats`: +```julia +for (name, value) in Base.pairs(pws) + @info name value +end +``` """ -hyperparam_metrics(model, sampler; kwargs...) = String[] -function hyperparam_metrics(model, sampler, state; kwargs...) - return hyperparam_metrics(model, sampler; kwargs...) +function Base.pairs(pws::ParamsWithStats) + return Iterators.flatten((pairs(pws.params), pairs(pws.stats), pairs(pws.extras))) +end + +function Base.isempty(pws::ParamsWithStats) + return (isempty(pws.params) && isempty(pws.stats) && isempty(pws.extras)) end ################################# @@ -242,7 +278,7 @@ Create a TensorBoard logging callback. **Requires TensorBoardLogger.jl to be loa - `true` or `:default`: Use default statistics (Mean, Variance, KHist) - requires OnlineStats - An OnlineStat or tuple of OnlineStats - requires OnlineStats - `stats_options`: NamedTuple with `thin`, `skip`, `window` -- `name_filter`: NamedTuple with `include`, `exclude`, `extras`, `hyperparams` +- `name_filter`: NamedTuple with `include`, `exclude`, `stats`, `hyperparams` # Examples ```julia diff --git a/test/callbacks.jl b/test/callbacks.jl index df791fa8..24903b55 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -93,8 +93,8 @@ end @testset "DEFAULT_NAME_FILTER" begin @test AbstractMCMC.DEFAULT_NAME_FILTER.include == String[] @test AbstractMCMC.DEFAULT_NAME_FILTER.exclude == String[] + @test AbstractMCMC.DEFAULT_NAME_FILTER.stats == false @test AbstractMCMC.DEFAULT_NAME_FILTER.extras == false - @test AbstractMCMC.DEFAULT_NAME_FILTER.hyperparams == false end end @@ -154,16 +154,88 @@ end @test f("a", 1.0) == true @test f("c", 2.0) == false end -end -@testset "default_param_names_for_values" begin - names = collect(AbstractMCMC.default_param_names_for_values([1.0, 2.0, 3.0])) - @test names == ["θ[1]", "θ[2]", "θ[3]"] + @testset "Symbol names (from NamedTuple iteration)" begin + f = AbstractMCMC.NameFilter(; include=["a", "b"]) + @test f(:a) == true + @test f(:c) == false + end end -@testset "_names_and_values" begin - # Test that the internal unified function exists and has expected signature - @test hasmethod(AbstractMCMC._names_and_values, Tuple{Any,Any,Any,Any}) +######################### +### ParamsWithStats ### +######################### + +@testset "ParamsWithStats" begin + @testset "Constructor from NamedTuple" begin + pws = AbstractMCMC.ParamsWithStats((a=1.0, b=2.0), (lp=-10.0,), NamedTuple()) + @test pws isa AbstractMCMC.ParamsWithStats + @test pws.params == (a=1.0, b=2.0) + @test pws.stats == (lp=-10.0,) + @test pws.extras == NamedTuple() + 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) + end + + @testset "Constructor from Vector{Pair} - named" begin + pws = AbstractMCMC.ParamsWithStats( + ["μ" => 1.0, "σ" => 2.0], NamedTuple(), NamedTuple() + ) + @test pws.params == (μ=1.0, σ=2.0) + end + + @testset "Constructor from state" begin + state = 5 + pws = AbstractMCMC.ParamsWithStats( + MyModel(), MySampler(), nothing, state; params=true, stats=true + ) + @test pws isa AbstractMCMC.ParamsWithStats + @test pws.params == NamedTuple() # Empty vector becomes empty NamedTuple + @test pws.stats == (iteration=5,) + @test pws.extras == NamedTuple() + end + + @testset "Copy constructor with selection" begin + pws = AbstractMCMC.ParamsWithStats((a=1.0,), (lp=-10.0,), NamedTuple()) + + # Select only params + pws_params = AbstractMCMC.ParamsWithStats(pws; params=true, stats=false) + @test pws_params.params == (a=1.0,) + @test pws_params.stats == NamedTuple() + + # Select only stats + pws_stats = AbstractMCMC.ParamsWithStats(pws; params=false, stats=true) + @test pws_stats.params == NamedTuple() + @test pws_stats.stats == (lp=-10.0,) + end + + @testset "Base.pairs iteration" begin + pws = AbstractMCMC.ParamsWithStats((a=1.0, b=2.0), (lp=-10.0,), NamedTuple()) + pairs_list = collect(Base.pairs(pws)) + @test length(pairs_list) == 3 + @test (:a => 1.0) in pairs_list + @test (:b => 2.0) in pairs_list + @test (:lp => -10.0) in pairs_list + end + + @testset "Base.isempty" begin + pws_full = AbstractMCMC.ParamsWithStats((a=1.0,), (lp=-10.0,), NamedTuple()) + @test !isempty(pws_full) + + pws_empty = AbstractMCMC.ParamsWithStats(NamedTuple(), NamedTuple(), NamedTuple()) + @test isempty(pws_empty) + end + + @testset "Illegal states are unrepresentable" begin + # Should not be able to construct with arbitrary types + @test_throws MethodError AbstractMCMC.ParamsWithStats(1, 2, 3) + @test_throws MethodError AbstractMCMC.ParamsWithStats( + "bad", NamedTuple(), NamedTuple() + ) + end end using OnlineStats @@ -234,7 +306,7 @@ using TensorBoardLogger cb = mcmc_callback(; logger=logger, name_filter=( - include=["mu", "sigma"], exclude=["internal"], extras=true, hyperparams=true + include=["mu", "sigma"], exclude=["internal"], stats=true, extras=true ), ) @test cb isa AbstractMCMC.MultiCallback @@ -286,6 +358,21 @@ using TensorBoardLogger cb = mcmc_callback(; logger=logger, stats=Mean()) @test cb isa AbstractMCMC.MultiCallback end + + @testset "ParamsWithStats default implementation" begin + struct MockState + params::Vector{Float64} + end + AbstractMCMC.getparams(s::MockState) = s.params + + state = MockState([10.0, 20.0]) + pws = AbstractMCMC.ParamsWithStats(nothing, nothing, nothing, state; params=true) + result = collect(Base.pairs(pws)) + + @test length(result) == 2 + @test result[1] == (Symbol("θ[1]") => 10.0) + @test result[2] == (Symbol("θ[2]") => 20.0) + end end ######################### diff --git a/test/utils.jl b/test/utils.jl index 9ed92ba7..35fa5ed9 100644 --- a/test/utils.jl +++ b/test/utils.jl @@ -134,3 +134,6 @@ function AbstractMCMC.step( _state = state === nothing ? 1 : state + 1 return MySample(θ, logdensity_θ), _state end + +AbstractMCMC.getparams(state::Integer) = Float64[] +AbstractMCMC.getstats(state::Integer) = (iteration=state,)