Improve acc tests

Author: mhauru

src/DynamicPPL.jl

@@ -66,6 +66,7 @@ export AbstractVarInfo,
     setlogprior!!,
     setlogjac!!,
     setloglikelihood!!,
+    acclogp,
     acclogp!!,
     acclogjac!!,
     acclogprior!!,

src/accumulators.jl

@@ -203,11 +203,9 @@ Implemented as a generated function to enabled constant propagation of the resul
 @generated function _joint_keys(
     nt1::NamedTuple{names1}, nt2::NamedTuple{names2}
 ) where {names1,names2}
-    set_names1 = Set(names1)
-    set_names2 = Set(names2)
-    only_in_nt1 = tuple(setdiff(set_names1, set_names2)...)
-    only_in_nt2 = tuple(setdiff(set_names2, set_names1)...)
-    in_both = tuple(intersect(set_names1, set_names2)...)
+    only_in_nt1 = tuple(setdiff(names1, names2)...)
+    only_in_nt2 = tuple(setdiff(names2, names1)...)
+    in_both = tuple(intersect(names1, names2)...)
     return :($only_in_nt1, $only_in_nt2, $in_both)
 end
 
@@ -226,7 +224,7 @@ function Base.merge(at1::AccumulatorTuple, at2::AccumulatorTuple)
     accs_in_both = (
         merge(getfield(at1.nt, key), getfield(at2.nt, key)) for key in keys_in_both
     )
-    return AccumulatorTuple(accs_in_at1..., accs_in_at2..., accs_in_both...)
+    return AccumulatorTuple(accs_in_at1..., accs_in_both..., accs_in_at2...)
 end
 
 """

test/accumulators.jl

@@ -109,13 +109,13 @@ using DynamicPPL:
                 VariableOrderAccumulator(2)
         end
 
-        @testset "merge and subset" begin
+        @testset "merge" begin
             @test merge(LogPriorAccumulator(1.0), LogPriorAccumulator(2.0)) ==
-                LogPriorAccumulator(3.0)
+                LogPriorAccumulator(2.0)
             @test merge(LogJacobianAccumulator(1.0), LogJacobianAccumulator(2.0)) ==
-                LogJacobianAccumulator(3.0)
+                LogJacobianAccumulator(2.0)
             @test merge(LogLikelihoodAccumulator(1.0), LogLikelihoodAccumulator(2.0)) ==
-                LogLikelihoodAccumulator(3.0)
+                LogLikelihoodAccumulator(2.0)
 
             @test merge(
                 VariableOrderAccumulator(1, Dict{VarName,Int}()),
@@ -132,6 +132,49 @@ using DynamicPPL:
                 1, Dict{VarName,Int}((@varname(a) => 2, @varname(b) => 2, @varname(c) => 3))
             )
         end
+
+        @testset "subset" begin
+            @test subset(LogPriorAccumulator(1.0), VarName[]) == LogPriorAccumulator(1.0)
+            @test subset(LogJacobianAccumulator(1.0), VarName[]) ==
+                LogJacobianAccumulator(1.0)
+            @test subset(LogLikelihoodAccumulator(1.0), VarName[]) ==
+                LogLikelihoodAccumulator(1.0)
+
+            @test subset(
+                VariableOrderAccumulator(1, Dict{VarName,Int}()),
+                VarName[@varname(a), @varname(b)],
+            ) == VariableOrderAccumulator(1, Dict{VarName,Int}())
+            @test subset(
+                VariableOrderAccumulator(
+                    2, Dict{VarName,Int}((@varname(a) => 1, @varname(b) => 2))
+                ),
+                VarName[@varname(a)],
+            ) == VariableOrderAccumulator(2, Dict{VarName,Int}((@varname(a) => 1)))
+            @test subset(
+                VariableOrderAccumulator(
+                    2, Dict{VarName,Int}((@varname(a) => 1, @varname(b) => 2))
+                ),
+                VarName[],
+            ) == VariableOrderAccumulator(2, Dict{VarName,Int}())
+            @test subset(
+                VariableOrderAccumulator(
+                    2,
+                    Dict{VarName,Int}((
+                        @varname(a) => 1,
+                        @varname(a.b.c) => 2,
+                        @varname(a.b.c.d[1]) => 2,
+                        @varname(b) => 3,
+                        @varname(c[1]) => 4,
+                    )),
+                ),
+                VarName[@varname(a.b), @varname(b)],
+            ) == VariableOrderAccumulator(
+                2,
+                Dict{VarName,Int}((
+                    @varname(a.b.c) => 2, @varname(a.b.c.d[1]) => 2, @varname(b) => 3
+                )),
+            )
+        end
     end
 
     @testset "accumulator tuples" begin
@@ -140,7 +183,7 @@ using DynamicPPL:
         lp_f32 = LogPriorAccumulator(1.0f0)
         ll_f64 = LogLikelihoodAccumulator(1.0)
         ll_f32 = LogLikelihoodAccumulator(1.0f0)
-        np_i64 = VariableOrderAccumulator(1)
+        vo_i64 = VariableOrderAccumulator(1)
 
         @testset "constructors" begin
             @test AccumulatorTuple(lp_f64, ll_f64) == AccumulatorTuple((lp_f64, ll_f64))
@@ -154,22 +197,22 @@ using DynamicPPL:
         end
 
         @testset "basic operations" begin
-            at_all64 = AccumulatorTuple(lp_f64, ll_f64, np_i64)
+            at_all64 = AccumulatorTuple(lp_f64, ll_f64, vo_i64)
 
             @test at_all64[:LogPrior] == lp_f64
             @test at_all64[:LogLikelihood] == ll_f64
-            @test at_all64[:VariableOrder] == np_i64
+            @test at_all64[:VariableOrder] == vo_i64
 
-            @test haskey(AccumulatorTuple(np_i64), Val(:VariableOrder))
-            @test ~haskey(AccumulatorTuple(np_i64), Val(:LogPrior))
-            @test length(AccumulatorTuple(lp_f64, ll_f64, np_i64)) == 3
+            @test haskey(AccumulatorTuple(vo_i64), Val(:VariableOrder))
+            @test ~haskey(AccumulatorTuple(vo_i64), Val(:LogPrior))
+            @test length(AccumulatorTuple(lp_f64, ll_f64, vo_i64)) == 3
             @test keys(at_all64) == (:LogPrior, :LogLikelihood, :VariableOrder)
-            @test collect(at_all64) == [lp_f64, ll_f64, np_i64]
+            @test collect(at_all64) == [lp_f64, ll_f64, vo_i64]
 
             # Replace the existing LogPriorAccumulator
             @test setacc!!(at_all64, lp_f32)[:LogPrior] == lp_f32
             # Check that setacc!! didn't modify the original
-            @test at_all64 == AccumulatorTuple(lp_f64, ll_f64, np_i64)
+            @test at_all64 == AccumulatorTuple(lp_f64, ll_f64, vo_i64)
             # Add a new accumulator type.
             @test setacc!!(AccumulatorTuple(lp_f64), ll_f64) ==
                 AccumulatorTuple(lp_f64, ll_f64)
@@ -197,6 +240,52 @@ using DynamicPPL:
                 acc -> convert_eltype(Float64, acc), accs, Val(:LogLikelihood)
             ) == AccumulatorTuple(lp_f32, LogLikelihoodAccumulator(1.0))
         end
+
+        @testset "merge" begin
+            vo1 = VariableOrderAccumulator(
+                1, Dict{VarName,Int}(@varname(a) => 1, @varname(b) => 1)
+            )
+            vo2 = VariableOrderAccumulator(
+                2, Dict{VarName,Int}(@varname(a) => 2, @varname(c) => 2)
+            )
+            accs1 = AccumulatorTuple(lp_f64, ll_f64, vo1)
+            accs2 = AccumulatorTuple(lp_f32, vo2)
+            @test merge(accs1, accs2) == AccumulatorTuple(
+                ll_f64,
+                lp_f32,
+                VariableOrderAccumulator(
+                    2,
+                    Dict{VarName,Int}(@varname(a) => 2, @varname(b) => 1, @varname(c) => 2),
+                ),
+            )
+            @test merge(AccumulatorTuple(), accs1) == accs1
+            @test merge(accs1, AccumulatorTuple()) == accs1
+            @test merge(accs1, accs1) == accs1
+        end
+
+        @testset "subset" begin
+            accs = AccumulatorTuple(
+                lp_f64,
+                ll_f64,
+                VariableOrderAccumulator(
+                    1,
+                    Dict{VarName,Int}(
+                        @varname(a.b) => 1, @varname(a.b[1]) => 2, @varname(b) => 1
+                    ),
+                ),
+            )
+
+            @test subset(accs, VarName[]) == AccumulatorTuple(
+                lp_f64, ll_f64, VariableOrderAccumulator(1, Dict{VarName,Int}())
+            )
+            @test subset(accs, VarName[@varname(a)]) == AccumulatorTuple(
+                lp_f64,
+                ll_f64,
+                VariableOrderAccumulator(
+                    1, Dict{VarName,Int}(@varname(a.b) => 1, @varname(a.b[1]) => 2)
+                ),
+            )
+        end
     end
 end