Improve Callbacks

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"

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.

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

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

src/AbstractMCMC.jl

@@ -23,7 +23,7 @@ export sample
 export MCMCThreads, MCMCDistributed, MCMCSerial
 
 # Callback API
-export mcmc_callback
+export mcmc_callback, ParamsWithStats
 
 """
     AbstractChains