From 12ff8125dc84e2d5021b0660c8420d4e8cb5f413 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 18 Jan 2026 00:55:48 +0530 Subject: [PATCH 01/13] getparams can return named pairs or just values to avoid double-nested pairs & Symbol --> String & KHist require Float64 --- ext/AbstractMCMCOnlineStatsExt.jl | 17 +++++++++++++---- src/callbacks.jl | 8 ++++++-- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/ext/AbstractMCMCOnlineStatsExt.jl b/ext/AbstractMCMCOnlineStatsExt.jl index c9020df1..bc7e4747 100644 --- a/ext/AbstractMCMCOnlineStatsExt.jl +++ b/ext/AbstractMCMCOnlineStatsExt.jl @@ -145,17 +145,26 @@ end Update and log statistics. Called from TensorBoard callback. """ function log_stat_impl!(stats::AbstractDict, prototype, key, val, prefix) + actual_val = val isa Pair ? last(val) : val + + if !(actual_val isa Real) + return nothing + end + float_val = Float64(actual_val) + + 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 diff --git a/src/callbacks.jl b/src/callbacks.jl index f7138385..bd5fb300 100644 --- a/src/callbacks.jl +++ b/src/callbacks.jl @@ -158,7 +158,11 @@ function _names_and_values( if params try p = getparams(state) - push!(iters, zip(default_param_names_for_values(p), p)) + if !isempty(p) && first(p) isa Pair + push!(iters, p) + else + push!(iters, zip(default_param_names_for_values(p), p)) + end catch # No params available end @@ -175,7 +179,7 @@ function _names_and_values( try stats = getstats(state) if stats isa NamedTuple - push!(iters, pairs(stats)) + push!(iters, (string(k) => v for (k, v) in pairs(stats))) end catch # No extras available From ffcc0b50153b30dd840a4b1c38a84f7815eb1249 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 18 Jan 2026 00:56:17 +0530 Subject: [PATCH 02/13] bump version --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 13a1bc47..6bd1f144 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.11.1" [deps] BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" From 4d4616fd976638a5cef9cb442c8dc19300caeece Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 18 Jan 2026 02:03:35 +0530 Subject: [PATCH 03/13] Add safety comment for OnlineStats Extenstion --- ext/AbstractMCMCOnlineStatsExt.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/ext/AbstractMCMCOnlineStatsExt.jl b/ext/AbstractMCMCOnlineStatsExt.jl index bc7e4747..6d373941 100644 --- a/ext/AbstractMCMCOnlineStatsExt.jl +++ b/ext/AbstractMCMCOnlineStatsExt.jl @@ -145,6 +145,7 @@ end Update and log statistics. Called from TensorBoard callback. """ function log_stat_impl!(stats::AbstractDict, prototype, key, val, prefix) + # Safety: extract value if val is a Pair (can happen with nested iteration) actual_val = val isa Pair ? last(val) : val if !(actual_val isa Real) From 2aa0f90e685793780e229d45a9bd66128744e5d0 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Fri, 23 Jan 2026 10:45:01 +0530 Subject: [PATCH 04/13] export names_and_values as public API and removed hyperparam_metrics entirely --- docs/src/callbacks.md | 77 ++++++++++++++++++----- ext/AbstractMCMCTensorBoardLoggerExt.jl | 39 ++++-------- src/AbstractMCMC.jl | 2 +- src/callbacks.jl | 83 ++++++++++--------------- test/callbacks.jl | 21 +++++-- 5 files changed, 121 insertions(+), 101 deletions(-) diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md index 7745e480..5dfd9725 100644 --- a/docs/src/callbacks.md +++ b/docs/src/callbacks.md @@ -211,6 +211,67 @@ function my_callback(rng, model, sampler, transition, state, iteration; kwargs.. end ``` +## names_and_values - Public API + +The `names_and_values` function is the **public API** for extracting named values from MCMC states. +Override this in downstream packages to provide meaningful variable names for logging and visualization. + +### Signature + +```julia +names_and_values(model, sampler, transition, state; + params=true, stats=false, hyperparams=false, extras=false) +``` + +Returns an iterator of `(name, value)` pairs. + +### Default Behavior + +- Uses `getparams(state)` for parameter values with `θ[1], θ[2], ...` naming +- Uses `getstats(state)` for extra statistics when `stats=true` +- Returns empty for hyperparameters (samplers should override) +- Returns empty for extras (samplers should override) + +### Overriding for Your Package + +```julia +function AbstractMCMC.names_and_values( + model::MyPackage.MyModel, sampler, transition, state; + params=true, stats=false, hyperparams=false +) + iters = [] + + if params + # Return actual variable names + push!(iters, [ + "μ" => state.mu, + "σ" => state.sigma, + ]) + end + + if stats + s = getstats(state) + push!(iters, (string(k) => v for (k, v) in pairs(s))) + end + + if hyperparams + push!(iters, ["step_size" => sampler.step_size]) + end + + return Iterators.flatten(iters) +end +``` + +### Usage in TensorBoard Callback + +The TensorBoard callback uses `names_and_values` internally: + +```julia +for (k, val) in names_and_values(model, sampler, t, state; params=true, stats=true) + @info "$k" val +end +``` + ## Internals !!! note @@ -236,19 +297,3 @@ When using statistics, AbstractMCMC provides wrappers that modify how samples ar 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/AbstractMCMCTensorBoardLoggerExt.jl b/ext/AbstractMCMCTensorBoardLoggerExt.jl index ab60e6d2..51bb3cbf 100644 --- a/ext/AbstractMCMCTensorBoardLoggerExt.jl +++ b/ext/AbstractMCMCTensorBoardLoggerExt.jl @@ -4,8 +4,7 @@ using AbstractMCMC using AbstractMCMC: MultiCallback, NameFilter, - _names_and_values, - hyperparam_metrics, + names_and_values, 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,7 +28,7 @@ struct TensorBoardCallback{L,S,P,F} stats::S stat_prototype::P variable_filter::F - include_extras::Bool + include_stats::Bool include_hyperparams::Bool param_prefix::String extras_prefix::String @@ -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`, `hyperparams` - `num_bins`: Number of histogram bins (default: 100) # Examples @@ -85,7 +88,7 @@ function AbstractMCMC.mcmc_callback(; stats_dict, prototype, variable_filter, - merged_name_filter.extras, + merged_name_filter.stats, merged_name_filter.hyperparams, "", "extras/", @@ -105,35 +108,15 @@ 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( + all_values = names_and_values( model, sampler, transition, state; params=true, - hyperparams=false, - extra=cb.include_extras, - kwargs..., + stats=cb.include_stats, + hyperparams=(iteration == 1 && cb.include_hyperparams), ) for (k, val) in Iterators.filter(filter_fn, all_values) diff --git a/src/AbstractMCMC.jl b/src/AbstractMCMC.jl index 79398d24..5b2df093 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, names_and_values """ AbstractChains diff --git a/src/callbacks.jl b/src/callbacks.jl index bd5fb300..c9af8cd5 100644 --- a/src/callbacks.jl +++ b/src/callbacks.jl @@ -62,7 +62,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, hyperparams=false ) """ @@ -122,36 +122,37 @@ Return an iterator of `θ[i]` for each element in `x`. default_param_names_for_values(x) = ("θ[$i]" for i in 1:length(x)) """ - _names_and_values( - model, - sampler, - transition, - state; - params::Bool = true, - hyperparams::Bool = false, - extra::Bool = false, - kwargs... - ) + names_and_values(model, sampler, transition, state; params=true, stats=false, hyperparams=false, extras=false) -Return an iterator over parameter names and values. +Return an iterator of `(name, value)` pairs for MCMC logging and visualization. -This function is not part of the public API and may change or break at any time. +This is the **public API** for extracting named values from MCMC states. +Downstream packages should override this to provide meaningful variable names. +If not overridden, it defaults to: +- `params=true`: `θ[1], θ[2], ...` from `getparams(state)` +- `stats=true`: `lp`, etc. from `getstats(state)` +- `hyperparams=true`: empty (unless overridden) +- `extras=true`: empty (unless overridden) -## Keywords -- `params`: include model parameters. -- `hyperparams`: include sampler hyperparameters -- `extra`: include additional statistics. -- `kwargs...`: reserved for internal extensibility. +# Arguments +- `model`: The probabilistic model being sampled +- `sampler`: The MCMC sampler +- `transition`: The current transition +- `state`: The current sampler state +- `params::Bool=true`: Include model parameters +- `stats::Bool=false`: Include extra statistics (e.g., log probability) +- `hyperparams::Bool=false`: Include sampler hyperparameters +- `extras::Bool=false`: Include extra transition information """ -function _names_and_values( +function names_and_values( model, sampler, transition, state; params::Bool=true, + stats::Bool=false, hyperparams::Bool=false, - extra::Bool=false, - kwargs..., + extras::Bool=false, ) iters = [] @@ -159,8 +160,10 @@ function _names_and_values( try p = getparams(state) if !isempty(p) && first(p) isa Pair + # Already named pairs - use directly push!(iters, p) else + # Raw values - add default θ[i] names push!(iters, zip(default_param_names_for_values(p), p)) end catch @@ -168,41 +171,21 @@ function _names_and_values( end end - if hyperparams - hp = _hyperparams_impl(model, sampler, state; kwargs...) - if !isempty(hp) - push!(iters, hp) - end - end - - if extra + if stats try - stats = getstats(state) - if stats isa NamedTuple - push!(iters, (string(k) => v for (k, v) in pairs(stats))) + s = getstats(state) + if s isa NamedTuple && !isempty(s) + push!(iters, (string(k) => v for (k, v) in pairs(s))) end catch - # No extras available + # No stats available end end - return Iterators.flatten(iters) -end - -# Internal helper for hyperparams extraction -function _hyperparams_impl(model, sampler, state; kwargs...) - return Pair{String,Any}[] -end + # hyperparams and extras: default returns empty + # samplers should override names_and_values to provide these -""" - hyperparam_metrics(model, sampler[, state]; kwargs...) - -Return a Vector{String} of metrics for hyperparameters. -Override this to specify which logged values should be used as hyperparam metrics in TensorBoard. -""" -hyperparam_metrics(model, sampler; kwargs...) = String[] -function hyperparam_metrics(model, sampler, state; kwargs...) - return hyperparam_metrics(model, sampler; kwargs...) + return Iterators.flatten(iters) end ################################# @@ -246,7 +229,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..b6fecd79 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -93,7 +93,7 @@ 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.extras == false + @test AbstractMCMC.DEFAULT_NAME_FILTER.stats == false @test AbstractMCMC.DEFAULT_NAME_FILTER.hyperparams == false end end @@ -161,11 +161,6 @@ end @test names == ["θ[1]", "θ[2]", "θ[3]"] 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}) -end - using OnlineStats ############################# @@ -286,6 +281,20 @@ using TensorBoardLogger cb = mcmc_callback(; logger=logger, stats=Mean()) @test cb isa AbstractMCMC.MultiCallback end + + @testset "names_and_values default implementation" begin + struct MockState + params::Vector{Float64} + end + AbstractMCMC.getparams(s::MockState) = s.params + + state = MockState([10.0, 20.0]) + result = collect(AbstractMCMC.names_and_values(nothing, nothing, nothing, state)) + + @test length(result) == 2 + @test result[1] == ("θ[1]", 10.0) + @test result[2] == ("θ[2]", 20.0) + end end ######################### From af00196c60287a30e39fdf4f5e54dd983afe719d Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Fri, 23 Jan 2026 11:11:56 +0530 Subject: [PATCH 05/13] update callbacks.md --- docs/src/callbacks.md | 22 +++------------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md index 5dfd9725..90727393 100644 --- a/docs/src/callbacks.md +++ b/docs/src/callbacks.md @@ -213,31 +213,16 @@ end ## names_and_values - Public API -The `names_and_values` function is the **public API** for extracting named values from MCMC states. -Override this in downstream packages to provide meaningful variable names for logging and visualization. - -### Signature - -```julia -names_and_values(model, sampler, transition, state; - params=true, stats=false, hyperparams=false, extras=false) +```@docs +AbstractMCMC.names_and_values ``` -Returns an iterator of `(name, value)` pairs. - -### Default Behavior - -- Uses `getparams(state)` for parameter values with `θ[1], θ[2], ...` naming -- Uses `getstats(state)` for extra statistics when `stats=true` -- Returns empty for hyperparameters (samplers should override) -- Returns empty for extras (samplers should override) - ### Overriding for Your Package ```julia function AbstractMCMC.names_and_values( model::MyPackage.MyModel, sampler, transition, state; - params=true, stats=false, hyperparams=false + params=true, stats=false, hyperparams=false, extras=false ) iters = [] @@ -296,4 +281,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. - From ee079845989f72bb8c61dd25f637ecbaf7e321e2 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 25 Jan 2026 15:38:51 +0530 Subject: [PATCH 06/13] replace names_and_values with ParamsWithStats --- Project.toml | 2 +- docs/src/callbacks.md | 63 ++++++----- ext/AbstractMCMCOnlineStatsExt.jl | 20 +++- ext/AbstractMCMCTensorBoardLoggerExt.jl | 15 +-- src/AbstractMCMC.jl | 2 +- src/callbacks.jl | 140 ++++++++++++++++-------- test/callbacks.jl | 62 ++++++++++- test/utils.jl | 3 + 8 files changed, 216 insertions(+), 91 deletions(-) diff --git a/Project.toml b/Project.toml index 6bd1f144..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.1" +version = "5.12.0" [deps] BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md index 90727393..a78c5d5f 100644 --- a/docs/src/callbacks.md +++ b/docs/src/callbacks.md @@ -198,8 +198,7 @@ mcmc_callback |--------------|------------|----------------------------------| | `include` | `String[]` | Only log these (empty=all) | | `exclude` | `String[]` | Don't log these | -| `extras` | `false` | Include extra stats | -| `hyperparams`| `false` | Include hyperparameters | +| `extras` | `false` | Include extra diagnostics | ## Implementing Custom Callbacks @@ -211,48 +210,52 @@ function my_callback(rng, model, sampler, transition, state, iteration; kwargs.. end ``` -## names_and_values - Public API +## ParamsWithStats - Public API + +`ParamsWithStats` is the first-class container for extracting and iterating over MCMC parameters, statistics, and extras. ```@docs -AbstractMCMC.names_and_values +AbstractMCMC.ParamsWithStats +``` + +### 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 -function AbstractMCMC.names_and_values( - model::MyPackage.MyModel, sampler, transition, state; - params=true, stats=false, hyperparams=false, extras=false -) - iters = [] - - if params - # Return actual variable names - push!(iters, [ - "μ" => state.mu, - "σ" => state.sigma, - ]) - end - - if stats - s = getstats(state) - push!(iters, (string(k) => v for (k, v) in pairs(s))) - end - - if hyperparams - push!(iters, ["step_size" => sampler.step_size]) - end - - return Iterators.flatten(iters) +# Override getparams to return named pairs +function AbstractMCMC.getparams(state::MyState) + return ["μ" => state.mu, "σ" => state.sigma] +end + +# Override getstats to return statistics +function AbstractMCMC.getstats(state::MyState) + return (lp=state.logp, acceptance_rate=state.accept_rate) end ``` ### Usage in TensorBoard Callback -The TensorBoard callback uses `names_and_values` internally: +The TensorBoard callback uses `ParamsWithStats` with `Base.pairs`: ```julia -for (k, val) in names_and_values(model, sampler, t, state; params=true, stats=true) +pws = ParamsWithStats(model, sampler, t, state; params=true, stats=true) +for (k, val) in Base.pairs(pws) @info "$k" val end ``` diff --git a/ext/AbstractMCMCOnlineStatsExt.jl b/ext/AbstractMCMCOnlineStatsExt.jl index 6d373941..ab7e7b5e 100644 --- a/ext/AbstractMCMCOnlineStatsExt.jl +++ b/ext/AbstractMCMCOnlineStatsExt.jl @@ -144,14 +144,12 @@ end Update and log statistics. Called from TensorBoard callback. """ -function log_stat_impl!(stats::AbstractDict, prototype, key, val, prefix) - # Safety: extract value if val is a Pair (can happen with nested iteration) - actual_val = val isa Pair ? last(val) : val - - if !(actual_val isa Real) +function log_stat_impl!(stats::AbstractDict, prototype, key, val::Real, prefix) + float_val = try + Float64(val) + catch return nothing end - float_val = Float64(actual_val) str_key = string(key) @@ -169,6 +167,16 @@ function log_stat_impl!(stats::AbstractDict, prototype, key, val, prefix) 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 51bb3cbf..0839bf72 100644 --- a/ext/AbstractMCMCTensorBoardLoggerExt.jl +++ b/ext/AbstractMCMCTensorBoardLoggerExt.jl @@ -4,7 +4,7 @@ using AbstractMCMC using AbstractMCMC: MultiCallback, NameFilter, - names_and_values, + ParamsWithStats, merge_with_defaults, create_stats_with_options, DEFAULT_STATS_OPTIONS, @@ -29,7 +29,7 @@ struct TensorBoardCallback{L,S,P,F} stat_prototype::P variable_filter::F include_stats::Bool - include_hyperparams::Bool + include_extras::Bool param_prefix::String extras_prefix::String end @@ -46,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`, `stats`, `hyperparams` +- `name_filter`: NamedTuple with `include`, `exclude`, `stats`, `extras` - `num_bins`: Number of histogram bins (default: 100) # Examples @@ -89,7 +89,7 @@ function AbstractMCMC.mcmc_callback(; prototype, variable_filter, merged_name_filter.stats, - merged_name_filter.hyperparams, + merged_name_filter.extras, "", "extras/", ) @@ -109,17 +109,18 @@ function (cb::TensorBoardCallback)( filter_fn = Base.Fix1(filter_name_and_value, cb) with_logger(lg) do - all_values = names_and_values( + # Use ParamsWithStats container with Base.pairs iteration + pws = ParamsWithStats( model, sampler, transition, state; params=true, stats=cb.include_stats, - hyperparams=(iteration == 1 && cb.include_hyperparams), + 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 5b2df093..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, names_and_values +export mcmc_callback, ParamsWithStats """ AbstractChains diff --git a/src/callbacks.jl b/src/callbacks.jl index c9af8cd5..a5e8da57 100644 --- a/src/callbacks.jl +++ b/src/callbacks.jl @@ -62,7 +62,7 @@ end const DEFAULT_STATS_OPTIONS = (; thin=0, skip=0, window=typemax(Int)) const DEFAULT_NAME_FILTER = (; - include=String[], exclude=String[], stats=false, hyperparams=false + include=String[], exclude=String[], stats=false, extras=false ) """ @@ -122,72 +122,126 @@ Return an iterator of `θ[i]` for each element in `x`. default_param_names_for_values(x) = ("θ[$i]" for i in 1:length(x)) """ - names_and_values(model, sampler, transition, state; params=true, stats=false, hyperparams=false, extras=false) + ParamsWithStats{P,S,E} -Return an iterator of `(name, value)` pairs for MCMC logging and visualization. +A container for MCMC parameters, statistics, and extras. This is the **public API** for extracting named values from MCMC states. -Downstream packages should override this to provide meaningful variable names. -If not overridden, it defaults to: -- `params=true`: `θ[1], θ[2], ...` from `getparams(state)` -- `stats=true`: `lp`, etc. from `getstats(state)` -- `hyperparams=true`: empty (unless overridden) -- `extras=true`: empty (unless overridden) +Use `Base.pairs(pws)` to iterate over `(name, value)` pairs. + +# Fields +- `params::P`: Parameter values (Vector, Dict, or OrderedDict) +- `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) +``` +""" +struct ParamsWithStats{P,S,E} + params::P + stats::S + extras::E +end -# Arguments -- `model`: The probabilistic model being sampled -- `sampler`: The MCMC sampler -- `transition`: The current transition -- `state`: The current sampler state -- `params::Bool=true`: Include model parameters -- `stats::Bool=false`: Include extra statistics (e.g., log probability) -- `hyperparams::Bool=false`: Include sampler hyperparameters -- `extras::Bool=false`: Include extra transition information """ -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. + +- `params=true`: Include model parameters via `getparams(state)` +- `stats=true`: Include statistics via `getstats(state)` +- `extras=true`: Include extras (empty by default; samplers override to provide) +""" +function ParamsWithStats( model, sampler, transition, state; params::Bool=true, stats::Bool=false, - hyperparams::Bool=false, extras::Bool=false, ) + p = params ? getparams(state) : nothing + s = stats ? getstats(state) : NamedTuple() + e = extras ? NamedTuple() : NamedTuple() # Samplers can override for actual extras + return ParamsWithStats(p, s, e) +end + +""" + ParamsWithStats(pws::ParamsWithStats; params=true, stats=true, extras=true) + +Create a new `ParamsWithStats` by selecting subsets of an existing one. + +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 : nothing + s = stats ? pws.stats : NamedTuple() + e = extras ? pws.extras : NamedTuple() + return ParamsWithStats(p, s, e) +end + +""" + Base.pairs(pws::ParamsWithStats) + +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 +``` +""" +function Base.pairs(pws::ParamsWithStats) iters = [] - if params - try - p = getparams(state) - if !isempty(p) && first(p) isa Pair - # Already named pairs - use directly - push!(iters, p) - else - # Raw values - add default θ[i] names - push!(iters, zip(default_param_names_for_values(p), p)) - end - catch - # No params available + # Handle params + if pws.params !== nothing && !isempty(pws.params) + if first(pws.params) isa Pair + # Already named pairs - use directly + push!(iters, pws.params) + else + # Raw values - add default θ[i] names + push!(iters, zip(default_param_names_for_values(pws.params), pws.params)) end end - if stats - try - s = getstats(state) - if s isa NamedTuple && !isempty(s) - push!(iters, (string(k) => v for (k, v) in pairs(s))) - end - catch - # No stats available - end + # Handle stats + if !isempty(pws.stats) + push!(iters, (string(k) => v for (k, v) in Base.pairs(pws.stats))) end - # hyperparams and extras: default returns empty - # samplers should override names_and_values to provide these + # Handle extras + if !isempty(pws.extras) + push!(iters, (string(k) => v for (k, v) in Base.pairs(pws.extras))) + end return Iterators.flatten(iters) end +function Base.isempty(pws::ParamsWithStats) + return ( + (pws.params === nothing || isempty(pws.params)) && + isempty(pws.stats) && + isempty(pws.extras) + ) +end + ################################# ### Unified mcmc_callback API ### ################################# diff --git a/test/callbacks.jl b/test/callbacks.jl index b6fecd79..2f6b8bc5 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -94,7 +94,7 @@ end @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.hyperparams == false + @test AbstractMCMC.DEFAULT_NAME_FILTER.extras == false end end @@ -161,6 +161,61 @@ end @test names == ["θ[1]", "θ[2]", "θ[3]"] end +######################### +### ParamsWithStats ### +######################### + +@testset "ParamsWithStats" begin + @testset "Constructor from state" begin + # Use a simple Integer state (has getparams/getstats in utils.jl) + state = 5 + pws = AbstractMCMC.ParamsWithStats( + MyModel(), MySampler(), nothing, state; params=true, stats=true + ) + @test pws isa AbstractMCMC.ParamsWithStats + @test pws.params == Float64[] # Our Integer state returns empty params + @test pws.stats == (iteration=5,) # Our Integer state returns iteration + @test pws.extras == NamedTuple() + end + + @testset "Copy constructor with selection" begin + state = 5 + pws = AbstractMCMC.ParamsWithStats( + MyModel(), MySampler(), nothing, state; params=true, stats=true + ) + + # Select only params + pws_params = AbstractMCMC.ParamsWithStats(pws; params=true, stats=false) + @test pws_params.params == Float64[] + @test pws_params.stats == NamedTuple() + + # Select only stats + pws_stats = AbstractMCMC.ParamsWithStats(pws; params=false, stats=true) + @test pws_stats.params === nothing + @test pws_stats.stats == (iteration=5,) + end + + @testset "Base.pairs iteration" begin + state = 5 + pws = AbstractMCMC.ParamsWithStats( + MyModel(), MySampler(), nothing, state; params=true, stats=true + ) + + # Collect pairs + pairs_list = collect(Base.pairs(pws)) + @test length(pairs_list) == 1 # Only stats (iteration), no params + @test ("iteration" => 5) in pairs_list + end + + @testset "Base.isempty" begin + state = 5 + pws_full = AbstractMCMC.ParamsWithStats( + MyModel(), MySampler(), nothing, state; params=true, stats=true + ) + @test !isempty(pws_full) + end +end + using OnlineStats ############################# @@ -282,14 +337,15 @@ using TensorBoardLogger @test cb isa AbstractMCMC.MultiCallback end - @testset "names_and_values default implementation" begin + @testset "ParamsWithStats default implementation" begin struct MockState params::Vector{Float64} end AbstractMCMC.getparams(s::MockState) = s.params state = MockState([10.0, 20.0]) - result = collect(AbstractMCMC.names_and_values(nothing, nothing, nothing, state)) + pws = AbstractMCMC.ParamsWithStats(nothing, nothing, nothing, state; params=true) + result = collect(Base.pairs(pws)) @test length(result) == 2 @test result[1] == ("θ[1]", 10.0) 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,) From aeaaca7f7fb3214c7811641dc90847035f09e75f Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 25 Jan 2026 17:25:23 +0530 Subject: [PATCH 07/13] fix params handling --- src/callbacks.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/callbacks.jl b/src/callbacks.jl index a5e8da57..c75dca9f 100644 --- a/src/callbacks.jl +++ b/src/callbacks.jl @@ -156,9 +156,9 @@ end Construct a `ParamsWithStats` by extracting values from the MCMC state. -- `params=true`: Include model parameters via `getparams(state)` -- `stats=true`: Include statistics via `getstats(state)` -- `extras=true`: Include extras (empty by default; samplers override to provide) +- `params=true`: Include model parameters via `getparams(state)`. +- `stats=true`: Include step-level statistics via `getstats(state)`. +- `extras=true`: Include constant or iteration-level metadata (e.g. hyperparams). """ function ParamsWithStats( model, @@ -217,7 +217,7 @@ function Base.pairs(pws::ParamsWithStats) push!(iters, pws.params) else # Raw values - add default θ[i] names - push!(iters, zip(default_param_names_for_values(pws.params), pws.params)) + push!(iters, (n => v for (n, v) in zip(default_param_names_for_values(pws.params), pws.params))) end end From bd24e809b7396584e89ba5ad1cf3f7527bf60f58 Mon Sep 17 00:00:00 2001 From: Shravan Goswami <123811742+shravanngoswamii@users.noreply.github.com> Date: Sun, 25 Jan 2026 17:27:50 +0530 Subject: [PATCH 08/13] format src/callbacks.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/callbacks.jl | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/callbacks.jl b/src/callbacks.jl index c75dca9f..ca2b5bb3 100644 --- a/src/callbacks.jl +++ b/src/callbacks.jl @@ -217,7 +217,13 @@ function Base.pairs(pws::ParamsWithStats) push!(iters, pws.params) else # Raw values - add default θ[i] names - push!(iters, (n => v for (n, v) in zip(default_param_names_for_values(pws.params), pws.params))) + push!( + iters, + ( + n => v for + (n, v) in zip(default_param_names_for_values(pws.params), pws.params) + ), + ) end end From 37bf4969ee42354b081c69ca7f4a7f8ebb33a561 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 25 Jan 2026 22:58:10 +0530 Subject: [PATCH 09/13] fix iterator to yield Pairs --- test/callbacks.jl | 4 ++-- test/runtests.jl | 12 ++++++------ 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/test/callbacks.jl b/test/callbacks.jl index 2f6b8bc5..41894138 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -348,8 +348,8 @@ using TensorBoardLogger result = collect(Base.pairs(pws)) @test length(result) == 2 - @test result[1] == ("θ[1]", 10.0) - @test result[2] == ("θ[2]", 20.0) + @test result[1] == ("θ[1]" => 10.0) + @test result[2] == ("θ[2]" => 20.0) end end diff --git a/test/runtests.jl b/test/runtests.jl index c7c793a8..4ac2c88b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -20,12 +20,12 @@ const CURRENT_LOGGER = Logging.current_logger() include("utils.jl") -@testset "AbstractMCMC" begin - include("sample.jl") - include("stepper.jl") - include("transducer.jl") - include("logdensityproblems.jl") -end +# @testset "AbstractMCMC" begin +# include("sample.jl") +# include("stepper.jl") +# include("transducer.jl") +# include("logdensityproblems.jl") +# end @testset "Callbacks" begin include("callbacks.jl") From 943e377457f759ab1399e6a87653c70589c471b6 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Sun, 25 Jan 2026 22:58:40 +0530 Subject: [PATCH 10/13] uncomment tests --- test/runtests.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index 4ac2c88b..c7c793a8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -20,12 +20,12 @@ const CURRENT_LOGGER = Logging.current_logger() include("utils.jl") -# @testset "AbstractMCMC" begin -# include("sample.jl") -# include("stepper.jl") -# include("transducer.jl") -# include("logdensityproblems.jl") -# end +@testset "AbstractMCMC" begin + include("sample.jl") + include("stepper.jl") + include("transducer.jl") + include("logdensityproblems.jl") +end @testset "Callbacks" begin include("callbacks.jl") From 38a587f83e6fe99d73e8fa9b342a9ab3ff99080e Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Sun, 25 Jan 2026 20:08:23 +0000 Subject: [PATCH 11/13] Update callbacks.md --- docs/src/callbacks.md | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md index a78c5d5f..681207da 100644 --- a/docs/src/callbacks.md +++ b/docs/src/callbacks.md @@ -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 @@ -210,13 +209,9 @@ function my_callback(rng, model, sampler, transition, state, iteration; kwargs.. end ``` -## ParamsWithStats - Public API - -`ParamsWithStats` is the first-class container for extracting and iterating over MCMC parameters, statistics, and extras. +## ParamsWithStats -```@docs -AbstractMCMC.ParamsWithStats -``` +`ParamsWithStats` is a container for extracting and iterating over MCMC parameters, statistics, and extras. ### Basic Usage From 69ed6a03977e7438d0723d014484bbc6f748f874 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Mon, 26 Jan 2026 15:01:14 +0530 Subject: [PATCH 12/13] implement new suggestions --- docs/src/callbacks.md | 24 ++++++++++-- src/callbacks.jl | 91 +++++++++++++++++++++---------------------- test/callbacks.jl | 78 ++++++++++++++++++++++++++----------- 3 files changed, 120 insertions(+), 73 deletions(-) diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md index 681207da..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 ), ) ``` @@ -197,6 +197,7 @@ AbstractMCMC.ParamsWithStats |--------------|------------|----------------------------------| | `include` | `String[]` | Only log these (empty=all) | | `exclude` | `String[]` | Don't log these | +| `stats` | `false` | Include step-level statistics | | `extras` | `false` | Include extra diagnostics | ## Implementing Custom Callbacks @@ -233,17 +234,32 @@ pws_params = ParamsWithStats(pws; params=true, stats=false, extras=false) To provide meaningful variable names, override the extraction hooks: ```julia -# Override getparams to return named pairs +# 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 statistics +# 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`: diff --git a/src/callbacks.jl b/src/callbacks.jl index ca2b5bb3..7b9d69a2 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 ############################## @@ -114,23 +115,16 @@ end ### Parameter Extraction API ### ################################ -""" - default_param_names_for_values(x) - -Return an iterator of `θ[i]` for each element in `x`. -""" -default_param_names_for_values(x) = ("θ[$i]" for i in 1:length(x)) - """ ParamsWithStats{P,S,E} A container for MCMC parameters, statistics, and extras. -This is the **public API** for extracting named values from MCMC states. +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 (Vector, Dict, or OrderedDict) +- `params::P`: Parameter values as a NamedTuple - `stats::S`: Statistics as a NamedTuple (e.g., `(lp=...,)`) - `extras::E`: Extra diagnostics as a NamedTuple @@ -145,20 +139,50 @@ end pws2 = ParamsWithStats(pws; params=true, stats=false) ``` """ -struct ParamsWithStats{P,S,E} +struct ParamsWithStats{P<:NamedTuple,S<:NamedTuple,E<:NamedTuple} params::P stats::S extras::E end +# 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 + +# 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 + """ 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)`. -- `extras=true`: Include constant or iteration-level metadata (e.g. hyperparams). +- `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, @@ -171,7 +195,7 @@ function ParamsWithStats( ) p = params ? getparams(state) : nothing s = stats ? getstats(state) : NamedTuple() - e = extras ? NamedTuple() : NamedTuple() # Samplers can override for actual extras + e = extras ? NamedTuple() : NamedTuple() return ParamsWithStats(p, s, e) end @@ -189,7 +213,7 @@ 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 : nothing + p = params ? pws.params : NamedTuple() s = stats ? pws.stats : NamedTuple() e = extras ? pws.extras : NamedTuple() return ParamsWithStats(p, s, e) @@ -208,41 +232,16 @@ end ``` """ function Base.pairs(pws::ParamsWithStats) - iters = [] - - # Handle params - if pws.params !== nothing && !isempty(pws.params) - if first(pws.params) isa Pair - # Already named pairs - use directly - push!(iters, pws.params) - else - # Raw values - add default θ[i] names - push!( - iters, - ( - n => v for - (n, v) in zip(default_param_names_for_values(pws.params), pws.params) - ), - ) - end - end - - # Handle stats - if !isempty(pws.stats) - push!(iters, (string(k) => v for (k, v) in Base.pairs(pws.stats))) - end - - # Handle extras - if !isempty(pws.extras) - push!(iters, (string(k) => v for (k, v) in Base.pairs(pws.extras))) - end - - return Iterators.flatten(iters) + return Iterators.flatten(( + pairs(pws.params), + pairs(pws.stats), + pairs(pws.extras), + )) end function Base.isempty(pws::ParamsWithStats) return ( - (pws.params === nothing || isempty(pws.params)) && + isempty(pws.params) && isempty(pws.stats) && isempty(pws.extras) ) diff --git a/test/callbacks.jl b/test/callbacks.jl index 41894138..8a0dcc4f 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -154,11 +154,12 @@ 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 ######################### @@ -166,53 +167,84 @@ end ######################### @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 - # Use a simple Integer state (has getparams/getstats in utils.jl) state = 5 pws = AbstractMCMC.ParamsWithStats( MyModel(), MySampler(), nothing, state; params=true, stats=true ) @test pws isa AbstractMCMC.ParamsWithStats - @test pws.params == Float64[] # Our Integer state returns empty params - @test pws.stats == (iteration=5,) # Our Integer state returns iteration + @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 - state = 5 pws = AbstractMCMC.ParamsWithStats( - MyModel(), MySampler(), nothing, state; params=true, stats=true + (a=1.0,), (lp=-10.0,), NamedTuple() ) # Select only params pws_params = AbstractMCMC.ParamsWithStats(pws; params=true, stats=false) - @test pws_params.params == Float64[] + @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 === nothing - @test pws_stats.stats == (iteration=5,) + @test pws_stats.params == NamedTuple() + @test pws_stats.stats == (lp=-10.0,) end @testset "Base.pairs iteration" begin - state = 5 pws = AbstractMCMC.ParamsWithStats( - MyModel(), MySampler(), nothing, state; params=true, stats=true + (a=1.0, b=2.0), (lp=-10.0,), NamedTuple() ) - - # Collect pairs pairs_list = collect(Base.pairs(pws)) - @test length(pairs_list) == 1 # Only stats (iteration), no params - @test ("iteration" => 5) in pairs_list + @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 - state = 5 pws_full = AbstractMCMC.ParamsWithStats( - MyModel(), MySampler(), nothing, state; params=true, stats=true + (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 @@ -284,7 +316,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 @@ -348,8 +380,8 @@ using TensorBoardLogger result = collect(Base.pairs(pws)) @test length(result) == 2 - @test result[1] == ("θ[1]" => 10.0) - @test result[2] == ("θ[2]" => 20.0) + @test result[1] == (Symbol("θ[1]") => 10.0) + @test result[2] == (Symbol("θ[2]") => 20.0) end end From 279322063663f85ff340718df25e670ee38ff884 Mon Sep 17 00:00:00 2001 From: Shravan Goswami Date: Mon, 26 Jan 2026 15:05:27 +0530 Subject: [PATCH 13/13] format --- src/callbacks.jl | 16 +++------------- test/callbacks.jl | 28 +++++++++------------------- 2 files changed, 12 insertions(+), 32 deletions(-) diff --git a/src/callbacks.jl b/src/callbacks.jl index 7b9d69a2..393f54a6 100644 --- a/src/callbacks.jl +++ b/src/callbacks.jl @@ -165,9 +165,7 @@ function ParamsWithStats( end # Constructor for nothing params (when params=false) -function ParamsWithStats( - ::Nothing, stats::S, extras::E -) where {S<:NamedTuple,E<:NamedTuple} +function ParamsWithStats(::Nothing, stats::S, extras::E) where {S<:NamedTuple,E<:NamedTuple} return ParamsWithStats(NamedTuple(), stats, extras) end @@ -232,19 +230,11 @@ end ``` """ function Base.pairs(pws::ParamsWithStats) - return Iterators.flatten(( - pairs(pws.params), - pairs(pws.stats), - pairs(pws.extras), - )) + 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) - ) + return (isempty(pws.params) && isempty(pws.stats) && isempty(pws.extras)) end ################################# diff --git a/test/callbacks.jl b/test/callbacks.jl index 8a0dcc4f..24903b55 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -168,9 +168,7 @@ end @testset "ParamsWithStats" begin @testset "Constructor from NamedTuple" begin - pws = AbstractMCMC.ParamsWithStats( - (a=1.0, b=2.0), (lp=-10.0,), NamedTuple() - ) + 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,) @@ -178,9 +176,7 @@ end end @testset "Constructor from Vector{Real} - default names" begin - pws = AbstractMCMC.ParamsWithStats( - [1.0, 2.0, 3.0], NamedTuple(), NamedTuple() - ) + 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 @@ -203,9 +199,7 @@ end end @testset "Copy constructor with selection" begin - pws = AbstractMCMC.ParamsWithStats( - (a=1.0,), (lp=-10.0,), NamedTuple() - ) + pws = AbstractMCMC.ParamsWithStats((a=1.0,), (lp=-10.0,), NamedTuple()) # Select only params pws_params = AbstractMCMC.ParamsWithStats(pws; params=true, stats=false) @@ -219,9 +213,7 @@ end end @testset "Base.pairs iteration" begin - pws = AbstractMCMC.ParamsWithStats( - (a=1.0, b=2.0), (lp=-10.0,), NamedTuple() - ) + 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 @@ -230,21 +222,19 @@ end end @testset "Base.isempty" begin - pws_full = AbstractMCMC.ParamsWithStats( - (a=1.0,), (lp=-10.0,), NamedTuple() - ) + pws_full = AbstractMCMC.ParamsWithStats((a=1.0,), (lp=-10.0,), NamedTuple()) @test !isempty(pws_full) - pws_empty = AbstractMCMC.ParamsWithStats( - NamedTuple(), NamedTuple(), NamedTuple() - ) + 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()) + @test_throws MethodError AbstractMCMC.ParamsWithStats( + "bad", NamedTuple(), NamedTuple() + ) end end