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
using FlexiChains: FlexiChain, Parameter

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. AdvancedMH's random-walk Metropolis is the smallest one to reach for; a gradient backend added through LogDensityProblemsAD opens up AdvancedHMC the same way. The wrapper below returns -Inf off shape's positive support, which a random-walk proposal does not respect on its own.

using AdvancedMH
using LinearAlgebra: I

model = AdvancedMH.DensityModel() do x
    any(<=(0), x) ? -Inf : DistributionsInference.logdensity(prob, x)
end
sampler = RWMH(MvNormal(zeros(1), 0.05^2 * I))
transitions = sample(Xoshiro(1), model, sampler, 2000;
    param_names = ["shape"], progress = false)
draws = [t.params for t in transitions][1001:end]
length(draws)

1000

Reading the fit back onto the distribution

The readback reads a FlexiChain keyed by the estimated rows' dotted names, the naming contract every readback verb here uses. Building that chain out of a sampler's raw draws is FlexiChains' own constructor rather than a step this package owns.

niter = length(draws)
chain = FlexiChain{Symbol}(niter, 1,
    Dict(Parameter(:shape) => reshape(first.(draws), niter, 1)))

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

point_estimate reduces the chain straight back to a ToyDelay, posterior mean by default.

fitted = point_estimate(delay, chain)
fitted.shape

2.091519040483683

The fixed row came back at its template value.

fitted.scale

2.5

distribution_params is the params-first primitive underneath, the same reduction keyed by dotted name before the object is rebuilt.

distribution_params(delay, chain)

(shape = 2.091519040483683,)

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

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

2-element Vector{Float64}:
 2.2140185527904634
 2.27278889857716

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 on survival_prob.

survival_model = AdvancedMH.DensityModel() do x
    any(<=(0), x) ? -Inf : DistributionsInference.logdensity(survival_prob, x)
end
survival_transitions = sample(Xoshiro(1), survival_model, sampler, 2000;
    param_names = ["shape"], progress = false)
survival_draws = [t.params for t in survival_transitions][1001:end]
survival_chain = FlexiChain{Symbol}(length(survival_draws), 1,
    Dict(Parameter(:shape) => reshape(
        first.(survival_draws), length(survival_draws), 1)))
point_estimate(delay, survival_chain).shape

1.9388763392990207

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 starts it somewhere else, and loglik swaps the reducer as it does for distribution_to_logdensity.

The steps are still individually available. logdensity_to_objective gives the callable, minimise runs the optimiser, and objective_to_distribution maps the minimiser back onto the distribution.

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

2.2712643851493333

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