How to extract priors from a Turing model?

[DominiqueMakowski]

I would just like to mention here this dicourse post, which I might have posted at the wrong place.

Assuming the following model:

# Bayesian linear regression.
@model function linear_regression(x, y)
    # Set variance prior.
    σ² ~ truncated(Normal(0, 100); lower=0)

    # Set intercept prior.
    intercept ~ Normal(0, sqrt(3))

    # Set the priors on our coefficients.
    nfeatures = size(x, 2)
    coefficients ~ MvNormal(Zeros(nfeatures), 10.0 * I)

    # Calculate all the mu terms.
    mu = intercept .+ x * coefficients
    return y ~ MvNormal(mu, σ² * I)
end

Is there a way to extract the priors set in the model? Something like:

(σ²=truncated(Normal(0, 100); lower=0), 
 intercept=Normal(0, sqrt(3)))

I couldn't find any solutions in the Turing documentation, so thanks for any pointers!

[torfjelde]

May I ask why you want the priors? We might already have some functionality for this:)

[DominiqueMakowski]

I find it very convenient to plot the priors when fitting/describing a model (to give a sense of the different distribution types). Hence, I would like to implement a function that you could conveniently run on a model, that extracts the priors of the model and plots them :)

But even besides plotting, extracting them would still be useful: I could imagine a function to create a textual "report" of the priors used by creating text like "For the parameter x, the prior used was ...".

TLDR; lots of use cases for extracting priors from a fitted model ☺️

[torfjelde]

Oh I most certainly agree that it's a useful thing to have:) It's just that in general it's difficult to get something that's "good".

For example, in the case of visualizing the prior, I'd instead suggest to just sample from the prior directly, i.e.

chain = sample(model, Prior(), 1000)

and then inspect this. This will work even in cases where there are dependencies between the prior variables, e.g. hierarchical models, while explicitly extracting the priors might not be as useful here since some of them will change depending on the particular realization.

But with the following:

using DynamicPPL: OrderedDict, SamplingContext, AbstractContext, IsParent, VarName, Distribution, evaluate!!, VarInfo
import DynamicPPL: tilde_assume, dot_tilde_assume, childcontext, setchildcontext, NodeTrait

Base.@kwdef struct PriorExtractorContext{D,Ctx} <: AbstractContext
    priors::D=OrderedDict{VarName,Any}()
    context::Ctx=SamplingContext()
end

NodeTrait(::PriorExtractorContext) = IsParent()
childcontext(context::PriorExtractorContext) = context.context
setchildcontext(parent::PriorExtractorContext, child) = PriorExtractorContext(parent.priors, child)

function tilde_assume(context::PriorExtractorContext, right, vn, vi)
    setprior!(context, vn, right)
    return tilde_assume(childcontext(context), right, vn, vi)
end

function dot_tilde_assume(context::PriorExtractorContext, right, left, vn, vi)
    setprior!(context, vn, right)
    return dot_tilde_assume(childcontext(context), right, left, vn, vi)
end

function setprior!(context::PriorExtractorContext, vn::VarName, dist::Distribution)
    context.priors[vn] = dist
end

function setprior!(context::PriorExtractorContext, vns::AbstractArray{<:VarName}, dist::Distribution)
    for vn in vns
        context.priors[vn] = dist
    end
end

function setprior!(context::PriorExtractorContext, vns::AbstractArray{<:VarName}, dists::AbstractArray{<:Distribution})
    # TODO: Support broadcasted expressions properly.
    for (vn, dist) in zip(vns, dists)
        context.priors[vn] = dist
    end
end

"""
    extract_priors(model::Model)

Extract the priors from a model. This is done by sampling from the model and
recording the distributions that are used to generate the samples.
"""
function extract_priors(model::Model)
    context = PriorExtractorContext()
    evaluate!!(model, VarInfo(), context)
    return context.priors
end

you should be able to do so:

julia> model = DynamicPPL.TestUtils.DEMO_MODELS[1]
Model{typeof(DynamicPPL.TestUtils.demo_dot_assume_dot_observe), (:x, Symbol("##arg#378")), (), (), Tuple{Vector{Float64}, DataType}, Tuple{}, DefaultContext}(DynamicPPL.TestUtils.demo_dot_assume_dot_observe, (x = [1.5, 2.0], var"##arg#378" = Vector{Float64}), NamedTuple(), DefaultContext())

julia> extract_priors(model)
OrderedDict{VarName, Any} with 4 entries:
  s[1] => InverseGamma{Float64}(…
  s[2] => InverseGamma{Float64}(…
  m[1] => Normal{Float64}(μ=0.0, σ=3.78366)
  m[2] => Normal{Float64}(μ=0.0, σ=1.02669)

For this particular model, it's indeed a hierarchical prior and so the sigmas are different.

If you want something that extracts the actual graph of the model, it's less easy to do (still possible in many cases, just requires some more complex machinery).

[DominiqueMakowski]

Thanks for the help! Just wanted to mention a slight problem, but maybe it's not the most appropriate to do it here, just let me know.
Sampling from the prior leads to an error :/

using Turing

@model function lm(x, y)
    σ² ~ InverseGamma(2, 3)
    m ~ Normal(0, sqrt(σ²))
    x ~ Normal(m, sqrt(σ²))
    y ~ Normal(m, sqrt(σ²))
    return x, y
end

chain = sample(lm, Prior(), 1000)

lm (generic function with 2 methods)

ERROR: MethodError: no method matching sample(::typeof(lm), ::Prior, ::Int64)

Closest candidates are:
  sample(::AbstractMCMC.AbstractModel, ::Turing.Inference.InferenceAlgorithm, ::Integer; kwargs...)
   @ Turing C:\Users\domma\.julia\packages\Turing\rXq36\src\inference\Inference.jl:130
  sample(::Any, ::AbstractMCMC.AbstractSampler, ::Any; kwargs...)
   @ AbstractMCMC C:\Users\domma\.julia\packages\AbstractMCMC\bE6VB\src\sample.jl:15
  sample(::Random.AbstractRNG, ::Any, ::AbstractMCMC.AbstractSampler, ::AbstractMCMC.AbstractMCMCEnsemble, ::Integer, ::Integer; kwargs...)
   @ AbstractMCMC C:\Users\domma\.julia\packages\AbstractMCMC\bE6VB\src\logdensityproblems.jl:64
  ...

Is that a known issue or am I missing something obvious?

[sunxd3]

The call to sample should be something like chain = sample(lm(0, 0), Prior(), 1000) as a model is constructed when you call lm(x, y) where x, y are observations. But I don't think the values of x and y matter here.

[DominiqueMakowski]

Based on this PR it would seems like extract_priors() have been added to DynamicPPL but after installing and loading that package it does not seem to be exported 🤔

[DominiqueMakowski]

kind tag of @torfjelde as GH discussions notifications can be easily missed :)

[devmotion]

If you install the latest version of DynamicPPL, extract_priors is available if you load DynamicPPL. Based on your question I assume though that you did not install DynamicPPL in a clean environment but probably also installed Turing. Since Turing does not support Bijectors 0.13 yet but recent releases of DynamicPPL only support Bijectors 0.13, this means probably you ended up with an old version of DynamicPPL that does not contain extract_priors.