DistributionsInference

Documentation	Build Status	Code Quality	License & DOI	Downloads

ForwardDiff	ReverseDiff (tape)	Enzyme forward	Enzyme reverse	Mooncake reverse	Mooncake forward

The inference layer for the EpiAware composable-modelling stack: a PPL-neutral fit protocol and log-density engine, so any object that names its own parameters is fittable without committing to a probabilistic programming language.

Why DistributionsInference?

• Fitting an object today usually means committing to one PPL's macros; here a type opts in by naming its own scalar parameters, and it becomes fittable everywhere, hand-rolled sampler and PPL alike.
• The log-density this produces has no PPL dependency, so it evaluates through whatever LogDensityProblems-compatible sampler a project already uses.
• A parameter becomes estimated by attaching a prior at the row level; nothing else about a type needs to change to be fitted.
• Reading a fitted chain back onto a concrete object is the same one call whether the chain came from a hand-rolled sampler or from Turing.
• Turing and Bijectors support are opt-in layers over the same protocol, not requirements, so a project can start with the bare log-density and add a PPL later without rewriting its model.
• Ported from ComposedDistributions.jl's own fit protocol, so a composed distribution and a plain hand-written type share one estimation surface.

Getting started

See documentation for a full walkthrough.

A type becomes fittable by naming its scalar parameters and how to rebuild itself from a flat vector — no other change needed.

using DistributionsInference, Distributions, Random

struct ToyDelay
    shape::Float64
    scale::Float64
end

Distributions.logpdf(d::ToyDelay, y::Real) = logpdf(Gamma(d.shape, d.scale), y)

function DistributionsInference.parameter_rows(d::ToyDelay)
    return [(name = :shape, value = d.shape,
            prior = LogNormal(log(2.0), 0.2), support = (0.0, Inf)),
        (name = :scale, value = d.scale, prior = nothing,
            support = (0.0, Inf))]
end

function DistributionsInference.reconstruct(d::ToyDelay, x::AbstractVector)
    return ToyDelay(x[1], d.scale)
end

as_logdensity packages a template object and data into a log-density over just the one estimated parameter (shape; scale stays fixed).

leaf = ToyDelay(2.0, 1.0)
data = [1.5, 2.0, 3.2, 1.8, 2.6]
prob = DistributionsInference.as_logdensity(leaf, data)
DistributionsInference.flat_dimension(leaf)

Any LogDensityProblems-compatible sampler can drive prob; here is the tiniest one, a ten-line random-walk Metropolis, so this stays self-contained.

function toy_sample(prob, x0, n; step = 0.2, rng = Xoshiro(1))
    x, lp = copy(x0), DistributionsInference.logdensity(prob, x0)
    draws = Vector{Vector{Float64}}(undef, n)
    for i in 1:n
        prop = x .+ step .* randn(rng, length(x))
        if all(>(0), prop)
            lp_prop = DistributionsInference.logdensity(prob, prop)
            log(rand(rng)) < lp_prop - lp && ((x, lp) = (prop, lp_prop))
        end
        draws[i] = copy(x)
    end
    return draws
end

draws = toy_sample(prob, [2.0], 500)

readback reduces the draws to a fitted ToyDelay, through the same dotted-name chain a real PPL's sampler would hand back.

chain = DistributionsInference.to_flexichain(leaf, draws)
fit = DistributionsInference.readback(leaf, chain)
fit.shape

The getting started guide carries this same object further: reading every draw with readback_draws, and sampling with Turing instead of the toy sampler above.

Where to learn more

• Want to get started running code? See the getting started guide.
• Want to understand the API? See the API reference.
• Want to see the code? Check out our GitHub repository.

Related packages

• ComposedDistributions.jl is the package this fit protocol was ported from; a package extension here reads a composed tree's generated codec directly, so its estimated leaves (including pooled and shared parameters) are fittable with no extra glue.
• ModifiedDistributions.jl support is landing next: a standalone extension will let a modified or weighted distribution opt into the same protocol.

Getting help

For usage questions, ask on the Julia Discourse (the SciML or usage categories) or the epinowcast community forum, our home for epidemiological modelling questions. Please use GitHub issues for bug reports and feature requests only.

Contributing

We welcome contributions and new contributors! Please open an issue or pull request on GitHub. This package follows ColPrac and the SciML style.

How to cite

If you use DistributionsInference in your work, please cite it. Citation metadata lives in CITATION.cff, which GitHub renders as a "Cite this repository" button on the repository page.

Code of conduct

Please note that the DistributionsInference project is released with a Contributor Code of Conduct. By contributing, you agree to abide by its terms.