Port TuringCallbacks to AbstractMCMC

.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

.gitignore

@@ -2,4 +2,5 @@
 *.jl.*.cov
 *.jl.mem
 deps/deps.jl
-Manifest.toml
+Manifest.toml
+tmp/

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"]

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"
 

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,
 )

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.