Add NamedTuple methods for get/set/acclogp

Author: mhauru

HISTORY.md

@@ -15,6 +15,8 @@ This release overhauls how VarInfo objects track variables such as the log joint
   - For literal observation statements like `0.0 ~ Normal(blahblah)` we used to call `tilde_observe!!` without the `vn` argument. This method no longer exists. Rather we call `tilde_observe!!` with `vn` set to `nothing`.
   - `set/reset/increment_num_produce!` have become `set/reset/increment_num_produce!!` (note the second exclamation mark). They are no longer guaranteed to modify the `VarInfo` in place, and one should always use the return value.
   - `@addlogprob!` now _always_ adds to the log likelihood. Previously it added to the log probability that the execution context specified, e.g. the log prior when using `PriorContext`.
+  - `getlogp` now returns a `NamedTuple` with keys `logprior` and `loglikelihood`. If you want the log joint probability, which is what `getlogp` used to return, use `getlogjoint`.
+  - Correspondingly `setlogp!!` and `acclogp!!` should now be called with a `NamedTuple` with keys `logprior` and `loglikelihood`. The method with a single scalar value has been deprecated, and falls back on `setloglikelihood!!` or `accloglikelihood!!`. Corresponding setter/accumulator functions exist for the log prior as well.
 
 ## 0.36.0
 

docs/src/api.md

@@ -366,12 +366,15 @@ NumProduce
 #### Accumulation of log-probabilities
 
 ```@docs
-getlogprior
-getloglikelihood
+getlogp
+setlogp!!
+acclogp!!
 getlogjoint
+getlogprior
 setlogprior!!
-setloglikelihood!!
 acclogprior!!
+getloglikelihood
+setloglikelihood!!
 accloglikelihood!!
 resetlogp!!
 ```

src/abstract_varinfo.jl

@@ -98,9 +98,17 @@ Return the log of the joint probability of the observed data and parameters in `
 See also: [`getlogprior`](@ref), [`getloglikelihood`](@ref).
 """
 getlogjoint(vi::AbstractVarInfo) = getlogprior(vi) + getloglikelihood(vi)
+
+"""
+    getlogp(vi::AbstractVarInfo)
+
+Return a NamedTuple of the log prior and log likelihood probabilities.
+
+The keys are called `logprior` and `loglikelihood`. If either one is not present in `vi` an
+error will be thrown.
+"""
 function getlogp(vi::AbstractVarInfo)
-    Base.depwarn("getlogp is deprecated, use getlogjoint instead", :getlogp)
-    return getlogjoint(vi)
+    return (; logprior=getlogprior(vi), loglikelihood=getloglikelihood(vi))
 end
 
 """
@@ -198,23 +206,31 @@ See also: [`setlogprior!!`](@ref), [`setlogp!!`](@ref), [`getloglikelihood`](@re
 setloglikelihood!!(vi::AbstractVarInfo, logp) = setacc!!(vi, LogLikelihood(logp))
 
 """
-    setlogp!!(vi::AbstractVarInfo, logp)
+    setlogp!!(vi::AbstractVarInfo, logp::NamedTuple)
 
-Set the log of the joint probability of the observed data and parameters sampled in
-`vi` to `logp`, mutating if it makes sense.
+Set both the log prior and the log likelihood probabilities in `vi`.
+
+`logp` should have fields `logprior` and `loglikelihood` and no other fields.
 
 See also: [`setlogprior!!`](@ref), [`setloglikelihood!!`](@ref), [`getlogp`](@ref).
 """
-function setlogp!!(vi::AbstractVarInfo, logp)
-    Base.depwarn(
-        "setlogp!! is deprecated, use setlogprior!! or setloglikelihood!! instead",
-        :setlogp!!,
-    )
-    vi = setlogprior!!(vi, zero(logp))
-    vi = setloglikelihood!!(vi, logp)
+function setlogp!!(vi::AbstractVarInfo, logp::NamedTuple{names}) where {names}
+    if !(names == (:logprior, :loglikelihood) || names == (:loglikelihood, :logprior))
+        error("logp must have the fields logprior and loglikelihood and no other fields.")
+    end
+    vi = setlogprior!!(vi, logp.logprior)
+    vi = setloglikelihood!!(vi, logp.loglikelihood)
     return vi
 end
 
+function setlogp!!(vi::AbstractVarInfo, logp::Number)
+    depwarn(
+        "`setlogp!!(vi::AbstractVarInfo, logp::Number)` is deprecated. Use `setloglikelihood!!(vi, logp)` instead.",
+        :setlogp,
+    )
+    return setloglikelihood!!(vi, logp)
+end
+
 """
     getacc(vi::AbstractVarInfo, ::Val{accname})
 
@@ -303,15 +319,34 @@ function accloglikelihood!!(vi::AbstractVarInfo, logp)
 end
 
 """
-    acclogp!!(vi::AbstractVarInfo, logp)
+    acclogp!!(vi::AbstractVarInfo, logp::NamedTuple)
+
+Add to both the log prior and the log likelihood probabilities in `vi`.
 
-Add `logp` to the value of the log of the joint probability of the observed data and
-parameters sampled in `vi`, mutating if it makes sense.
+`logp` should have fields `logprior` and/or `loglikelihood`, and no other fields.
 """
-function acclogp!!(vi::AbstractVarInfo, logp)
-    Base.depwarn(
-        "acclogp!! is deprecated, use acclogprior!! or accloglikelihood!! instead",
-        :acclogp!!,
+function acclogp!!(vi::AbstractVarInfo, logp::NamedTuple{names}) where {names}
+    if !(
+        names == (:logprior, :loglikelihood) ||
+        names == (:loglikelihood, :logprior) ||
+        names == (:logprior,) ||
+        names == (:loglikelihood,)
+    )
+        error("logp must have fields logprior and/or loglikelihood and no other fields.")
+    end
+    if haskey(logp, :logprior)
+        vi = acclogprior!!(vi, logp.logprior)
+    end
+    if haskey(logp, :loglikelihood)
+        vi = accloglikelihood!!(vi, logp.loglikelihood)
+    end
+    return vi
+end
+
+function acclogp!!(vi::AbstractVarInfo, logp::Number)
+    depwarn(
+        "`acclogp!!(vi::AbstractVarInfo, logp::Number)` is deprecated. Use `accloglikelihood!!(vi, logp)` instead.",
+        :acclogp,
     )
     return accloglikelihood!!(vi, logp)
 end

test/submodels.jl

@@ -35,7 +35,7 @@ using Test
                 @test model()[1] == x_val
                 # Test that the logp was correctly set
                 vi = VarInfo(model)
-                @test getlogp(vi) == x_logp + logpdf(Normal(), vi[@varname(a.y)])
+                @test getlogjoint(vi) == x_logp + logpdf(Normal(), vi[@varname(a.y)])
                 # Check the keys
                 @test Set(keys(VarInfo(model))) == Set([@varname(a.y)])
             end
@@ -67,7 +67,7 @@ using Test
                 @test model()[1] == x_val
                 # Test that the logp was correctly set
                 vi = VarInfo(model)
-                @test getlogp(vi) == x_logp + logpdf(Normal(), vi[@varname(y)])
+                @test getlogjoint(vi) == x_logp + logpdf(Normal(), vi[@varname(y)])
                 # Check the keys
                 @test Set(keys(VarInfo(model))) == Set([@varname(y)])
             end
@@ -99,7 +99,7 @@ using Test
                 @test model()[1] == x_val
                 # Test that the logp was correctly set
                 vi = VarInfo(model)
-                @test getlogp(vi) == x_logp + logpdf(Normal(), vi[@varname(b.y)])
+                @test getlogjoint(vi) == x_logp + logpdf(Normal(), vi[@varname(b.y)])
                 # Check the keys
                 @test Set(keys(VarInfo(model))) == Set([@varname(b.y)])
             end
@@ -148,7 +148,7 @@ using Test
             # No conditioning
             vi = VarInfo(h())
             @test Set(keys(vi)) == Set([@varname(a.b.x), @varname(a.b.y)])
-            @test getlogp(vi) ==
+            @test getlogjoint(vi) ==
                 logpdf(Normal(), vi[@varname(a.b.x)]) +
                   logpdf(Normal(), vi[@varname(a.b.y)])
 
@@ -174,7 +174,7 @@ using Test
             @testset "$name" for (name, model) in models
                 vi = VarInfo(model)
                 @test Set(keys(vi)) == Set([@varname(a.b.y)])
-                @test getlogp(vi) == x_logp + logpdf(Normal(), vi[@varname(a.b.y)])
+                @test getlogjoint(vi) == x_logp + logpdf(Normal(), vi[@varname(a.b.y)])
             end
         end
     end

test/varinfo.jl

@@ -175,29 +175,68 @@ end
         vi = last(DynamicPPL.evaluate!!(m, deepcopy(vi)))
         @test getlogprior(vi) == lp_a + lp_b
         @test getloglikelihood(vi) == lp_c + lp_d
+        @test getlogp(vi) == (; logprior=lp_a + lp_b, loglikelihood=lp_c + lp_d)
+        @test getlogjoint(vi) == lp_a + lp_b + lp_c + lp_d
         @test get_num_produce(vi) == 2
+        @test begin
+            vi = acclogprior!!(vi, 1.0)
+            getlogprior(vi) == lp_a + lp_b + 1.0
+        end
+        @test begin
+            vi = accloglikelihood!!(vi, 1.0)
+            getloglikelihood(vi) == lp_c + lp_d + 1.0
+        end
+        @test begin
+            vi = setlogprior!!(vi, -1.0)
+            getlogprior(vi) == -1.0
+        end
+        @test begin
+            vi = setloglikelihood!!(vi, -1.0)
+            getloglikelihood(vi) == -1.0
+        end
+        @test begin
+            vi = setlogp!!(vi, (logprior=-3.0, loglikelihood=-3.0))
+            getlogp(vi) == (; logprior=-3.0, loglikelihood=-3.0)
+        end
+        @test begin
+            vi = acclogp!!(vi, (logprior=1.0, loglikelihood=1.0))
+            getlogp(vi) == (; logprior=-2.0, loglikelihood=-2.0)
+        end
+        @test getlogp(setlogp!!(vi, getlogp(vi))) == getlogp(vi)
 
         vi = last(
             DynamicPPL.evaluate!!(m, DynamicPPL.setaccs!!(deepcopy(vi), (LogPrior(),)))
         )
         @test getlogprior(vi) == lp_a + lp_b
         @test_throws "has no field LogLikelihood" getloglikelihood(vi)
+        @test_throws "has no field LogLikelihood" getlogp(vi)
         @test_throws "has no field LogLikelihood" getlogjoint(vi)
         @test_throws "has no field NumProduce" get_num_produce(vi)
+        @test begin
+            vi = acclogprior!!(vi, 1.0)
+            getlogprior(vi) == lp_a + lp_b + 1.0
+        end
+        @test begin
+            vi = setlogprior!!(vi, -1.0)
+            getlogprior(vi) == -1.0
+        end
+        @test_throws "has no field LogLikelihood" setlogp!!(getlogp(vi))
 
         vi = last(
             DynamicPPL.evaluate!!(m, DynamicPPL.setaccs!!(deepcopy(vi), (NumProduce(),)))
         )
         @test_throws "has no field LogPrior" getlogprior(vi)
-        @test_throws "has no field LogPrior" getlogjoint(vi)
         @test_throws "has no field LogLikelihood" getloglikelihood(vi)
+        @test_throws "has no field LogPrior" getlogp(vi)
+        @test_throws "has no field LogPrior" getlogjoint(vi)
         @test get_num_produce(vi) == 2
 
         # Test evaluating without any accumulators.
         vi = last(DynamicPPL.evaluate!!(m, DynamicPPL.setaccs!!(deepcopy(vi), ())))
         @test_throws "has no field LogPrior" getlogprior(vi)
-        @test_throws "has no field LogPrior" getlogjoint(vi)
         @test_throws "has no field LogLikelihood" getloglikelihood(vi)
+        @test_throws "has no field LogPrior" getlogp(vi)
+        @test_throws "has no field LogPrior" getlogjoint(vi)
         @test_throws "has no field NumProduce" get_num_produce(vi)
         @test_throws "has no field NumProduce" reset_num_produce!!(vi)
     end