Advanced Usage
How to Define a Customized Distribution
Turing.jl
supports the use of distributions from the Distributions.jl package. By extension, it also supports the use of customized distributions by defining them as subtypes of Distribution
type of the Distributions.jl package, as well as corresponding functions.
Below shows a workflow of how to define a customized distribution, using our own implementation of a simple Uniform
distribution as a simple example.
1. Define the Distribution Type
First, define a type of the distribution, as a subtype of a corresponding distribution type in the Distributions.jl package.
struct CustomUniform <: ContinuousUnivariateDistribution end
2. Implement Sampling and Evaluation of the log-pdf
Second, define rand
and logpdf
, which will be used to run the model.
# sample in [0, 1]
Distributions.rand(rng::AbstractRNG, d::CustomUniform) = rand(rng)
# p(x) = 1 → logp(x) = 0
Distributions.logpdf(d::CustomUniform, x::Real) = zero(x)
3. Define Helper Functions
In most cases, it may be required to define some helper functions.
3.1 Domain Transformation
Certain samplers, such as HMC
, require the domain of the priors to be unbounded. Therefore, to use our CustomUniform
as a prior in a model we also need to define how to transform samples from [0, 1]
to ℝ
. To do this, we simply need to define the corresponding Bijector
from Bijectors.jl
, which is what Turing.jl
uses internally to deal with constrained distributions.
To transform from [0, 1]
to ℝ
we can use the Logit
bijector:
Bijectors.bijector(d::CustomUniform) = Logit(0.0, 1.0)
You'd do the exact same thing for ContinuousMultivariateDistribution
and ContinuousMatrixDistribution
. For example, Wishart
defines a distribution over positive-definite matrices and so bijector
returns a PDBijector
when called with a Wishart
distribution as an argument. For discrete distributions, there is no need to define a bijector; the Identity
bijector is used by default.
Alternatively, for UnivariateDistribution
we can define the minimum
and maximum
of the distribution
Distributions.minimum(d::CustomUniform) = 0.0
Distributions.maximum(d::CustomUniform) = 1.0
and Bijectors.jl
will return a default Bijector
called TruncatedBijector
which makes use of minimum
and maximum
derive the correct transformation.
Internally, Turing basically does the following when it needs to convert a constrained distribution to an unconstrained distribution, e.g. when sampling using HMC
:
dist = Gamma(2,3)
b = bijector(dist)
transformed_dist = transformed(dist, b) # results in distribution with transformed support + correction for logpdf
Bijectors.UnivariateTransformed{Distributions.Gamma{Float64}, Base.Fix1{typ
eof(broadcast), typeof(log)}}(
dist: Distributions.Gamma{Float64}(α=2.0, θ=3.0)
transform: Base.Fix1{typeof(broadcast), typeof(log)}(broadcast, log)
)
and then we can call rand
and logpdf
as usual, where
rand(transformed_dist)
returns a sample in the unconstrained space, andlogpdf(transformed_dist, y)
returns the log density of the original distribution, but withy
living in the unconstrained space.
To read more about Bijectors.jl, check out the project README.
Update the accumulated log probability in the model definition
Turing accumulates log probabilities internally in an internal data structure that is accessible through
the internal variable __varinfo__
inside of the model definition (see below for more details about model internals).
However, since users should not have to deal with internal data structures, a macro Turing.@addlogprob!
is provided
that increases the accumulated log probability. For instance, this allows you to
include arbitrary terms in the likelihood
using Turing
myloglikelihood(x, μ) = loglikelihood(Normal(μ, 1), x)
@model function demo(x)
μ ~ Normal()
Turing.@addlogprob! myloglikelihood(x, μ)
end
demo (generic function with 2 methods)
and to reject samples:
using Turing
using LinearAlgebra
@model function demo(x)
m ~ MvNormal(zero(x), I)
if dot(m, x) < 0
Turing.@addlogprob! -Inf
# Exit the model evaluation early
return nothing
end
x ~ MvNormal(m, I)
return nothing
end
demo (generic function with 2 methods)
Note that @addlogprob!
always increases the accumulated log probability, regardless of the provided
sampling context. For instance, if you do not want to apply Turing.@addlogprob!
when evaluating the
prior of your model but only when computing the log likelihood and the log joint probability, then you
should check the type of the internal variable __context_
such as
if DynamicPPL.leafcontext(__context__) !== Turing.PriorContext()
Turing.@addlogprob! myloglikelihood(x, μ)
end
Model Internals
The @model
macro accepts a function definition and rewrites it such that call of the function generates a Model
struct for use by the sampler.
Models can be constructed by hand without the use of a macro.
Taking the gdemo
model as an example, the macro-based definition
using Turing
@model function gdemo(x)
# Set priors.
s² ~ InverseGamma(2, 3)
m ~ Normal(0, sqrt(s²))
# Observe each value of x.
@. x ~ Normal(m, sqrt(s²))
end
model = gdemo([1.5, 2.0])
DynamicPPL.Model{typeof(Main.var"##WeaveSandBox#326".gdemo), (:x,), (), (),
Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}(Main.var"##Wea
veSandBox#326".gdemo, (x = [1.5, 2.0],), NamedTuple(), DynamicPPL.DefaultCo
ntext())
can be implemented also (a bit less generally) with the macro-free version
using Turing
# Create the model function.
function gdemo(model, varinfo, context, x)
# Assume s² has an InverseGamma distribution.
s², varinfo = DynamicPPL.tilde_assume!!(
context, InverseGamma(2, 3), Turing.@varname(s²), varinfo
)
# Assume m has a Normal distribution.
m, varinfo = DynamicPPL.tilde_assume!!(
context, Normal(0, sqrt(s²)), Turing.@varname(m), varinfo
)
# Observe each value of x[i] according to a Normal distribution.
return DynamicPPL.dot_tilde_observe!!(
context, Normal(m, sqrt(s²)), x, Turing.@varname(x), varinfo
)
end
gdemo(x) = Turing.Model(gdemo, (; x))
# Instantiate a Model object with our data variables.
model = gdemo([1.5, 2.0])
DynamicPPL.Model{typeof(Main.var"##WeaveSandBox#326".gdemo), (:x,), (), (),
Tuple{Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext}(Main.var"##Wea
veSandBox#326".gdemo, (x = [1.5, 2.0],), NamedTuple(), DynamicPPL.DefaultCo
ntext())
Task Copying
Turing copies Julia tasks to
deliver efficient inference algorithms, but it also provides alternative slower
implementation as a fallback. Task copying is enabled by default. Task copying
requires us to use the TapedTask
facility which is provided by
Libtask to create tasks.