diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 18ab53bc..a39a8579 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -23,13 +23,10 @@ jobs: fail-fast: false matrix: runner: + # x64 - version: '1' os: ubuntu-latest arch: x64 - # Minimum supported version - - version: 'min' - os: ubuntu-latest - arch: x64 # x86 - version: '1' os: ubuntu-latest diff --git a/.gitignore b/.gitignore index 83d89f72..89501654 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,5 @@ *.jl.*.cov *.jl.mem deps/deps.jl -Manifest.toml \ No newline at end of file +Manifest.toml +tmp/ \ No newline at end of file diff --git a/Project.toml b/Project.toml index 173bfaab..13a1bc47 100644 --- a/Project.toml +++ b/Project.toml @@ -3,11 +3,12 @@ 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.10.0" +version = "5.11.0" [deps] BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66" ConsoleProgressMonitor = "88cd18e8-d9cc-4ea6-8889-5259c0d15c8b" +Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c" @@ -20,24 +21,36 @@ TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed" Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999" UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" +[weakdeps] +OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e" +TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f" + +[extensions] +AbstractMCMCOnlineStatsExt = "OnlineStats" +AbstractMCMCTensorBoardLoggerExt = "TensorBoardLogger" + [compat] BangBang = "0.3.19, 0.4" ConsoleProgressMonitor = "0.1" FillArrays = "1" LogDensityProblems = "2" LoggingExtras = "0.4, 0.5, 1" +OnlineStats = "1.5" ProgressLogging = "0.1" StatsBase = "0.32, 0.33, 0.34" +TensorBoardLogger = "0.1" TerminalLoggers = "0.1" Transducers = "0.4.30" UUIDs = "<0.0.1, 1" -julia = "1.6" +julia = "1.10" [extras] FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" IJulia = "7073ff75-c697-5162-941a-fcdaad2a7d2a" +OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +TensorBoardLogger = "899adc3e-224a-11e9-021f-63837185c80f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["FillArrays", "IJulia", "Statistics", "Test"] +test = ["FillArrays", "IJulia", "OnlineStats", "Statistics", "TensorBoardLogger", "Test"] diff --git a/docs/Project.toml b/docs/Project.toml index f74dfb58..d5fc343e 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,4 +1,5 @@ [deps] +AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" diff --git a/docs/make.jl b/docs/make.jl index 1e45cbff..30398820 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -8,6 +8,6 @@ makedocs(; sitename="AbstractMCMC", format=Documenter.HTML(), modules=[AbstractMCMC], - pages=["Home" => "index.md", "api.md", "design.md"], + pages=["Home" => "index.md", "api.md", "callbacks.md", "design.md"], checkdocs=:exports, ) diff --git a/docs/src/assets/tensorboard_demo_distributions_screen.png b/docs/src/assets/tensorboard_demo_distributions_screen.png new file mode 100644 index 00000000..c78f6ef7 Binary files /dev/null and b/docs/src/assets/tensorboard_demo_distributions_screen.png differ diff --git a/docs/src/assets/tensorboard_demo_histograms_screen.png b/docs/src/assets/tensorboard_demo_histograms_screen.png new file mode 100644 index 00000000..fccb4106 Binary files /dev/null and b/docs/src/assets/tensorboard_demo_histograms_screen.png differ diff --git a/docs/src/assets/tensorboard_demo_scalars_screen.png b/docs/src/assets/tensorboard_demo_scalars_screen.png new file mode 100644 index 00000000..a765b391 Binary files /dev/null and b/docs/src/assets/tensorboard_demo_scalars_screen.png differ diff --git a/docs/src/assets/tensorboard_demo_time-series_screen.png b/docs/src/assets/tensorboard_demo_time-series_screen.png new file mode 100644 index 00000000..1830d53a Binary files /dev/null and b/docs/src/assets/tensorboard_demo_time-series_screen.png differ diff --git a/docs/src/callbacks.md b/docs/src/callbacks.md new file mode 100644 index 00000000..ec3c0c61 --- /dev/null +++ b/docs/src/callbacks.md @@ -0,0 +1,243 @@ +# Callbacks + +AbstractMCMC provides a unified callback API for monitoring and logging MCMC sampling. + +## Basic Usage + +The `mcmc_callback` function is the main entry point for creating callbacks: + +```julia +using AbstractMCMC + +# Simple callback with a function +cb = mcmc_callback() do rng, model, sampler, transition, state, iteration + println("Iteration: $iteration") +end + +chain = sample(model, sampler, 1000; callback=cb) +``` + +## Combining Multiple Callbacks + +Pass multiple callbacks to `mcmc_callback` to combine them: + +```julia +cb1 = (args...; kwargs...) -> println("Callback 1") +cb2 = (args...; kwargs...) -> println("Callback 2") + +cb = mcmc_callback(cb1, cb2) +``` + +You can also add callbacks dynamically using `BangBang.push!!`: + +```julia +using BangBang + +cb = mcmc_callback(cb1) +cb = push!!(cb, cb2) +``` + +## TensorBoard Logging + +TensorBoard logging requires `TensorBoardLogger`. Statistics collection also requires `OnlineStats`. + +### Basic Logging (No Statistics) + +```julia +using AbstractMCMC +using TensorBoardLogger + +logger = TBLogger("runs/experiment1") +cb = mcmc_callback(logger=logger) + +chain = sample(model, sampler, 1000; callback=cb) +``` + +### Logging with Statistics + +To collect running statistics (mean, variance, histograms), load `OnlineStats` and use the `stats` argument: + +```julia +using AbstractMCMC +using TensorBoardLogger +using OnlineStats + +logger = TBLogger("runs/experiment1") + +# Use default statistics (Mean, Variance, KHist) +cb = mcmc_callback(logger=logger, stats=true) + +# Or specify custom statistics +cb = mcmc_callback( + logger=logger, + stats=(Mean(), Variance(), KHist(50)), +) +``` + +!!! note + If you request statistics without loading `OnlineStats`, you will get a helpful error: + `"Statistics collection requires OnlineStats.jl. Please load OnlineStats before enabling statistics."` + +### Stats Processing Options + +Control how samples are processed before computing statistics with `stats_options`: + +```julia +cb = mcmc_callback( + logger=logger, + stats=true, + stats_options=( + skip=100, # Skip first 100 samples (burn-in) + thin=5, # Use every 5th sample + window=1000, # Rolling window of 1000 samples + ), +) +``` + +Options merge with defaults, so you only need to specify what you want to change: + +```julia +# Only change thin, skip and window use defaults (0 and typemax(Int)) +cb = mcmc_callback(logger=logger, stats=true, stats_options=(thin=10,)) +``` + +### Name Filtering + +Use `name_filter` to control which parameters and statistics are logged: + +```julia +cb = mcmc_callback( + logger=logger, + 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) + ), +) +``` + +### Complete Example + +```julia +using AbstractMCMC +using TensorBoardLogger +using OnlineStats + +logger = TBLogger("runs/full_example") + +cb = mcmc_callback( + logger=logger, + stats=true, + stats_options=(skip=50, thin=2), + name_filter=( + exclude=["_internal"], + extras=true, + hyperparams=true, + ), +) + +chain = sample(model, sampler, 10000; callback=cb) +``` + +Then view in TensorBoard: +```bash +tensorboard --logdir=runs/full_example +``` + +Navigate to `localhost:6006` in your browser to see the dashboard. You'll see real-time plots of your parameter distributions, histograms, and other statistics as sampling progresses. + +![TensorBoard Time Series Tab](assets/tensorboard_demo_time-series_screen.png) + +*The Time Series tab provides detailed traces of parameter values throughout the sampling process.* + +![TensorBoard Scalars Tab](assets/tensorboard_demo_scalars_screen.png) + +*The Scalars tab shows time series of parameter values and statistics over the sampling iterations.* + +![TensorBoard Distributions Tab](assets/tensorboard_demo_distributions_screen.png) + +*The Distributions tab displays the marginal distributions of each parameter.* + +![TensorBoard Histograms Tab](assets/tensorboard_demo_histograms_screen.png) + +*The Histograms tab shows the evolution of parameter distributions over time.* + +## API Reference + +### Main Functions + +```@docs +mcmc_callback +``` + +## Default Values + +### stats_options defaults + +| Option | Default | Description | +|----------|----------------|--------------------------------| +| `skip` | `0` | Skip first n samples (burn-in) | +| `thin` | `0` | Use every nth sample (0=all) | +| `window` | `typemax(Int)` | Window size for rolling stats | + +### name_filter defaults + +| Option | Default | Description | +|--------------|------------|----------------------------------| +| `include` | `String[]` | Only log these (empty=all) | +| `exclude` | `String[]` | Don't log these | +| `extras` | `false` | Include extra stats | +| `hyperparams`| `false` | Include hyperparameters | + +## Implementing Custom Callbacks + +Any callable with the following signature can be used as a callback: + +```julia +function my_callback(rng, model, sampler, transition, state, iteration; kwargs...) + # Your callback logic here +end +``` + +## Internals + +!!! note + These types and methods are used internally. They are not part of the public API and + may change or break at any time without notice. + +### Types + +```@docs +AbstractMCMC.MultiCallback +AbstractMCMC.NameFilter +``` + +### OnlineStats Wrappers + +When using statistics, AbstractMCMC provides wrappers that modify how samples are processed: + +| Wrapper | Description | +|---------|-------------| +| `Skip(n, stat)` | Skip first `n` observations before fitting `stat` | +| `Thin(n, stat)` | Only fit every `n`-th observation to `stat` | +| `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 new file mode 100644 index 00000000..c9020df1 --- /dev/null +++ b/ext/AbstractMCMCOnlineStatsExt.jl @@ -0,0 +1,174 @@ +module AbstractMCMCOnlineStatsExt + +using AbstractMCMC +using OnlineStats +using OnlineStats: + OnlineStat, fit!, value, nobs, MovingWindow, Series, Mean, Variance, KHist +using Logging: @info + +""" + Skip(b::Int, stat::OnlineStat) + +Skips the first `b` observations before passing them on to `stat`. +""" +mutable struct Skip{T,O<:OnlineStat{T}} <: OnlineStat{T} + b::Int + n::Int + stat::O +end + +Skip(b::Int, stat::OnlineStat{T}) where {T} = Skip{T,typeof(stat)}(b, 0, stat) + +OnlineStats.nobs(o::Skip) = OnlineStats.nobs(o.stat) +OnlineStats.value(o::Skip) = OnlineStats.value(o.stat) + +function OnlineStats._fit!(o::Skip, x) + if o.n >= o.b + OnlineStats._fit!(o.stat, x) + end + o.n += 1 + return o +end + +Base.show(io::IO, o::Skip) = print(io, "Skip($(o.b)): n=$(o.n) | $(o.stat)") + +""" + Thin(b::Int, stat::OnlineStat) + +Thins `stat` with an interval `b`, i.e. only passes every b-th observation to `stat`. +""" +mutable struct Thin{T,O<:OnlineStat{T}} <: OnlineStat{T} + b::Int + n::Int + stat::O +end + +Thin(b::Int, stat::OnlineStat{T}) where {T} = Thin{T,typeof(stat)}(b, 0, stat) + +OnlineStats.nobs(o::Thin) = OnlineStats.nobs(o.stat) +OnlineStats.value(o::Thin) = OnlineStats.value(o.stat) + +function OnlineStats._fit!(o::Thin, x) + if (o.n % o.b) == 0 + OnlineStats._fit!(o.stat, x) + end + o.n += 1 + return o +end + +Base.show(io::IO, o::Thin) = print(io, "Thin($(o.b)): n=$(o.n) | $(o.stat)") + +""" + WindowStat(b::Int, stat::OnlineStat) + +Wraps `stat` in a `MovingWindow` of length `b`. +""" +struct WindowStat{T,O} <: OnlineStat{T} + window::MovingWindow{T} + stat::O +end + +function WindowStat(b::Int, stat::OnlineStat{T}) where {T} + return WindowStat{T,typeof(stat)}(MovingWindow(b, T), stat) +end + +OnlineStats.nobs(o::WindowStat) = OnlineStats.nobs(o.window) + +function OnlineStats._fit!(o::WindowStat, x) + OnlineStats._fit!(o.window, x) + return o +end + +function OnlineStats.value(o::WindowStat) + stat_new = deepcopy(o.stat) + fit!(stat_new, OnlineStats.value(o.window)) + return stat_new +end + +function Base.show(io::IO, o::WindowStat) + return print(io, "WindowStat($(o.window.b)): nobs=$(nobs(o)) | $(o.stat)") +end + +""" + create_stats_with_options_impl(stats, stats_options, num_bins) + +Create stats dictionary and prototype, applying Skip/Thin/WindowStat wrappers. +""" +function create_stats_with_options_impl(stats, stats_options, num_bins) + base_stat = if stats === true || stats === :default + Series(Mean(), Variance(), KHist(num_bins)) + elseif stats isa OnlineStat + stats + elseif stats isa Tuple + Series(stats...) + else + Series(Mean(), Variance(), KHist(num_bins)) + end + + wrapped_stat = wrap_stat(base_stat, stats_options) + + if wrapped_stat isa OnlineStat + nobs(wrapped_stat) > 0 && + @warn("using statistic with observations as a base: $(wrapped_stat)") + return (Dict{String,typeof(wrapped_stat)}(), deepcopy(wrapped_stat)) + else + return (wrapped_stat, nothing) + end +end + +""" + wrap_stat(stat, stats_options) + +Apply Skip, Thin, and WindowStat wrappers to a statistic based on options. +""" +function wrap_stat(stat, stats_options) + result = stat + + if stats_options.window < typemax(Int) + result = WindowStat(stats_options.window, result) + end + + if stats_options.skip > 0 + result = Skip(stats_options.skip, result) + end + + if stats_options.thin > 0 + result = Thin(stats_options.thin, result) + end + + return result +end + +""" + log_stat_impl!(stats, prototype, key, val, prefix) + +Update and log statistics. Called from TensorBoard callback. +""" +function log_stat_impl!(stats::AbstractDict, prototype, key, val, prefix) + stat = if prototype !== nothing + get!(stats, key) do + deepcopy(prototype) + end + else + get(stats, key, nothing) + end + + if stat !== nothing + fit!(stat, val) + @info "$(prefix)$key" stat + end +end + +log_stat_impl!(::Nothing, prototype, key, val, prefix) = nothing + +# tb_name helpers for formatting stat names in TensorBoard +tb_name(arg) = string(arg) +tb_name(stat::OnlineStat) = string(nameof(typeof(stat))) +tb_name(o::Skip) = "Skip($(o.b))" +tb_name(o::Thin) = "Thin($(o.b))" +tb_name(o::WindowStat) = "WindowStat($(o.window.b))" +tb_name(s1::String, s2::String) = s1 * "/" * s2 +tb_name(arg1, arg2) = tb_name(arg1) * "/" * tb_name(arg2) +tb_name(arg, args...) = tb_name(arg) * "/" * tb_name(args...) + +end diff --git a/ext/AbstractMCMCTensorBoardLoggerExt.jl b/ext/AbstractMCMCTensorBoardLoggerExt.jl new file mode 100644 index 00000000..ab60e6d2 --- /dev/null +++ b/ext/AbstractMCMCTensorBoardLoggerExt.jl @@ -0,0 +1,158 @@ +module AbstractMCMCTensorBoardLoggerExt + +using AbstractMCMC +using AbstractMCMC: + MultiCallback, + NameFilter, + _names_and_values, + hyperparam_metrics, + merge_with_defaults, + create_stats_with_options, + DEFAULT_STATS_OPTIONS, + DEFAULT_NAME_FILTER +using TensorBoardLogger +using TensorBoardLogger: TBLogger +using Logging: AbstractLogger, with_logger, @info + +""" + TensorBoardCallback + +A callback for logging MCMC samples to TensorBoard. +Supports statistics collection when OnlineStats.jl is loaded. +""" +struct TensorBoardCallback{L,S,P,F} + logger::L + stats::S + stat_prototype::P + variable_filter::F + include_extras::Bool + include_hyperparams::Bool + param_prefix::String + extras_prefix::String +end + +""" + mcmc_callback(; logger, stats=nothing, stats_options=nothing, name_filter=nothing, num_bins=100) + +Create a TensorBoard logging callback. + +# Arguments +- `logger`: An `AbstractLogger` instance (e.g., `TBLogger` from TensorBoardLogger.jl) +- `stats`: Statistics to collect. Can be: + - `nothing`: No statistics (default) + - `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` +- `num_bins`: Number of histogram bins (default: 100) + +# Examples +```julia +using TensorBoardLogger +lg = TBLogger("runs/exp") +cb = mcmc_callback(logger=lg) + +# With default stats (requires OnlineStats) +using TensorBoardLogger, OnlineStats +lg = TBLogger("runs/exp") +cb = mcmc_callback(logger=lg, stats=true) +``` +""" +function AbstractMCMC.mcmc_callback(; + logger::AbstractLogger, + stats=nothing, + stats_options=nothing, + name_filter=nothing, + num_bins::Int=100, +) + merged_stats_options = merge_with_defaults(stats_options, DEFAULT_STATS_OPTIONS) + merged_name_filter = merge_with_defaults(name_filter, DEFAULT_NAME_FILTER) + + processed_stats = create_stats_with_options(stats, merged_stats_options, num_bins) + + variable_filter = NameFilter(; + include=merged_name_filter.include, exclude=merged_name_filter.exclude + ) + + stats_dict, prototype = if processed_stats === nothing + (nothing, nothing) + else + processed_stats + end + + callback = TensorBoardCallback( + logger, + stats_dict, + prototype, + variable_filter, + merged_name_filter.extras, + merged_name_filter.hyperparams, + "", + "extras/", + ) + + return MultiCallback((callback,)) +end + +function filter_name_and_value(cb::TensorBoardCallback, name_and_value) + return cb.variable_filter(first(name_and_value), last(name_and_value)) +end + +function (cb::TensorBoardCallback)( + rng, model, sampler, transition, state, iteration; kwargs... +) + stats = cb.stats + 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( + model, + sampler, + transition, + state; + params=true, + hyperparams=false, + extra=cb.include_extras, + kwargs..., + ) + + for (k, val) in Iterators.filter(filter_fn, all_values) + @info "$(cb.param_prefix)$k" val + + if stats !== nothing + _log_stat!(stats, cb.stat_prototype, k, val, cb.param_prefix) + end + end + + TensorBoardLogger.increment_step!(lg, 1) + end +end + +function _log_stat!(stats, prototype, key, val, prefix) + ext = Base.get_extension(AbstractMCMC, :AbstractMCMCOnlineStatsExt) + if ext !== nothing + ext.log_stat_impl!(stats, prototype, key, val, prefix) + end +end + +end diff --git a/src/AbstractMCMC.jl b/src/AbstractMCMC.jl index bceb04fd..834cdf3d 100644 --- a/src/AbstractMCMC.jl +++ b/src/AbstractMCMC.jl @@ -22,6 +22,9 @@ export sample # Parallel sampling types export MCMCThreads, MCMCDistributed, MCMCSerial +# Callback API +export mcmc_callback + """ AbstractChains @@ -189,6 +192,7 @@ include("sample.jl") include("stepper.jl") include("transducer.jl") include("logdensityproblems.jl") +include("callbacks.jl") if isdefined(Base.Experimental, :register_error_hint) function __init__() @@ -204,6 +208,18 @@ if isdefined(Base.Experimental, :register_error_hint) ) end end + + Base.Experimental.register_error_hint(MethodError) do io, exc, argtypes, _ + if exc.f === mcmc_callback && length(argtypes) == 0 + printstyled( + io, + "\n\n`mcmc_callback(logger=...)` requires TensorBoardLogger.jl to be loaded.\n" * + "Please run `using TensorBoardLogger` before creating a logger callback.\n"; + color=:cyan, + bold=true, + ) + end + end end end diff --git a/src/callbacks.jl b/src/callbacks.jl new file mode 100644 index 00000000..f7138385 --- /dev/null +++ b/src/callbacks.jl @@ -0,0 +1,272 @@ +# Callbacks for AbstractMCMC +# This module provides the unified callback API and supporting types. + +########################### +### Core Callback Types ### +########################### + +""" + MultiCallback + +A callback that combines multiple callbacks into one. + +Supports `push!!` from [BangBang.jl](https://github.com/JuliaFolds/BangBang.jl) to add callbacks, +returning a new `MultiCallback` with the added callback. +""" +struct MultiCallback{Cs<:Tuple} + callbacks::Cs +end + +MultiCallback() = MultiCallback(()) +MultiCallback(callbacks...) = MultiCallback(callbacks) + +(c::MultiCallback)(args...; kwargs...) = foreach(c -> c(args...; kwargs...), c.callbacks) + +function BangBang.push!!(c::MultiCallback, callback) + return MultiCallback((c.callbacks..., callback)) +end + +""" + NameFilter(; include=Set{String}(), exclude=Set{String}()) + +A filter for variable names. + +- If `include` is non-empty, only names in `include` will pass the filter. +- Names in `exclude` will be excluded. +- Throws an error if `include` and `exclude` have overlapping elements. +""" +struct NameFilter + include::Set{String} + exclude::Set{String} + + function NameFilter(; include=Set{String}(), exclude=Set{String}()) + inc_set = include isa Set ? include : Set{String}(include) + exc_set = exclude isa Set ? exclude : Set{String}(exclude) + overlap = intersect(inc_set, exc_set) + if !isempty(overlap) + error("NameFilter: include and exclude have overlapping elements: $overlap") + end + return new(inc_set, exc_set) + end +end + +(f::NameFilter)(name, value) = f(name) +function (f::NameFilter)(name) + return name ∉ f.exclude && (isempty(f.include) || name ∈ f.include) +end + +############################## +### Defaults and Utilities ### +############################## + +const DEFAULT_STATS_OPTIONS = (; thin=0, skip=0, window=typemax(Int)) + +const DEFAULT_NAME_FILTER = (; + include=String[], exclude=String[], extras=false, hyperparams=false +) + +""" + merge_with_defaults(user_options::NamedTuple, defaults::NamedTuple) + +Merge user-provided options with defaults, where user options take precedence. +""" +function merge_with_defaults(user_options::NamedTuple, defaults::NamedTuple) + return merge(defaults, user_options) +end +merge_with_defaults(::Nothing, defaults::NamedTuple) = defaults + +################################ +### Statistics Creation API ### +################################ + +""" + create_stats_with_options(stats, stats_options, num_bins) + +Internal constructor for statistics handlers. + +If `stats === nothing`, no statistics are collected and `nothing` is returned. +If `stats` is provided, this function requires the OnlineStats extension to be +loaded; otherwise, an error is thrown. + +Supports special values: +- `stats=true` or `stats=:default`: Use default statistics (Mean, Variance, KHist) +- `stats=`: Use the provided OnlineStat (requires OnlineStats to be loaded) +- `stats=`: Use multiple stats + +This function is not part of the public API and may change or break at any time. +""" +create_stats_with_options(::Nothing, stats_options, num_bins) = nothing + +function create_stats_with_options(stats, stats_options, num_bins) + ext = Base.get_extension(@__MODULE__, :AbstractMCMCOnlineStatsExt) + if ext === nothing + error( + "Statistics collection requires OnlineStats.jl. " * + "Please load OnlineStats before enabling statistics: `using OnlineStats`", + ) + end + + # Delegate to OnlineStatsExt for actual creation + return ext.create_stats_with_options_impl(stats, stats_options, num_bins) +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)) + +""" + _names_and_values( + model, + sampler, + transition, + state; + params::Bool = true, + hyperparams::Bool = false, + extra::Bool = false, + kwargs... + ) + +Return an iterator over parameter names and values. + +This function is not part of the public API and may change or break at any time. + +## Keywords +- `params`: include model parameters. +- `hyperparams`: include sampler hyperparameters +- `extra`: include additional statistics. +- `kwargs...`: reserved for internal extensibility. +""" +function _names_and_values( + model, + sampler, + transition, + state; + params::Bool=true, + hyperparams::Bool=false, + extra::Bool=false, + kwargs..., +) + 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 + + if extra + try + stats = getstats(state) + if stats isa NamedTuple + push!(iters, pairs(stats)) + end + catch + # No extras 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 + +""" + 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...) +end + +################################# +### Unified mcmc_callback API ### +################################# + +""" + mcmc_callback(callback) + mcmc_callback(callbacks...) + +Create a callback or combine multiple callbacks into one. + +Any callable (function or callable struct) with the signature +`(rng, model, sampler, transition, state, iteration; kwargs...)` can be used. + +# Example +```julia +cb = mcmc_callback() do rng, model, sampler, transition, state, iteration + println("Iteration: \$iteration") +end +``` +""" +function mcmc_callback(cb1, callbacks...) + return MultiCallback((cb1, callbacks...)) +end + +""" + mcmc_callback(; + logger, + stats = nothing, + stats_options = nothing, + name_filter = nothing, + ) + +Create a TensorBoard logging callback. **Requires TensorBoardLogger.jl to be loaded.** + +# Arguments +- `logger`: An `AbstractLogger` instance (e.g., `TBLogger` from TensorBoardLogger.jl) +- `stats`: Statistics to collect. Can be: + - `nothing`: No statistics (default) + - `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` + +# Examples +```julia +using TensorBoardLogger +lg = TBLogger("runs/exp") +cb = mcmc_callback(logger=lg) + +# With default stats (requires OnlineStats) +using TensorBoardLogger, OnlineStats +lg = TBLogger("runs/exp") +cb = mcmc_callback(logger=lg, stats=true) +``` + +!!! note + This method is defined in the TensorBoardLogger extension. You must load + TensorBoardLogger before using it: `using TensorBoardLogger` +""" +function mcmc_callback end + +""" + mcmc_callback(existing::MultiCallback, new_callbacks...) + +Add callbacks to an existing MultiCallback. +""" +function mcmc_callback(existing::MultiCallback, new_callbacks...) + return MultiCallback((existing.callbacks..., new_callbacks...)) +end diff --git a/test/callbacks.jl b/test/callbacks.jl new file mode 100644 index 00000000..df791fa8 --- /dev/null +++ b/test/callbacks.jl @@ -0,0 +1,332 @@ +#################################### +### Basic Callback Functionality ### +#################################### + +@testset "Basic mcmc_callback" begin + @testset "Function callback" begin + count = Ref(0) + cb = mcmc_callback() do rng, model, sampler, transition, state, iteration + count[] += 1 + end + + @test cb isa AbstractMCMC.MultiCallback + chain = sample(MyModel(), MySampler(), 100; callback=cb) + @test count[] == 100 + end + + @testset "Multiple function callbacks" begin + counts = [Ref(0), Ref(0)] + cb1 = (args...; kwargs...) -> counts[1][] += 1 + cb2 = (args...; kwargs...) -> counts[2][] += 1 + + cb = mcmc_callback(cb1, cb2) + @test cb isa AbstractMCMC.MultiCallback + + chain = sample(MyModel(), MySampler(), 50; callback=cb) + @test counts[1][] == 50 + @test counts[2][] == 50 + end + + @testset "Error without callback type" begin + @test_throws MethodError mcmc_callback() + end + + @testset "Adding callbacks with push!!" begin + counts = [Ref(0), Ref(0)] + cb1 = (args...; kwargs...) -> counts[1][] += 1 + cb2 = (args...; kwargs...) -> counts[2][] += 1 + + cb = mcmc_callback(cb1) + cb = BangBang.push!!(cb, cb2) + + for _ in 1:5 + cb(nothing, nothing, nothing, nothing, nothing, 1) + end + @test counts[1][] == 5 + @test counts[2][] == 5 + end + + @testset "Callable struct callback" begin + struct CountingCallback + count::Ref{Int} + end + function (cb::CountingCallback)(args...; kwargs...) + return cb.count[] += 1 + end + + counter = CountingCallback(Ref(0)) + cb = mcmc_callback(counter) + + @test cb isa AbstractMCMC.MultiCallback + chain = sample(MyModel(), MySampler(), 25; callback=cb) + @test counter.count[] == 25 + end +end + +######################## +### Defaults Merging ### +######################## + +@testset "Defaults merging" begin + @testset "merge_with_defaults" begin + defaults = (; a=1, b=2, c=3) + + # Partial override + result = AbstractMCMC.merge_with_defaults((; b=10), defaults) + @test result == (; a=1, b=10, c=3) + + # Full override + result = AbstractMCMC.merge_with_defaults((; a=10, b=20, c=30), defaults) + @test result == (; a=10, b=20, c=30) + + # No override (nothing) + result = AbstractMCMC.merge_with_defaults(nothing, defaults) + @test result == defaults + end + + @testset "DEFAULT_STATS_OPTIONS" begin + @test AbstractMCMC.DEFAULT_STATS_OPTIONS.thin == 0 + @test AbstractMCMC.DEFAULT_STATS_OPTIONS.skip == 0 + @test AbstractMCMC.DEFAULT_STATS_OPTIONS.window == typemax(Int) + 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.hyperparams == false + end +end + +###################### +### Internal Types ### +###################### + +@testset "MultiCallback" begin + counts = [Ref(0), Ref(0)] + cb1 = (args...; kwargs...) -> counts[1][] += 1 + cb2 = (args...; kwargs...) -> counts[2][] += 1 + + multi = AbstractMCMC.MultiCallback(cb1, cb2) + multi(nothing, nothing, nothing, nothing, nothing, 1) + + @test counts[1][] == 1 + @test counts[2][] == 1 +end + +@testset "NameFilter" begin + @testset "Include only" begin + f = AbstractMCMC.NameFilter(; include=["a", "b"]) + @test f("a") == true + @test f("b") == true + @test f("c") == false + end + + @testset "Exclude only" begin + f = AbstractMCMC.NameFilter(; exclude=["x", "y"]) + @test f("a") == true + @test f("x") == false + @test f("y") == false + end + + @testset "Include and exclude (non-overlapping)" begin + f = AbstractMCMC.NameFilter(; include=["a", "b"], exclude=["x", "y"]) + @test f("a") == true + @test f("b") == true + @test f("x") == false + @test f("y") == false + @test f("c") == false + end + + @testset "Include and exclude (overlapping errors)" begin + @test_throws ErrorException AbstractMCMC.NameFilter(; + include=["a", "b", "c"], exclude=["c"] + ) + end + + @testset "No filter" begin + f = AbstractMCMC.NameFilter() + @test f("anything") == true + end + + @testset "Two argument form" begin + f = AbstractMCMC.NameFilter(; include=["a", "b"]) + @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]"] +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 + +############################# +### TensorBoard Extension ### +############################# + +using TensorBoardLogger + +@testset "TensorBoard Extension" begin + @testset "mcmc_callback with explicit TBLogger" begin + logdir = mktempdir() + logger = TBLogger(logdir) + cb = mcmc_callback(; logger=logger) + @test cb isa AbstractMCMC.MultiCallback + end + + @testset "mcmc_callback requires logger argument" begin + @test_throws UndefKeywordError mcmc_callback() + end + + @testset "mcmc_callback with stats=true" begin + logdir = mktempdir() + logger = TBLogger(logdir) + cb = mcmc_callback(; logger=logger, stats=true) + @test cb isa AbstractMCMC.MultiCallback + end + + @testset "mcmc_callback with stats=:default" begin + logdir = mktempdir() + logger = TBLogger(logdir) + cb = mcmc_callback(; logger=logger, stats=:default) + @test cb isa AbstractMCMC.MultiCallback + end + + @testset "mcmc_callback with explicit OnlineStats" begin + logdir = mktempdir() + logger = TBLogger(logdir) + cb = mcmc_callback(; logger=logger, stats=(Mean(), Variance())) + @test cb isa AbstractMCMC.MultiCallback + end + + @testset "mcmc_callback with stats_options" begin + logdir = mktempdir() + logger = TBLogger(logdir) + + # Test partial stats_options (merges with defaults) + cb = mcmc_callback(; logger=logger, stats=true, stats_options=(thin=5,)) + @test cb isa AbstractMCMC.MultiCallback + + # Test full stats_options + cb = mcmc_callback(; + logger=logger, stats=true, stats_options=(thin=5, skip=100, window=1000) + ) + @test cb isa AbstractMCMC.MultiCallback + end + + @testset "mcmc_callback with name_filter" begin + logdir = mktempdir() + logger = TBLogger(logdir) + + # Test partial name_filter + cb = mcmc_callback(; logger=logger, name_filter=(include=["mu", "sigma"],)) + @test cb isa AbstractMCMC.MultiCallback + + # Test full name_filter + cb = mcmc_callback(; + logger=logger, + name_filter=( + include=["mu", "sigma"], exclude=["internal"], extras=true, hyperparams=true + ), + ) + @test cb isa AbstractMCMC.MultiCallback + end + + @testset "mcmc_callback with all options" begin + logdir = mktempdir() + logger = TBLogger(logdir) + cb = mcmc_callback(; + logger=logger, + stats=true, + stats_options=(skip=100, thin=5), + name_filter=(exclude=["_internal"], extras=true), + ) + @test cb isa AbstractMCMC.MultiCallback + end + + @testset "TensorBoard callback with sample" begin + logdir = mktempdir() + logger = TBLogger(logdir) + cb = mcmc_callback(; logger=logger) + + # Should complete without error + chain = sample(MyModel(), MySampler(), 20; callback=cb) + @test length(chain) == 20 + end + + @testset "mcmc_callback with custom AbstractLogger" begin + # Use TBLogger as our custom logger (could be any AbstractLogger) + logdir = mktempdir() + custom_logger = TensorBoardLogger.TBLogger(logdir; min_level=Logging.Info) + + cb = mcmc_callback(; logger=custom_logger) + @test cb isa AbstractMCMC.MultiCallback + + # Should work with sampling + chain = sample(MyModel(), MySampler(), 20; callback=cb) + @test length(chain) == 20 + end + + @testset "Stats with stats=true works when OnlineStats loaded" begin + logdir = mktempdir() + logger = TBLogger(logdir) + # OnlineStats is already loaded, so this should work + cb = mcmc_callback(; logger=logger, stats=true) + @test cb isa AbstractMCMC.MultiCallback + + # Also works with explicit OnlineStat + cb = mcmc_callback(; logger=logger, stats=Mean()) + @test cb isa AbstractMCMC.MultiCallback + end +end + +######################### +### Integration Tests ### +######################### + +@testset "Integration" begin + @testset "Callback receives correct iteration" begin + iterations = Int[] + cb = mcmc_callback() do rng, model, sampler, transition, state, iteration + push!(iterations, iteration) + end + + chain = sample(MyModel(), MySampler(), 50; callback=cb) + @test iterations == 1:50 + end + + @testset "Multiple callbacks all execute" begin + results = Dict{Symbol,Int}(:cb1 => 0, :cb2 => 0, :cb3 => 0) + + cb = mcmc_callback( + (args...; kwargs...) -> results[:cb1] += 1, + (args...; kwargs...) -> results[:cb2] += 1, + (args...; kwargs...) -> results[:cb3] += 1, + ) + + chain = sample(MyModel(), MySampler(), 30; callback=cb) + @test results[:cb1] == 30 + @test results[:cb2] == 30 + @test results[:cb3] == 30 + end + + @testset "Combining TensorBoard with custom callback" begin + logdir = mktempdir() + logger = TBLogger(logdir) + count = Ref(0) + custom = (args...; kwargs...) -> count[] += 1 + + tb_cb = mcmc_callback(; logger=logger) + combined = mcmc_callback(tb_cb, custom) + + @test length(combined.callbacks) == 2 + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 909ae8b3..c7c793a8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,5 @@ using AbstractMCMC +using BangBang using ConsoleProgressMonitor: ProgressLogger using IJulia using LogDensityProblems @@ -25,3 +26,7 @@ include("utils.jl") include("transducer.jl") include("logdensityproblems.jl") end + +@testset "Callbacks" begin + include("callbacks.jl") +end