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 distribution names its own parameters and becomes fittable through a plain log-density that any sampler can read.

Why DistributionsInference?

• Fitting a distribution usually means rewriting it inside one probabilistic programming language's macros. Here it names its own scalar parameters once and every sampler reads that declaration.
• distribution_to_logdensity turns a distribution and its data into a LogDensityProblems problem, which any consumer of that interface can drive.
• A distribution declares its parameters as a table of rows, one row per scalar parameter carrying its name, value, prior and support. Attaching a prior to a row is what makes that parameter estimated. The rows are plain NamedTuples, so the inventory is a row table any Tables.jl consumer reads, without this package depending on Tables.jl.
• optimise_distribution fits a distribution to data with an optimiser of your choosing in one call, for a point estimate rather than a posterior, and the pieces it composes stay available for a fit you drive yourself.
• inference_to_distribution post-processes sampler output onto a fitted distribution in one call: the Monte Carlo posterior predictive by default, or a plug-in point estimate with a reduction named at the call site.
• Supporting a new chain type takes one added method.
• The distribution that comes back is the same kind of object that went in, so a fitted Gamma is a Gamma.

Getting started

See documentation for a full walkthrough.

Start with a plain Distributions.jl Gamma. The package ships no parameter_rows/reconstruct methods for Distributions.jl types, so write those two once per family. Nothing else about the distribution changes.

using DistributionsInference, Distributions, Random

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

function DistributionsInference.reconstruct(::Gamma, x::AbstractVector)
    return Gamma(x[1], x[2])
end

parameter_rows sets the estimation boundary. A row carrying a prior is estimated; a row with prior = nothing stays at its value. distribution_to_logdensity packages a template distribution and the data into a log-density over the estimated rows.

template = Gamma(2.0, 1.0)
data = rand(Xoshiro(1), Gamma(2.0, 1.0), 200)
prob = distribution_to_logdensity(template, data)
DistributionsInference.flat_dimension(template)

prob is a LogDensityProblems problem, so any sampler that consumes that interface can drive it. distribution_to_advancedmh drives one for you: AdvancedMH's random-walk Metropolis, sampled on the unconstrained scale so a proposal never lands outside a prior's support and needs no hand-written guard for it. It hands back a FlexiChains.FlexiChain keyed by the names parameter_rows declared.

using AdvancedMH, Bijectors
using LinearAlgebra: I

dim = DistributionsInference.flat_dimension(template)
chain = distribution_to_advancedmh(
    template, data, RWMH(MvNormal(zeros(dim), 0.05^2 * I)), 4000; burnin = 2000)

inference_to_distribution reads the chain back onto the distribution. Naming a reduction (here mean) returns the plug-in point estimate, a Gamma; leaving it off returns the Monte Carlo posterior predictive instead, an equal-weight mixture over every draw. The chain readback is a package extension, so add and load FlexiChains for this last step.

using FlexiChains

inference_to_distribution(template, chain, mean)

The getting started guide goes further, with a distribution type of your own and sampling through Turing.

Related packages

• Distributions.jl defines the distributions this fits; a type that has a logpdf and can name its scalar parameters is a candidate.
• ComposedDistributions.jl builds a distribution by composing others into chains, branches and outcomes; a package extension here reads a composed tree's generated codec directly, so its estimated leaves, pooled and shared parameters included, are fittable as they stand.
• ModifiedDistributions.jl wraps a distribution to change one behaviour, such as rescaling, likelihood weighting or a hazard shift; a modifier used as a leaf inside a composed tree is already fittable, and a package extension here makes a standalone modifier fittable too.
• ReparameterisedDistributions.jl switches a family between parameter conventions, so a distribution can be fitted in the coordinates its priors were elicited in rather than the family's native ones.
• CensoredDistributions.jl applies primary event censoring, interval censoring and right truncation to a delay, and returns a Distributions.jl distribution.
• ConvolvedDistributions.jl builds the distribution of a sum, difference or product of independent delays, again as a Distributions.jl distribution.

ComposedDistributions is the only one of these covered by an extension today. For the rest, as for a plain Distributions.jl distribution, the two protocol methods are yours to write.

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.

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.

Part of the EpiAware ecosystem

DistributionsInference is part of EpiAware, a set of composable tools for infectious disease modelling. See the other packages in the ecosystem.

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.