resetaccs!! before evaluating (#1013)

* resetaccs

* Add a test, plus a bunch of == methods

* Fix tests

* Fix test (properly)

Author: penelopeysm

docs/src/api.md

@@ -380,7 +380,6 @@ acclogprior!!
 getloglikelihood
 setloglikelihood!!
 accloglikelihood!!
-resetlogp!!
 ```
 
 #### Variables and their realizations

src/DynamicPPL.jl

@@ -70,7 +70,6 @@ export AbstractVarInfo,
     acclogjac!!,
     acclogprior!!,
     accloglikelihood!!,
-    resetlogp!!,
     is_flagged,
     set_flag!,
     unset_flag!,

src/abstract_varinfo.jl

@@ -330,6 +330,15 @@ function map_accumulators!!(func::Function, vi::AbstractVarInfo)
     return setaccs!!(vi, map(func, getaccs(vi)))
 end
 
+"""
+    resetaccs!!(vi::AbstractVarInfo)
+
+Reset the values of all accumulators, using [`reset`](@ref).
+"""
+function resetaccs!!(vi::AbstractVarInfo)
+    return setaccs!!(vi, map(reset, getaccs(vi)))
+end
+
 """
     map_accumulator!!(func::Function, vi::AbstractVarInfo, ::Val{accname}) where {accname}
 
@@ -423,24 +432,6 @@ function acclogp!!(vi::AbstractVarInfo, logp::Number)
     return accloglikelihood!!(vi, logp)
 end
 
-"""
-    resetlogp!!(vi::AbstractVarInfo)
-
-Reset the values of the log probabilities (prior and likelihood) in `vi` to zero.
-"""
-function resetlogp!!(vi::AbstractVarInfo)
-    if hasacc(vi, Val(:LogPrior))
-        vi = map_accumulator!!(zero, vi, Val(:LogPrior))
-    end
-    if hasacc(vi, Val(:LogJacobian))
-        vi = map_accumulator!!(zero, vi, Val(:LogJacobian))
-    end
-    if hasacc(vi, Val(:LogLikelihood))
-        vi = map_accumulator!!(zero, vi, Val(:LogLikelihood))
-    end
-    return vi
-end
-
 # Variables and their realizations.
 @doc """
     keys(vi::AbstractVarInfo)
@@ -491,7 +482,7 @@ function getindex_internal end
 @doc """
     empty!!(vi::AbstractVarInfo)
 
-Empty `vi` of variables and reset any `logp` accumulators zeros.
+Empty `vi` of variables and reset all accumulators.
 
 This is useful when using a sampling algorithm that assumes an empty `vi`, e.g. `SMC`.
 """ BangBang.empty!!

src/accumulators.jl

@@ -13,6 +13,7 @@ An accumulator type `T <: AbstractAccumulator` must implement the following meth
 - `accumulator_name(acc::T)` or `accumulator_name(::Type{T})`
 - `accumulate_observe!!(acc::T, dist, val, vn)`
 - `accumulate_assume!!(acc::T, val, logjac, vn, dist)`
+- `reset(acc::T)`
 - `Base.copy(acc::T)`
 
 In these functions:
@@ -50,9 +51,7 @@ accumulator_name(acc::AbstractAccumulator) = accumulator_name(typeof(acc))
 
 Update `acc` in a `tilde_observe!!` call. Returns the updated `acc`.
 
-`vn` is the name of the variable being observed, `left` is the value of the variable, and
-`right` is the distribution on the RHS of the tilde statement. `vn` is `nothing` in the case
-of literal observations like `0.0 ~ Normal()`.
+See [`AbstractAccumulator`](@ref) for the meaning of the arguments.
 
 `accumulate_observe!!` may mutate `acc`, but not any of the other arguments.
 
@@ -65,26 +64,45 @@ function accumulate_observe!! end
 
 Update `acc` in a `tilde_assume!!` call. Returns the updated `acc`.
 
-`vn` is the name of the variable being assumed, `val` is the value of the variable (in the
-original, unlinked space), and `right` is the distribution on the RHS of the tilde
-statement. `logjac` is the log determinant of the Jacobian of the transformation that was
-done to convert the value of `vn` as it was given to `val`: for example, if the sampler is
-operating in linked (Euclidean) space, then logjac will be nonzero.
+See [`AbstractAccumulator`](@ref) for the meaning of the arguments.
 
 `accumulate_assume!!` may mutate `acc`, but not any of the other arguments.
 
 See also: [`accumulate_observe!!`](@ref)
 """
 function accumulate_assume!! end
 
+"""
+    reset(acc::AbstractAccumulator)
+
+Return a new accumulator like `acc`, but with its contents reset to the state that they
+should be at the beginning of model evaluation.
+
+Note that this may in general have very similar behaviour to [`split`](@ref), and may share
+the same implementation, but the difference is that `split` may in principle happen at any
+stage during model evaluation, whereas `reset` is only called at the beginning of model
+evaluation.
+"""
+function reset end
+
+@doc """
+    Base.copy(acc::AbstractAccumulator)
+
+Create a new accumulator that is a copy of `acc`, without aliasing (i.e., this should
+behave conceptually like a `deepcopy`).
+""" Base.copy
+
 """
     split(acc::AbstractAccumulator)
 
-Return a new accumulator like `acc` but empty.
+Return a new accumulator like `acc` suitable for use in a forked thread.
+
+The returned value should be such that `combine(acc, split(acc))` is equal to `acc`. This is
+used in the context of multi-threading where different threads may accumulate independently
+and the results are then combined.
 
-The precise meaning of "empty" is that that the returned value should be such that
-`combine(acc, split(acc))` is equal to `acc`. This is used in the context of multi-threading
-where different threads may accumulate independently and the results are then combined.
+Note that this may in general have very similar behaviour to [`reset`](@ref), but is
+semantically different. See [`reset`](@ref) for more details.
 
 See also: [`combine`](@ref)
 """

src/debug_utils.jl

@@ -156,12 +156,22 @@ end
 const _DEBUG_ACC_NAME = :Debug
 DynamicPPL.accumulator_name(::Type{<:DebugAccumulator}) = _DEBUG_ACC_NAME
 
-function split(acc::DebugAccumulator)
+function Base.:(==)(acc1::DebugAccumulator, acc2::DebugAccumulator)
+    return (
+        acc1.varnames_seen == acc2.varnames_seen &&
+        acc1.statements == acc2.statements &&
+        acc1.error_on_failure == acc2.error_on_failure
+    )
+end
+
+function _zero(acc::DebugAccumulator)
     return DebugAccumulator(
         OrderedDict{VarName,Int}(), Vector{Stmt}(), acc.error_on_failure
     )
 end
-function combine(acc1::DebugAccumulator, acc2::DebugAccumulator)
+DynamicPPL.reset(acc::DebugAccumulator) = _zero(acc)
+DynamicPPL.split(acc::DebugAccumulator) = _zero(acc)
+function DynamicPPL.combine(acc1::DebugAccumulator, acc2::DebugAccumulator)
     return DebugAccumulator(
         merge(acc1.varnames_seen, acc2.varnames_seen),
         vcat(acc1.statements, acc2.statements),
@@ -416,7 +426,7 @@ function check_model_and_trace(
     issuccess = check_model_pre_evaluation(model)
 
     # Force single-threaded execution.
-    DynamicPPL.evaluate_threadunsafe!!(model, varinfo)
+    _, varinfo = DynamicPPL.evaluate_threadunsafe!!(model, varinfo)
 
     # Perform checks after evaluating the model.
     debug_acc = DynamicPPL.getacc(varinfo, Val(_DEBUG_ACC_NAME))

src/default_accumulators.jl

@@ -47,8 +47,9 @@ end
 
 Base.hash(acc::T, h::UInt) where {T<:LogProbAccumulator} = hash((T, logp(acc)), h)
 
-split(::AccType) where {T,AccType<:LogProbAccumulator{T}} = AccType(zero(T))
-
+_zero(::Tacc) where {Tlogp,Tacc<:LogProbAccumulator{Tlogp}} = Tacc(zero(Tlogp))
+reset(acc::LogProbAccumulator) = _zero(acc)
+split(acc::LogProbAccumulator) = _zero(acc)
 function combine(acc::LogProbAccumulator, acc2::LogProbAccumulator)
     if basetypeof(acc) !== basetypeof(acc2)
         msg = "Cannot combine accumulators of different types: $(basetypeof(acc)) and $(basetypeof(acc2))"
@@ -59,8 +60,6 @@ end
 
 acclogp(acc::LogProbAccumulator, val) = basetypeof(acc)(logp(acc) + val)
 
-Base.zero(acc::T) where {T<:LogProbAccumulator} = T(zero(logp(acc)))
-
 function Base.convert(
     ::Type{AccType}, acc::LogProbAccumulator
 ) where {T,AccType<:LogProbAccumulator{T}}

src/extract_priors.jl

@@ -10,7 +10,13 @@ end
 
 accumulator_name(::PriorDistributionAccumulator) = :PriorDistributionAccumulator
 
-split(acc::PriorDistributionAccumulator) = PriorDistributionAccumulator(empty(acc.priors))
+function Base.:(==)(acc1::PriorDistributionAccumulator, acc2::PriorDistributionAccumulator)
+    return acc1.priors == acc2.priors
+end
+
+_zero(acc::PriorDistributionAccumulator) = PriorDistributionAccumulator(empty(acc.priors))
+reset(acc::PriorDistributionAccumulator) = _zero(acc)
+split(acc::PriorDistributionAccumulator) = _zero(acc)
 function combine(acc1::PriorDistributionAccumulator, acc2::PriorDistributionAccumulator)
     return PriorDistributionAccumulator(merge(acc1.priors, acc2.priors))
 end

src/model.jl

@@ -884,7 +884,7 @@ This method is not exposed and supposed to be used only internally in DynamicPPL
 See also: [`evaluate_threadsafe!!`](@ref)
 """
 function evaluate_threadunsafe!!(model, varinfo)
-    return _evaluate!!(model, resetlogp!!(varinfo))
+    return _evaluate!!(model, resetaccs!!(varinfo))
 end
 
 """
@@ -899,7 +899,7 @@ This method is not exposed and supposed to be used only internally in DynamicPPL
 See also: [`evaluate_threadunsafe!!`](@ref)
 """
 function evaluate_threadsafe!!(model, varinfo)
-    wrapper = ThreadSafeVarInfo(resetlogp!!(varinfo))
+    wrapper = ThreadSafeVarInfo(resetaccs!!(varinfo))
     result, wrapper_new = _evaluate!!(model, wrapper)
     # TODO(penelopeysm): If seems that if you pass a TSVI to this method, it
     # will return the underlying VI, which is a bit counterintuitive (because

src/pointwise_logdensities.jl

@@ -32,6 +32,12 @@ function PointwiseLogProbAccumulator{whichlogprob,KeyType}() where {whichlogprob
     return PointwiseLogProbAccumulator{whichlogprob,KeyType,typeof(logps)}(logps)
 end
 
+function Base.:(==)(
+    acc1::PointwiseLogProbAccumulator{wlp1}, acc2::PointwiseLogProbAccumulator{wlp2}
+) where {wlp1,wlp2}
+    return (wlp1 == wlp2 && acc1.logps == acc2.logps)
+end
+
 function Base.copy(acc::PointwiseLogProbAccumulator{whichlogprob}) where {whichlogprob}
     return PointwiseLogProbAccumulator{whichlogprob}(copy(acc.logps))
 end
@@ -56,10 +62,11 @@ function accumulator_name(
     return Symbol("PointwiseLogProbAccumulator{$whichlogprob}")
 end
 
-function split(acc::PointwiseLogProbAccumulator{whichlogprob}) where {whichlogprob}
+function _zero(acc::PointwiseLogProbAccumulator{whichlogprob}) where {whichlogprob}
     return PointwiseLogProbAccumulator{whichlogprob}(empty(acc.logps))
 end
-
+reset(acc::PointwiseLogProbAccumulator) = _zero(acc)
+split(acc::PointwiseLogProbAccumulator) = _zero(acc)
 function combine(
     acc::PointwiseLogProbAccumulator{whichlogprob},
     acc2::PointwiseLogProbAccumulator{whichlogprob},
@@ -223,23 +230,37 @@ function pointwise_logdensities(
     # Get the data by executing the model once
     vi = VarInfo(model)
 
+    # This accumulator tracks the pointwise log-probabilities in a single iteration.
     AccType = PointwiseLogProbAccumulator{whichlogprob,KeyType}
     vi = setaccs!!(vi, (AccType(),))
 
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
+
+    # Maintain a separate accumulator that isn't tied to a VarInfo but rather
+    # tracks _all_ iterations.
+    all_logps = AccType()
     for (sample_idx, chain_idx) in iters
         # Update the values
         setval!(vi, chain, sample_idx, chain_idx)
 
         # Execute model
         vi = last(evaluate!!(model, vi))
+
+        # Get the log-probabilities
+        this_iter_logps = getacc(vi, Val(accumulator_name(AccType))).logps
+
+        # Merge into main acc
+        for (varname, this_lp) in this_iter_logps
+            # Because `this_lp` is obtained from one model execution, it should only
+            # contain one variable, hence `only()`.
+            push!(all_logps, varname, only(this_lp))
+        end
     end
 
-    logps = getacc(vi, Val(accumulator_name(AccType))).logps
     niters = size(chain, 1)
     nchains = size(chain, 3)
     logdensities = OrderedDict(
-        varname => reshape(vals, niters, nchains) for (varname, vals) in logps
+        varname => reshape(vals, niters, nchains) for (varname, vals) in all_logps.logps
     )
     return logdensities
 end

src/simple_varinfo.jl

@@ -287,7 +287,7 @@ function unflatten(svi::SimpleVarInfo, x::AbstractVector)
 end
 
 function BangBang.empty!!(vi::SimpleVarInfo)
-    return resetlogp!!(Accessors.@set vi.values = empty!!(vi.values))
+    return resetaccs!!(Accessors.@set vi.values = empty!!(vi.values))
 end
 Base.isempty(vi::SimpleVarInfo) = isempty(vi.values)