Add @addlogprior! and @addloglikelihood!

Author: mhauru

docs/src/api.md

@@ -160,10 +160,12 @@ returned(::Model)
 
 ## Utilities
 
-It is possible to manually increase (or decrease) the accumulated log likelihood from within a model function.
+It is possible to manually increase (or decrease) the accumulated log likelihood or prior from within a model function.
 
 ```@docs
 @addlogprob!
+@addloglikelihood!
+@addlogprior!
 ```
 
 Return values of the model function for a collection of samples can be obtained with [`returned(model, chain)`](@ref).

src/DynamicPPL.jl

@@ -127,6 +127,8 @@ export AbstractVarInfo,
     to_submodel,
     # Convenience macros
     @addlogprob!,
+    @addlogprior!,
+    @addloglikelihood!,
     @submodel,
     value_iterator_from_chain,
     check_model,

src/abstract_varinfo.jl

@@ -320,13 +320,18 @@ function accloglikelihood!!(vi::AbstractVarInfo, logp)
 end
 
 """
-    acclogp!!(vi::AbstractVarInfo, logp::NamedTuple)
+    acclogp!!(vi::AbstractVarInfo, logp::NamedTuple; ignore_missing_accumulator::Bool=false)
 
 Add to both the log prior and the log likelihood probabilities in `vi`.
 
 `logp` should have fields `logprior` and/or `loglikelihood`, and no other fields.
+
+By default if the necessary accumulators are not in `vi` an error is thrown. If
+`ignore_missing_accumulator` is set to `true` then this is silently ignored instead.
 """
-function acclogp!!(vi::AbstractVarInfo, logp::NamedTuple{names}) where {names}
+function acclogp!!(
+    vi::AbstractVarInfo, logp::NamedTuple{names}; ignore_missing_accumulator=false
+) where {names}
     if !(
         names == (:logprior, :loglikelihood) ||
         names == (:loglikelihood, :logprior) ||
@@ -335,17 +340,19 @@ function acclogp!!(vi::AbstractVarInfo, logp::NamedTuple{names}) where {names}
     )
         error("logp must have fields logprior and/or loglikelihood and no other fields.")
     end
-    if haskey(logp, :logprior)
+    if haskey(logp, :logprior) &&
+        (!ignore_missing_accumulator || hasacc(vi, Val(:LogPrior)))
         vi = acclogprior!!(vi, logp.logprior)
     end
-    if haskey(logp, :loglikelihood)
+    if haskey(logp, :loglikelihood) &&
+        (!ignore_missing_accumulator || hasacc(vi, Val(:LogLikelihood)))
         vi = accloglikelihood!!(vi, logp.loglikelihood)
     end
     return vi
 end
 
 function acclogp!!(vi::AbstractVarInfo, logp::Number)
-    depwarn(
+    Base.depwarn(
         "`acclogp!!(vi::AbstractVarInfo, logp::Number)` is deprecated. Use `accloglikelihood!!(vi, logp)` instead.",
         :acclogp,
     )

src/utils.jl

@@ -18,8 +18,86 @@ const LogProbType = float(Real)
 """
     @addlogprob!(ex)
 
+Add a term to the log joint.
+
+If `ex` evaluates to a `NamedTuple` with keys `:loglikelihood` and/or `:logprior`, the
+values are added to the log likelihood and log prior respectively.
+
+If `ex` evaluates to a number it is added to the log likelihood. This use is deprecated
+and should be replaced with either the `NamedTuple` version or calls to
+[`@addloglikelihood!`](@ref).
+
+See also [`@addloglikelihood!`](@ref), [`@addlogprior!`](@ref).
+
+# Examples
+
+```jldoctest; setup = :(using Distributions)
+julia> mylogjoint(x, μ) = (; loglikelihood=loglikelihood(Normal(μ, 1), x), logprior=1.0);
+
+julia> @model function demo(x)
+           μ ~ Normal()
+           @addlogprob! mylogjoint(x, μ)
+       end;
+
+julia> x = [1.3, -2.1];
+
+julia> loglikelihood(demo(x), (μ=0.2,)) ≈ mylogjoint(x, 0.2).loglikelihood
+true
+
+julia> logprior(demo(x), (μ=0.2,)) ≈ logpdf(Normal(), 0.2) + mylogjoint(x, 0.2).logprior
+true
+```
+
+and to [reject samples](https://github.com/TuringLang/Turing.jl/issues/1328):
+
+```jldoctest; setup = :(using Distributions, LinearAlgebra)
+julia> @model function demo(x)
+           m ~ MvNormal(zero(x), I)
+           if dot(m, x) < 0
+               @addlogprob! (; loglikelihood=-Inf)
+               # Exit the model evaluation early
+               return
+           end
+           x ~ MvNormal(m, I)
+           return
+       end;
+
+julia> logjoint(demo([-2.1]), (m=[0.2],)) == -Inf
+true
+```
+"""
+macro addlogprob!(ex)
+    return quote
+        val = $(esc(ex))
+        vi = $(esc(:(__varinfo__)))
+        if val isa Number
+            Base.depwarn(
+                """
+       @addlogprob! with a single number argument is deprecated. Please use
+       @addlogprob! (; loglikelihood=x) or @addloglikelihood! instead.
+       """,
+                :addlogprob!,
+            )
+            if hasacc(vi, Val(:LogLikelihood))
+                $(esc(:(__varinfo__))) = accloglikelihood!!($(esc(:(__varinfo__))), val)
+            end
+        elseif !isa(val, NamedTuple)
+            error("logp must be a NamedTuple.")
+        else
+            $(esc(:(__varinfo__))) = acclogp!!(
+                $(esc(:(__varinfo__))), val; ignore_missing_accumulator=true
+            )
+        end
+    end
+end
+
+"""
+    @addloglikelihood!(ex)
+
 Add the result of the evaluation of `ex` to the log likelihood.
 
+See also [`@addlogprob!`](@ref), [`@addlogprior!`](@ref).
+
 # Examples
 
 This macro allows you to [include arbitrary terms in the likelihood](https://github.com/TuringLang/Turing.jl/issues/1332)
@@ -29,7 +107,7 @@ julia> myloglikelihood(x, μ) = loglikelihood(Normal(μ, 1), x);
 
 julia> @model function demo(x)
            μ ~ Normal()
-           @addlogprob! myloglikelihood(x, μ)
+           @addloglikelihood! myloglikelihood(x, μ)
        end;
 
 julia> x = [1.3, -2.1];
@@ -44,7 +122,7 @@ and to [reject samples](https://github.com/TuringLang/Turing.jl/issues/1328):
 julia> @model function demo(x)
            m ~ MvNormal(zero(x), I)
            if dot(m, x) < 0
-               @addlogprob! -Inf
+               @addloglikelihood! -Inf
                # Exit the model evaluation early
                return
            end
@@ -56,10 +134,43 @@ julia> logjoint(demo([-2.1]), (m=[0.2],)) == -Inf
 true
 ```
 """
-macro addlogprob!(ex)
+macro addloglikelihood!(ex)
+    return quote
+        if hasacc($(esc(:(__varinfo__))), Val(:LogLikelihood))
+            $(esc(:(__varinfo__))) = accloglikelihood!!($(esc(:(__varinfo__))), $(esc(ex)))
+        end
+    end
+end
+
+"""
+    @addlogprior!(ex)
+
+Add the result of the evaluation of `ex` to the log prior.
+
+See also [`@addloglikelihood!`](@ref), [`@addlogprob!`](@ref).
+
+# Examples
+
+This macro allows you to include arbitrary terms in the prior.
+
+```jldoctest; setup = :(using Distributions)
+julia> mylogpriorextraterm(μ) = μ > 0 ? -1.0 : 0.0;
+
+julia> @model function demo(x)
+           μ ~ Normal()
+           @addlogprior! mylogpriorextraterm(μ)
+       end;
+
+julia> x = [1.3, -2.1];
+
+julia> logprior(demo(x), (μ=0.2,)) ≈ logpdf(Normal(), 0.2) + mylogpriorextraterm(0.2)
+true
+```
+"""
+macro addlogprior!(ex)
     return quote
-        if $hasacc($(esc(:(__varinfo__))), Val(:LogLikelihood))
-            $(esc(:(__varinfo__))) = $accloglikelihood!!($(esc(:(__varinfo__))), $(esc(ex)))
+        if hasacc($(esc(:(__varinfo__))), Val(:LogPrior))
+            $(esc(:(__varinfo__))) = acclogprior!!($(esc(:(__varinfo__))), $(esc(ex)))
         end
     end
 end

test/utils.jl

@@ -6,10 +6,56 @@
             return global lp_after = getlogjoint(__varinfo__)
         end
 
-        model = testmodel()
-        varinfo = VarInfo(model)
+        varinfo = VarInfo(testmodel())
         @test iszero(lp_before)
         @test getlogjoint(varinfo) == lp_after == 42
+        @test getloglikelihood(varinfo) == 42
+
+        @model function testmodel_nt()
+            global lp_before = getlogjoint(__varinfo__)
+            @addlogprob! (; logprior=(pi + 1), loglikelihood=42)
+            return global lp_after = getlogjoint(__varinfo__)
+        end
+
+        varinfo = VarInfo(testmodel_nt())
+        @test iszero(lp_before)
+        @test getlogjoint(varinfo) == lp_after == 42 + 1 + pi
+        @test getloglikelihood(varinfo) == 42
+        @test getlogprior(varinfo) == pi + 1
+
+        @model function testmodel_nt2()
+            global lp_before = getlogjoint(__varinfo__)
+            llh_nt = (; loglikelihood=42)
+            @addlogprob! llh_nt
+            return global lp_after = getlogjoint(__varinfo__)
+        end
+
+        varinfo = VarInfo(testmodel_nt2())
+        @test iszero(lp_before)
+        @test getlogjoint(varinfo) == lp_after == 42
+        @test getloglikelihood(varinfo) == 42
+
+        @model function testmodel_likelihood()
+            global lp_before = getlogjoint(__varinfo__)
+            @addloglikelihood! 42
+            return global lp_after = getlogjoint(__varinfo__)
+        end
+
+        varinfo = VarInfo(testmodel_likelihood())
+        @test iszero(lp_before)
+        @test getlogjoint(varinfo) == lp_after == 42
+        @test getloglikelihood(varinfo) == 42
+
+        @model function testmodel_prior()
+            global lp_before = getlogjoint(__varinfo__)
+            @addlogprior! 42
+            return global lp_after = getlogjoint(__varinfo__)
+        end
+
+        varinfo = VarInfo(testmodel_prior())
+        @test iszero(lp_before)
+        @test getlogjoint(varinfo) == lp_after == 42
+        @test getlogprior(varinfo) == 42
     end
 
     @testset "getargs_dottilde" begin