Split TestUtils module up

Author: penelopeysm

src/DynamicPPL.jl

@@ -186,7 +186,7 @@ include("context_implementations.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")
 include("submodel_macro.jl")
-include("test_utils.jl")
+include("test_utils/main.jl")
 include("transforming.jl")
 include("logdensityfunction.jl")
 include("model_utils.jl")

src/test_utils/contexts.jl

@@ -0,0 +1,62 @@
+# contexts.jl
+# -----------
+#
+# Utilities for testing contexts.
+
+"""
+    test_context_interface(context)
+
+Test that `context` implements the `AbstractContext` interface.
+"""
+function test_context_interface(context)
+    # Is a subtype of `AbstractContext`.
+    @test context isa DynamicPPL.AbstractContext
+    # Should implement `NodeTrait.`
+    @test DynamicPPL.NodeTrait(context) isa Union{DynamicPPL.IsParent,DynamicPPL.IsLeaf}
+    # If it's a parent.
+    if DynamicPPL.NodeTrait(context) == DynamicPPL.IsParent
+        # Should implement `childcontext` and `setchildcontext`
+        @test DynamicPPL.setchildcontext(context, DynamicPPL.childcontext(context)) ==
+            context
+    end
+end
+
+"""
+Context that multiplies each log-prior by mod
+used to test whether varwise_logpriors respects child-context.
+"""
+struct TestLogModifyingChildContext{T,Ctx} <: DynamicPPL.AbstractContext
+    mod::T
+    context::Ctx
+end
+function TestLogModifyingChildContext(
+    mod=1.2, context::DynamicPPL.AbstractContext=DynamicPPL.DefaultContext()
+)
+    return TestLogModifyingChildContext{typeof(mod),typeof(context)}(mod, context)
+end
+
+DynamicPPL.NodeTrait(::TestLogModifyingChildContext) = DynamicPPL.IsParent()
+DynamicPPL.childcontext(context::TestLogModifyingChildContext) = context.context
+function DynamicPPL.setchildcontext(context::TestLogModifyingChildContext, child)
+    return TestLogModifyingChildContext(context.mod, child)
+end
+function DynamicPPL.tilde_assume(context::TestLogModifyingChildContext, right, vn, vi)
+    value, logp, vi = DynamicPPL.tilde_assume(context.context, right, vn, vi)
+    return value, logp * context.mod, vi
+end
+function DynamicPPL.dot_tilde_assume(
+    context::TestLogModifyingChildContext, right, left, vn, vi
+)
+    value, logp, vi = DynamicPPL.dot_tilde_assume(context.context, right, left, vn, vi)
+    return value, logp * context.mod, vi
+end
+function DynamicPPL.tilde_observe(context::TestLogModifyingChildContext, right, left, vi)
+    logp, vi = DynamicPPL.tilde_observe(context.context, right, left, vi)
+    return logp * context.mod, vi
+end
+function DynamicPPL.dot_tilde_observe(
+    context::TestLogModifyingChildContext, right, left, vi
+)
+    logp, vi = DynamicPPL.dot_tilde_observe(context.context, right, left, vi)
+    return logp * context.mod, vi
+end

src/test_utils/main.jl

@@ -0,0 +1,22 @@
+module TestUtils
+
+using AbstractMCMC
+using DynamicPPL
+using LinearAlgebra
+using Distributions
+using Test
+
+using Random: Random
+using Bijectors: Bijectors
+using Accessors: Accessors
+
+# For backwards compat.
+using DynamicPPL: varname_leaves, update_values!!
+
+include("model_interface.jl")
+include("models.jl")
+include("contexts.jl")
+include("varinfo.jl")
+include("sampler.jl")
+
+end

src/test_utils/model_interface.jl

@@ -0,0 +1,118 @@
+# model_interface.jl
+# ------------------
+#
+# This file contains the functions that the models inside models.jl should
+# implement.
+
+"""
+    logprior_true(model, args...)
+
+Return the `logprior` of `model` for `args`.
+
+This should generally be implemented by hand for every specific `model`.
+
+See also: [`logjoint_true`](@ref), [`loglikelihood_true`](@ref).
+"""
+function logprior_true end
+
+"""
+    loglikelihood_true(model, args...)
+
+Return the `loglikelihood` of `model` for `args`.
+
+This should generally be implemented by hand for every specific `model`.
+
+See also: [`logjoint_true`](@ref), [`logprior_true`](@ref).
+"""
+function loglikelihood_true end
+
+"""
+    logjoint_true(model, args...)
+
+Return the `logjoint` of `model` for `args`.
+
+Defaults to `logprior_true(model, args...) + loglikelihood_true(model, args..)`.
+
+This should generally be implemented by hand for every specific `model`
+so that the returned value can be used as a ground-truth for testing things like:
+
+1. Validity of evaluation of `model` using a particular implementation of `AbstractVarInfo`.
+2. Validity of a sampler when combined with DynamicPPL by running the sampler twice: once targeting ground-truth functions, e.g. `logjoint_true`, and once targeting `model`.
+
+And more.
+
+See also: [`logprior_true`](@ref), [`loglikelihood_true`](@ref).
+"""
+function logjoint_true(model::Model, args...)
+    return logprior_true(model, args...) + loglikelihood_true(model, args...)
+end
+
+"""
+    logjoint_true_with_logabsdet_jacobian(model::Model, args...)
+
+Return a tuple `(args_unconstrained, logjoint)` of `model` for `args`.
+
+Unlike [`logjoint_true`](@ref), the returned logjoint computation includes the
+log-absdet-jacobian adjustment, thus computing logjoint for the unconstrained variables.
+
+Note that `args` are assumed be in the support of `model`, while `args_unconstrained`
+will be unconstrained.
+
+This should generally not be implemented directly, instead one should implement
+[`logprior_true_with_logabsdet_jacobian`](@ref) for a given `model`.
+
+See also: [`logjoint_true`](@ref), [`logprior_true_with_logabsdet_jacobian`](@ref).
+"""
+function logjoint_true_with_logabsdet_jacobian(model::Model, args...)
+    args_unconstrained, lp = logprior_true_with_logabsdet_jacobian(model, args...)
+    return args_unconstrained, lp + loglikelihood_true(model, args...)
+end
+
+"""
+    logprior_true_with_logabsdet_jacobian(model::Model, args...)
+
+Return a tuple `(args_unconstrained, logprior_unconstrained)` of `model` for `args...`.
+
+Unlike [`logprior_true`](@ref), the returned logprior computation includes the
+log-absdet-jacobian adjustment, thus computing logprior for the unconstrained variables.
+
+Note that `args` are assumed be in the support of `model`, while `args_unconstrained`
+will be unconstrained.
+
+See also: [`logprior_true`](@ref).
+"""
+function logprior_true_with_logabsdet_jacobian end
+
+"""
+    varnames(model::Model)
+
+Return a collection of `VarName` as they are expected to appear in the model.
+
+Even though it is recommended to implement this by hand for a particular `Model`,
+a default implementation using [`SimpleVarInfo{<:Dict}`](@ref) is provided.
+"""
+function varnames(model::Model)
+    return collect(
+        keys(last(DynamicPPL.evaluate!!(model, SimpleVarInfo(Dict()), SamplingContext())))
+    )
+end
+
+"""
+    posterior_mean(model::Model)
+
+Return a `NamedTuple` compatible with `varnames(model)` where the values represent
+the posterior mean under `model`.
+
+"Compatible" means that a `varname` from `varnames(model)` can be used to extract the
+corresponding value using `get`, e.g. `get(posterior_mean(model), varname)`.
+"""
+function posterior_mean end
+
+"""
+    rand_prior_true([rng::AbstractRNG, ]model::DynamicPPL.Model)
+
+Return a `NamedTuple` of realizations from the prior of `model` compatible with `varnames(model)`.
+"""
+function rand_prior_true(model::DynamicPPL.Model)
+    return rand_prior_true(Random.default_rng(), model)
+end