Fitting a custom distribution

A distribution written by hand becomes fittable by answering two questions: what are its scalar parameters, and how does it rebuild itself from a flat vector. Nothing else changes.

This tutorial writes those two methods for a Weibull-backed delay, samples the posterior with a plain Metropolis sampler and reads the draws back onto the distribution. It then swaps the likelihood for a survival one and finds a maximum-a-posteriori point with an external optimiser.

using DistributionsInference, Distributions, Random

Declaring the parameters

parameter_rows returns one row per scalar parameter, carrying its name, current value, prior and support. A row with a prior is estimated; a row with prior = nothing is held fixed. reconstruct is the one place the type says how to rebuild itself from the flat vector the engine works over.

struct ToyDelay{T <: Real}
    shape::T
    scale::T
end

function Distributions.logpdf(d::ToyDelay, y::Real)
    return logpdf(Weibull(d.shape, d.scale), y)
end

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], oftype(x[1], d.scale))
end

delay = ToyDelay(2.0, 2.5)
data = [1.5, 2.0, 3.2, 1.8, 2.6]

parameter_rows(delay)

2-element Vector{NamedTuple{(:name, :value, :prior, :support)}}:
 (name = :shape, value = 2.0, prior = Distributions.LogNormal{Float64}(μ=0.6931471805599453, σ=0.2), support = (0.0, Inf))
 (name = :scale, value = 2.5, prior = nothing, support = (0.0, Inf))

The fields are typed T <: Real and the fixed scale is rebuilt through oftype, so reconstruct returns whatever number type it is handed. A gradient-based sampler differentiates through this call and passes a dual number rather than a Float64; a concretely typed field would reject it.

DistributionsInference.reconstruct(delay, [1 // 2])

Main.var"Main".ToyDelay{Rational{Int64}}(1//2, 5//2)

Scoring a parameter vector

distribution_to_logdensity packages the distribution and the data into a log-density over the estimated rows alone, and flat_dimension counts them.

prob = distribution_to_logdensity(delay, data)
DistributionsInference.flat_dimension(delay)

1

logdensity adds the prior's log-density at that value to the data likelihood of the distribution rebuilt there.

DistributionsInference.logdensity(prob, [2.0])

-6.133158008868996

Sampling the log-density directly

prob implements the LogDensityProblems interface, so any consumer of that interface can drive it. distribution_to_advancedmh is a ready-made one: it drives AdvancedMH's random-walk Metropolis, the smallest sampler to reach for, and hands back a FlexiChain keyed by the estimated rows' dotted names, the naming contract every readback verb here uses. A gradient backend added through LogDensityProblemsAD opens up AdvancedHMC the same way.

It samples on the unconstrained scale (via to_constrained, so Bijectors is loaded alongside AdvancedMH), so a random-walk proposal can never land outside shape's positive support in the first place — no hand-written -Inf guard needed. The proposal is sized to the estimated dimension, one row here.

using AdvancedMH, Bijectors
using LinearAlgebra: I

dim = DistributionsInference.flat_dimension(delay)
sampler = RWMH(MvNormal(zeros(dim), 0.05^2 * I))

Random.seed!(1)
chain = distribution_to_advancedmh(delay, data, sampler, 2000; burnin = 1000)

╭─FlexiChain (1000 iterations, 1 chain) ───────────────────────────────────────╮
│ ↓ iter  = 1:1000                                                             │
│ → chain = 1:1                                                                │
│                                                                              │
│ Parameters (1) ── Symbol                                                     │
│  Float64  shape                                                              │
│                                                                              │
│ Extras (0)                                                                   │
│  (none)                                                                      │
╰──────────────────────────────────────────────────────────────────────────────╯

Reading the fit back onto the distribution

inference_to_distribution reduces the chain straight back to a ToyDelay, the reduction (here the posterior mean) positional.

fitted = inference_to_distribution(delay, chain, mean)
fitted.shape

2.1900427054667624

The fixed row came back at its template value.

fitted.scale

2.5

inference_to_distributions keeps every draw instead of reducing them, for a per-draw posterior-predictive summary.

all_fitted = inference_to_distributions(delay, chain)
quantile([mean(Weibull(d.shape, d.scale)) for d in all_fitted], [0.025, 0.975])

2-element Vector{Float64}:
 2.2140114704768803
 2.252113721425234

Swapping the likelihood

loglik is any reducer (obj, data) -> Real. Right-censored records, known only to exceed a bound, score through logccdf rather than logpdf.

function survival_loglik(obj, records)
    return sum(y -> logccdf(Weibull(obj.shape, obj.scale), y), records)
end
bounds = [1.0, 1.5, 2.0]
survival_prob = distribution_to_logdensity(delay, bounds;
    loglik = survival_loglik)
DistributionsInference.logdensity(survival_prob, [2.0])

-1.162647801330518

The parameter inventory, the priors and the readback never see the reducer, so everything above works unchanged: distribution_to_advancedmh takes the same loglik keyword distribution_to_logdensity does.

Random.seed!(1)
survival_chain = distribution_to_advancedmh(delay, bounds, sampler, 2000;
    burnin = 1000, loglik = survival_loglik)
inference_to_distribution(delay, survival_chain, mean).shape

1.998080449142201

logccdf for a Weibull differentiates, so Sampling with Turing hands this same reducer to NUTS. A reducer that draws random replicates internally to marginalise a latent variable is noisy as well as usually non-differentiable, so keep it on a gradient-free sampler and raise its replicate count until the posterior summary stops moving.

A point estimate instead of a posterior

The objective is an unnormalised log-density, so an optimiser can find its mode directly. optimise_distribution runs that round trip in one call and hands back a ToyDelay, once Bijectors and an optimiser package are loaded. The optimiser itself stays the caller's choice.

using Bijectors, Optim

map_fit = optimise_distribution(delay, data, LBFGS())
map_fit.shape

2.2712643851493333

The fit starts at the template's own parameter values rather than at an arbitrary zero; init takes a starting point on the same (constrained) scale delay's own fields are in — no manual transform needed to try a different start — and loglik swaps the reducer as it does for distribution_to_logdensity.

optimise_distribution(delay, data, LBFGS(); init = [1.5]).shape

2.27126438512847

The steps are still individually available. logdensity_to_objective gives the callable, to_unconstrained maps a constrained starting point onto the scale the optimiser actually searches, minimise runs it, and objective_to_distribution maps the minimiser back onto the distribution — the same round trip optimise_distribution composes above, from the same starting point.

f = logdensity_to_objective(prob)
z0 = DistributionsInference.to_unconstrained(prob, [1.5])
objective_to_distribution(prob, DistributionsInference.minimise(
    f, z0, LBFGS())).shape

2.27126438512847

That is the maximum-a-posteriori point, because logdensity always adds an estimated row's own prior. Maximum likelihood needs a prior flat enough to leave the likelihood alone.

struct ToyDelayML{T <: Real}
    shape::T
    scale::T
end

function Distributions.logpdf(d::ToyDelayML, y::Real)
    return logpdf(Weibull(d.shape, d.scale), y)
end

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

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

ml_delay = ToyDelayML(2.0, 1.0)
optimise_distribution(ml_delay, data, LBFGS()).shape

0.9665449258546485

The data pull the shape above the prior mean of 2 either way, and the diffuse prior pulls it back less.

Next

• Sampling with Turing fits this same distribution through DynamicPPL and reads the chain back with the same calls.

• Fitting a composed distribution runs both routes against a tree built by ComposedDistributions. Its extension writes the protocol methods.

• Public API lists the rest of the protocol (with_priors, estimated_rows, extra_logprior).