Using LDFs in model evaluation

As mentioned on the previous page, a LogDensityFunction contains more information than just the mapping from vectors to log-densities. It also contains enough information to generate an initialisation strategy, and on top of that it directly wraps a transform strategy.

This means that, once you have constructed a LogDensityFunction, you can use it in conjunction with a parameter vector to evaluate models using the init!! framework established previously.

Let's start by generating the same LogDensityFunction as on the previous page:

using DynamicPPL, Distributions

@model function f()
    x ~ Normal()
    y ~ Beta(2, 2)
    return nothing
end
model = f()

accs = OnlyAccsVarInfo(VectorValueAccumulator())
_, accs = init!!(model, accs, InitFromPrior(), LinkAll())
vector_values = get_vector_values(accs)

ldf = LogDensityFunction(model, getlogjoint_internal, vector_values)

LogDensityFunction{Model{typeof(Main.f), (), (), (), Tuple{}, Tuple{}, DefaultContext, false}, Nothing, LinkAll, typeof(getlogjoint_internal), VarNamedTuple{(:x, :y), Tuple{DynamicPPL.RangeAndLinked{Tuple{}}, DynamicPPL.RangeAndLinked{Tuple{}}}}, Nothing, Vector{Float64}, DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}}(Model{typeof(Main.f), (), (), (), Tuple{}, Tuple{}, DefaultContext, false}(Main.f, NamedTuple(), NamedTuple(), DefaultContext()), nothing, LinkAll(), DynamicPPL.getlogjoint_internal, VarNamedTuple(x = DynamicPPL.RangeAndLinked{Tuple{}}(1:1, true, ()), y = DynamicPPL.RangeAndLinked{Tuple{}}(2:2, true, ())), nothing, 2, DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}((LogPrior = LogPriorAccumulator(0.0), LogJacobian = LogJacobianAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0))))

Evaluating models using vectorised parameters

You can regenerate the initialisation strategy using the constructor InitFromVector(vect, ldf):

init_strategy = InitFromVector([3.0, 4.0], ldf)

InitFromVector{Vector{Float64}, VarNamedTuple{(:x, :y), Tuple{DynamicPPL.RangeAndLinked{Tuple{}}, DynamicPPL.RangeAndLinked{Tuple{}}}}, LinkAll}([3.0, 4.0], VarNamedTuple(x = DynamicPPL.RangeAndLinked{Tuple{}}(1:1, true, ()), y = DynamicPPL.RangeAndLinked{Tuple{}}(2:2, true, ())), LinkAll())

and access the transform strategy via

ldf.transform_strategy

LinkAll()

With these two pieces of information, we can add whatever accumulators we like to the mix, and then evaluate using init!! as shown before. For example, let's say we'd like to know what raw values the vector [3.0, 4.0] corresponds to.

accs = OnlyAccsVarInfo(RawValueAccumulator(false))
_, accs = init!!(model, accs, init_strategy, ldf.transform_strategy)
get_raw_values(accs)

VarNamedTuple
├─ x => 3.0
└─ y => 0.9820137900379085

Unsurprisingly, the raw (i.e., untransformed) value of x is 3.0, and y is equal to:

using StatsFuns: logistic
logistic(4.0)

0.9820137900379085

Notice that you can include all the necessary log-density accumulators in the above:

accs = OnlyAccsVarInfo(
    LogPriorAccumulator(), LogLikelihoodAccumulator(), LogJacobianAccumulator()
)
_, accs = init!!(model, accs, init_strategy, ldf.transform_strategy)
accs

OnlyAccsVarInfo
 └─ AccumulatorTuple with 3 accumulators
    ├─ LogPrior => LogPriorAccumulator(-7.663478919812238)
    ├─ LogLikelihood => LogLikelihoodAccumulator(0.0)
    └─ LogJacobian => LogJacobianAccumulator(4.03629985583562)

and from this you can, in one fell swoop, obtain the log-prior, log-likelihood, and log-Jacobian corresponding to the vector [3.0, 4.0].

In fact, this is exactly what calling LogDensityProblems.logdensity(ldf, [3.0, 4.0] does, except that it also collapses this information into a single number based on which log-density getter you specified when creating the LogDensityFunction. For example, in the above ldf, we used getlogjoint_internal, which is equal to logprior + loglikelihood - logjacobian:

using LogDensityProblems

LogDensityProblems.logdensity(ldf, [3.0, 4.0])

-11.699778775647857

This represents a loss of information since we no longer know the individual components of the log-density. LogDensityFunction itself must do this to obey the LogDensityProblems interface, but as a user of DynamicPPL you need not be limited to it; you can use InitFromVector to obtain all the information you need.

Obtaining new vectorised parameters

We can also do the reverse process here, which is to obtain new vectorised parameters that are consistent with this LogDensityFunction.

Suppose that we want to evaluate the LogDensityFunction at the new raw values x = 5.0 and y = 0.6. The corresponding vectorised parameters should be

using StatsFuns: logit
[5.0, logit(0.6)]

2-element Vector{Float64}:
 5.0
 0.4054651081081642

(Of course, in general you can't just write these down by hand, since it will depend on the model.)

Naturally, we need to re-evaluate the model with the initialisation strategy we are interested in. We can use a special accumulator, VectorParamAccumulator, to get the newly vectorised set of parameters. This accumulator must take the LogDensityFunction as an argument, since it uses the information stored in there to ensure that the values it collects are consistent with the LogDensityFunction.

Note that we must use the same transform strategy as the one contained in the LogDensityFunction, or else we will be generating vectorised parameters that are inconsistent with the LogDensityFunction (an error will be thrown if that happens).

init_strategy = InitFromParams(VarNamedTuple(; x=5.0, y=0.6))
transform_strategy = ldf.transform_strategy

accs = OnlyAccsVarInfo(VectorParamAccumulator(ldf))
_, accs = init!!(model, accs, init_strategy, transform_strategy)
vec = get_vector_params(accs)

2-element Vector{Float64}:
 5.0
 0.4054651081081642

You can of course also bundle any extra accumulators you like into the above if you are interested not only in the vectorised parameters but also (e.g.) the log-density. This allows you to obtain all the information you need with only one model evaluation.

accs = OnlyAccsVarInfo(VectorValueAccumulator())
_, accs = init!!(model, accs, init_strategy, transform_strategy)
vec_vals = get_vector_values(accs)

VarNamedTuple
├─ x => LinkedVectorValue{Vector{Float64}, ComposedFunction{DynamicPPL.UnwrapSingletonTransform{Tuple{}}, ComposedFunction{typeof(identity), DynamicPPL.ReshapeTransform{Tuple{Int64}, Tuple{}}}}, Tuple{}}([5.0], DynamicPPL.UnwrapSingletonTransform{Tuple{}}(()) ∘ (identity ∘ DynamicPPL.ReshapeTransform{Tuple{Int64}, Tuple{}}((1,), ())), ())
└─ y => LinkedVectorValue{Vector{Float64}, ComposedFunction{DynamicPPL.UnwrapSingletonTransform{Tuple{}}, ComposedFunction{Bijectors.Inverse{Bijectors.Logit{Float64, Float64}}, DynamicPPL.ReshapeTransform{Tuple{Int64}, Tuple{}}}}, Tuple{}}([0.4054651081081642], DynamicPPL.UnwrapSingletonTransform{Tuple{}}(()) ∘ (Bijectors.Inverse{Bijectors.Logit{Float64, Float64}}(Bijectors.Logit{Float64, Float64}(0.0, 1.0)) ∘ DynamicPPL.ReshapeTransform{Tuple{Int64}, Tuple{}}((1,), ())), ())

to_vector_params(vec_vals, ldf)

2-element Vector{Float64}:
 5.0
 0.4054651081081642