Fitting

Running NUTS over a composed model, initialising from a viable prior draw, and the callbacks that stream progress to a log and to TensorBoard. fit_parallel runs independent fits side by side, bounded by the available threads.

Index

• BVDOutbreakSize.ViablePrior
• BVDOutbreakSize.combined_callback
• BVDOutbreakSize.default_adtype
• BVDOutbreakSize.enzyme_adtype
• BVDOutbreakSize.fit_callback
• BVDOutbreakSize.fit_parallel
• BVDOutbreakSize.nuts_sample
• BVDOutbreakSize.progress_callback
• BVDOutbreakSize.tensorboard_callback
• BVDOutbreakSize.viable_prior_init

Reference

BVDOutbreakSize.ViablePrior Type

struct ViablePrior

struct ViablePrior

Prior initialisation guarded against the prior predictive's unrecoverable tail. Each chain screens attempts independent prior draws and starts at the first whose initial log joint density is at or above that batch's median. The default init of nuts_sample.

InitFromPrior draws each parameter from its own prior independently. The joint model's product prior is dispersed enough that a sizeable minority of those draws put the whole latent trajectory hundreds of thousands of log units below the posterior. NUTS does not recover from such a point. Dual averaging shrinks the step size towards zero instead of moving and the chain crawls in place for the entire run. Nothing diverges, so the failure surfaces only as a split R-hat pinned near its ceiling.

The median rather than the batch maximum, because an argmax is an order statistic of the joint density rather than a draw from the prior. It keeps roughly the top eighth of the prior by density, concentrating the starting points and reducing the between-chain contrast R-hat is built on for a reason unrelated to mixing. Both rules clear the tail, and the median rule keeps more than twice the dispersion of starting log joints. Only forward density evaluations are used, so the guard costs milliseconds against a fit measured in hours.

• attempts::Int64: Prior draws screened per chain; the first above their median is used.

---

Fields

• attempts::Int64: Prior draws screened per chain; the first above their median is used.

source

BVDOutbreakSize.combined_callback Method

combined_callback(callbacks...) -> Any

Compose several nuts_sample step callbacks into one. Each argument is either a callback with the AbstractMCMC step signature or nothing. nothing entries are dropped. The composite invokes the surviving callbacks in order on every step. Returns the single callback unchanged when only one survives, and nothing when none do (so nuts_sample sees no callback at all rather than a no-op wrapper).

See also: fit_callback, progress_callback, tensorboard_callback.

source

BVDOutbreakSize.default_adtype Method

default_adtype() -> ADTypes.AutoMooncake{Mooncake.Config}

Mooncake reverse-mode AD with default Mooncake.Config(). Used as the NUTS adtype keyword.

source

BVDOutbreakSize.enzyme_adtype Function

enzyme_adtype()

Enzyme reverse-mode AD type, an opt-in alternative to the default default_adtype (Mooncake). Defined by the package's Enzyme weak-dependency extension (ext/BVDOutbreakSizeEnzymeExt.jl), so calling it without Enzyme loaded raises a MethodError. The SpecialFunctions.gamma EnzymeRule that the Beta and NegativeBinomial normalising constants reach comes from CensoredDistributions' own Enzyme extension. Enzyme differentiates the single-stream composers and matches Mooncake. Differentiating the full joint is platform-dependent and can hit an upstream Enzyme/LLVM compile failure, so Mooncake remains the package default for fitting.

source

BVDOutbreakSize.fit_callback Method

fit_callback(
    name::AbstractString;
    logdir,
    spec
) -> Union{Nothing, BVDOutbreakSize.var"#146#147"{BVDOutbreakSize.var"#148#149"{String, Int64, Base.RefValue{Int64}, ReentrantLock}}, BVDOutbreakSize.var"#combined_callback##0#combined_callback##1"}

Build the logging callback for a named model fit, selected by the BVD_FIT_LOG environment variable (or an explicit spec). This is the wiring the report build uses so every fit streams its progress without each call site repeating the callback construction.

Recognised spec values (case-insensitive), defaulting to "all" when BVD_FIT_LOG is unset:

• "all": both the dependency-free progress_callback (a <name>.log file under logdir) and the tensorboard_callback (a tensorboard/<name> run directory under logdir).

• "progress": the file progress stream only.

• "tensorboard" (or "tb"): the TensorBoard stream only.

• "none": no logging, returns nothing. CI sets this to keep release builds quiet (BVD_FIT_LOG=none).

TensorBoard logging needs TensorBoardLogger loaded (it activates the tensorboard_callback method through the package extension). When it is requested but not loaded, the TensorBoard stream is skipped with a warning rather than erroring, so a build without using TensorBoardLogger still gets the file progress stream.

See also: combined_callback, nuts_sample.

source

BVDOutbreakSize.fit_parallel Method

fit_parallel(thunks::AbstractVector; chains) -> Any

fit_parallel(thunks; chains = 2)

Run independent model fits, each a zero-argument thunk returning a chain, with model-level parallelism bounded by the available threads. At most Threads.nthreads() ÷ chains fits run at once, so each fit keeps chains threads for its own chains, clamped to the number of fits. With two threads it runs the fits sequentially, identical to a plain loop and with the same peak memory. On a many-core machine it fans them out. Each fit seeds its own RNG, so the results do not depend on the schedule. Returns the chains in input order.

source

BVDOutbreakSize.nuts_sample Method

nuts_sample(
    model;
    samples,
    chains,
    target_accept,
    max_depth,
    n_adapts,
    seed,
    progress,
    adtype,
    init,
    check_model,
    callback,
    warmup,
    kwargs...
) -> Any

NUTS on model, parallel chains via MCMCThreads. Chains initialise from the prior, each screening eight draws and taking the first at or above their median log joint density (ViablePrior), which keeps the sampler off the prior tail no chain recovers from without concentrating the starts. Pass init = Turing.DynamicPPL.InitFromPrior() for unguarded prior initialisation, or init = Turing.DynamicPPL.InitFromUniform() for unconstrained uniform initialisation. A vector of strategies, one per chain, starts each chain at its own point.

target_accept defaults to 0.85. The renewal joint conditions the confirmed counts on the observed analysed denominator and samples the random-walk and ascertainment blocks in non-centred form, so its geometry is benign. A lower target acceptance shortens the average NUTS trajectory and so the gradient evaluations per iteration, while staying above the conventional 0.8 floor. Raise it towards 0.9-0.99 if a model variant reintroduces divergences. The default is two chains of 500 post-warmup draws, trading some effective sample size for a shorter wall-clock.

check_model = false disables Turing's pre-sampling model check, which rejects any model with a sampled discrete variable even when its value feeds nothing downstream. Only exports_deaths_only_model needs the escape. It runs the exports submodel in predictive mode (exported_cases ~ Poisson with a missing count) purely for the export onsets, leaving a sampled discrete Poisson draw. The continuous parameters are unaffected.

Pass callback to stream live fit progress (iteration, log-density, divergences) instead of waiting for the whole fit. Use progress_callback for a dependency-free file/stdout stream, or tensorboard_callback for a TensorBoard backend (requires using TensorBoardLogger). The callback is forwarded to sample only when non-nothing. Any additional kwargs are passed through to sample.

n_adapts sets the NUTS warmup length (step-size and mass-matrix adaptation), run in addition to samples and discarded by default. It defaults to min(200, samples ÷ 2), capping the per-fit warmup below Turing's own default to speed the report build. The cap binds until a caller drops samples below 400.

A callback fires only on the samples that are kept, and NUTS discards its adaptation phase by default, so warmup is silent. Set warmup = true to keep the adaptation steps (discard_adapt = false), which streams them to the callback so step-size adaptation and early divergences are visible live. Those warmup draws are then also retained in the returned chain, so the first n_adapts draws are adaptation steps rather than posterior samples. Raise samples accordingly or drop them before summarising.

source

BVDOutbreakSize.progress_callback Method

progress_callback(; path, every)

A lightweight, dependency-free streaming progress callback for nuts_sample. Returns a closure matching the AbstractMCMC callback signature

callback(rng, model, sampler, transition, state, iteration; kwargs...)

which, every every iterations on each chain, appends one line to the file at path recording the iteration number, the log joint density, and a running count of divergent transitions.

The same closure is invoked from every thread MCMCThreads() spawns, so the divergence tally and file writes are shared across chains and guarded by a ReentrantLock. Tail the file live during a fit with tail -f <path>.

Step-level statistics are read through the sampler-agnostic AbstractMCMC.ParamsWithStats(model, sampler, transition, state; stats = true) interface rather than by reaching into transition fields, so the callback tracks Turing's transition format instead of a fixed field layout. The log density is taken from the logjoint statistic and the divergence flag from numerical_error. The whole body is wrapped in try/catch, so a transition that does not expose these statistics yields missing log-density / no divergence increment rather than crashing the fit. The running divergence count is a single total over all chains, reset implicitly each time a fresh callback is constructed.

See also: tensorboard_callback, nuts_sample.

source

BVDOutbreakSize.tensorboard_callback Method

tensorboard_callback(logdir; kwargs...)

TensorBoard streaming callback for nuts_sample, mirroring the optional Enzyme backend. tensorboard_callback(logdir; every = 20, histograms = true) opens a TensorBoardLogger.TBLogger(logdir) and, on every kept post-warmup draw, logs through the sampler-agnostic AbstractMCMC.ParamsWithStats interface to two grouped tag prefixes so the dashboard stays navigable:

• params/<name>: every sampled parameter

• diagnostics/<name>: log-density (logjoint), divergence flag (numerical_error), step size, tree depth, acceptance rate, ...

Each scalar streams every step as a .../value time series. With histograms = true (the default) a running histogram of the draws so far is also logged every every steps as .../distribution, populating the TensorBoard histograms and distributions dashboards. Set histograms = false for scalar traces only, or widen every to log histograms less often.

tensorboard_callback is a stub. Loading TensorBoardLogger (using TensorBoardLogger) activates the method via BVDOutbreakSizeTensorBoardLoggerExt. Calling it without TensorBoardLogger loaded raises an informative ErrorException.

Pass the result to nuts_sample:

using TensorBoardLogger
nuts_sample(model; callback = tensorboard_callback("logs/run"))

then view the run with tensorboard --logdir logs/run. Use chains = 1 for clean live traces. Parallel chains share one logger and interleave.

See also: progress_callback, nuts_sample.

source

BVDOutbreakSize.viable_prior_init Method

viable_prior_init(
    rng::Random.AbstractRNG,
    model;
    attempts,
    ldf
) -> Union{DynamicPPL.InitFromPrior, DynamicPPL.InitFromVector{Vector{Float64}}}

viable_prior_init(
    rng::Random.AbstractRNG,
    model;
    attempts,
    ldf
) -> Union{DynamicPPL.InitFromPrior, DynamicPPL.InitFromVector{Vector{Float64}}}

One starting point for model under ViablePrior semantics. The first of attempts independent prior draws whose log joint density is at or above that batch's median, returned as an initialisation strategy. Falls back to InitFromPrior() when no attempt gives a finite log density, or when the chosen draw cannot be wrapped as an initial vector for ldf, so a model this guard cannot evaluate still samples as before.

ldf defaults to a fresh LogDensityFunction(model). nuts_sample instead builds one and passes it to every chain, since construction re-evaluates model and so is not part of the guard's forward-evaluation-only cost.

source