Turn NumProduceAccumulator into VariableOrderAccumulator

Author: mhauru

docs/src/api.md

@@ -331,7 +331,7 @@ get_num_produce
 set_num_produce!!
 increment_num_produce!!
 reset_num_produce!!
-setorder!
+setorder!!
 set_retained_vns_del!
 ```
 
@@ -358,7 +358,7 @@ DynamicPPL provides the following default accumulators.
 ```@docs
 LogPriorAccumulator
 LogLikelihoodAccumulator
-NumProduceAccumulator
+VariableOrderAccumulator
 ```
 
 ### Common API

src/DynamicPPL.jl

@@ -50,7 +50,7 @@ export AbstractVarInfo,
     AbstractAccumulator,
     LogLikelihoodAccumulator,
     LogPriorAccumulator,
-    NumProduceAccumulator,
+    VariableOrderAccumulator,
     push!!,
     empty!!,
     subset,
@@ -73,7 +73,7 @@ export AbstractVarInfo,
     is_flagged,
     set_flag!,
     unset_flag!,
-    setorder!,
+    setorder!!,
     istrans,
     link,
     link!!,

src/abstract_varinfo.jl

@@ -374,6 +374,24 @@ function resetlogp!!(vi::AbstractVarInfo)
     return vi
 end
 
+"""
+    setorder!!(vi::AbstractVarInfo, vn::VarName, index::Integer)
+
+Set the `order` of `vn` in `vi` to `index`, where `order` is the number of `observe
+statements run before sampling `vn`.
+"""
+function setorder!!(vi::AbstractVarInfo, vn::VarName, index::Integer)
+    return map_accumulator!!(acc -> (acc.order[vn] = index; acc), vi, Val(:VariableOrder))
+end
+
+"""
+    getorder(vi::VarInfo, vn::VarName)
+
+Get the `order` of `vn` in `vi`, where `order` is the number of `observe` statements
+run before sampling `vn`.
+"""
+getorder(vi::AbstractVarInfo, vn::VarName) = getacc(vi, Val(:VariableOrder)).order[vn]
+
 # Variables and their realizations.
 @doc """
     keys(vi::AbstractVarInfo)
@@ -972,29 +990,37 @@ end
 
 Return the `num_produce` of `vi`.
 """
-get_num_produce(vi::AbstractVarInfo) = getacc(vi, Val(:NumProduce)).num
+get_num_produce(vi::AbstractVarInfo) = getacc(vi, Val(:VariableOrder)).num_produce
 
 """
     set_num_produce!!(vi::AbstractVarInfo, n::Int)
 
 Set the `num_produce` field of `vi` to `n`.
 """
-set_num_produce!!(vi::AbstractVarInfo, n::Int) = setacc!!(vi, NumProduceAccumulator(n))
+function set_num_produce!!(vi::AbstractVarInfo, n::Integer)
+    if hasacc(vi, Val(:VariableOrder))
+        acc = getacc(vi, Val(:VariableOrder))
+        acc = VariableOrderAccumulator(n, acc.order)
+    else
+        acc = VariableOrderAccumulator(n)
+    end
+    return setacc!!(vi, acc)
+end
 
 """
     increment_num_produce!!(vi::AbstractVarInfo)
 
 Add 1 to `num_produce` in `vi`.
 """
 increment_num_produce!!(vi::AbstractVarInfo) =
-    map_accumulator!!(increment, vi, Val(:NumProduce))
+    map_accumulator!!(increment, vi, Val(:VariableOrder))
 
 """
     reset_num_produce!!(vi::AbstractVarInfo)
 
 Reset the value of `num_produce` in `vi` to 0.
 """
-reset_num_produce!!(vi::AbstractVarInfo) = map_accumulator!!(zero, vi, Val(:NumProduce))
+reset_num_produce!!(vi::AbstractVarInfo) = set_num_produce!!(vi, zero(get_num_produce(vi)))
 
 """
     from_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])

src/context_implementations.jl

@@ -184,7 +184,6 @@ function assume(
             f = to_maybe_linked_internal_transform(vi, vn, dist)
             # TODO(mhauru) This should probably be call a function called setindex_internal!
             vi = BangBang.setindex!!(vi, f(r), vn)
-            setorder!(vi, vn, get_num_produce(vi))
         else
             # Otherwise we just extract it.
             r = vi[vn, dist]

src/default_accumulators.jl

@@ -41,52 +41,62 @@ LogLikelihoodAccumulator{T}() where {T<:Real} = LogLikelihoodAccumulator(zero(T)
 LogLikelihoodAccumulator() = LogLikelihoodAccumulator{LogProbType}()
 
 """
-    NumProduceAccumulator{T} <: AbstractAccumulator
+    VariableOrderAccumulator{T} <: AbstractAccumulator
 
-An accumulator that tracks the number of observations during model execution.
+An accumulator that tracks the order of variables in a `VarInfo`.
+
+This doesn't track the full ordering, but rather how many observations have taken place
+before the assume statement for each variable. This is needed for particle methods, where
+the model is segmented into parts by each observation, and we need to know which part each
+assume statement is in.
 
 # Fields
 $(TYPEDFIELDS)
 """
-struct NumProduceAccumulator{T<:Integer} <: AbstractAccumulator
+struct VariableOrderAccumulator{Eltype<:Integer,VNType<:VarName} <: AbstractAccumulator
     "the number of observations"
-    num::T
+    num_produce::Eltype
+    "mapping of variable names to their order in the model"
+    order::OrderedDict{VNType, Eltype}
 end
 
 """
-    NumProduceAccumulator{T<:Integer}()
+    VariableOrderAccumulator{T<:Integer}(n=zero(T))
 
-Create a new `NumProduceAccumulator` accumulator with the number of observations initialized to zero.
+Create a new `VariableOrderAccumulator` accumulator with the number of observations set to n
 """
-NumProduceAccumulator{T}() where {T<:Integer} = NumProduceAccumulator(zero(T))
-NumProduceAccumulator() = NumProduceAccumulator{Int}()
+VariableOrderAccumulator{T}(n=zero(T)) where {T<:Integer} = VariableOrderAccumulator(convert(T, n), OrderedDict{VarName, T}())
+VariableOrderAccumulator(n) = VariableOrderAccumulator{typeof(n)}(n)
+VariableOrderAccumulator() = VariableOrderAccumulator{Int}()
 
 function Base.show(io::IO, acc::LogPriorAccumulator)
     return print(io, "LogPriorAccumulator($(repr(acc.logp)))")
 end
 function Base.show(io::IO, acc::LogLikelihoodAccumulator)
     return print(io, "LogLikelihoodAccumulator($(repr(acc.logp)))")
 end
-function Base.show(io::IO, acc::NumProduceAccumulator)
-    return print(io, "NumProduceAccumulator($(repr(acc.num)))")
+function Base.show(io::IO, acc::VariableOrderAccumulator)
+    return print(io, "VariableOrderAccumulator($(repr(acc.num_produce)), $(repr(acc.order)))")
 end
 
 accumulator_name(::Type{<:LogPriorAccumulator}) = :LogPrior
 accumulator_name(::Type{<:LogLikelihoodAccumulator}) = :LogLikelihood
-accumulator_name(::Type{<:NumProduceAccumulator}) = :NumProduce
+accumulator_name(::Type{<:VariableOrderAccumulator}) = :VariableOrder
 
 split(::LogPriorAccumulator{T}) where {T} = LogPriorAccumulator(zero(T))
 split(::LogLikelihoodAccumulator{T}) where {T} = LogLikelihoodAccumulator(zero(T))
-split(acc::NumProduceAccumulator) = acc
+split(acc::VariableOrderAccumulator) = acc
 
 function combine(acc::LogPriorAccumulator, acc2::LogPriorAccumulator)
     return LogPriorAccumulator(acc.logp + acc2.logp)
 end
 function combine(acc::LogLikelihoodAccumulator, acc2::LogLikelihoodAccumulator)
     return LogLikelihoodAccumulator(acc.logp + acc2.logp)
 end
-function combine(acc::NumProduceAccumulator, acc2::NumProduceAccumulator)
-    return NumProduceAccumulator(max(acc.num, acc2.num))
+function combine(acc::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
+    # Note that assumptions are not allowed within in parallelised blocks, and thus the
+    # dictionaries should be identical.
+    return VariableOrderAccumulator(max(acc.num_produce, acc2.num_produce), merge(acc.order, acc2.order))
 end
 
 function Base.:+(acc1::LogPriorAccumulator, acc2::LogPriorAccumulator)
@@ -95,11 +105,10 @@ end
 function Base.:+(acc1::LogLikelihoodAccumulator, acc2::LogLikelihoodAccumulator)
     return LogLikelihoodAccumulator(acc1.logp + acc2.logp)
 end
-increment(acc::NumProduceAccumulator) = NumProduceAccumulator(acc.num + oneunit(acc.num))
+increment(acc::VariableOrderAccumulator) = VariableOrderAccumulator(acc.num_produce + oneunit(acc.num_produce), acc.order)
 
 Base.zero(acc::LogPriorAccumulator) = LogPriorAccumulator(zero(acc.logp))
 Base.zero(acc::LogLikelihoodAccumulator) = LogLikelihoodAccumulator(zero(acc.logp))
-Base.zero(acc::NumProduceAccumulator) = NumProduceAccumulator(zero(acc.num))
 
 function accumulate_assume!!(acc::LogPriorAccumulator, val, logjac, vn, right)
     return acc + LogPriorAccumulator(logpdf(right, val) + logjac)
@@ -114,8 +123,11 @@ function accumulate_observe!!(acc::LogLikelihoodAccumulator, right, left, vn)
     return acc + LogLikelihoodAccumulator(Distributions.loglikelihood(right, left))
 end
 
-accumulate_assume!!(acc::NumProduceAccumulator, val, logjac, vn, right) = acc
-accumulate_observe!!(acc::NumProduceAccumulator, right, left, vn) = increment(acc)
+function accumulate_assume!!(acc::VariableOrderAccumulator, val, logjac, vn, right)
+    acc.order[vn] = acc.num_produce
+    return acc
+end
+accumulate_observe!!(acc::VariableOrderAccumulator, right, left, vn) = increment(acc)
 
 function Base.convert(::Type{LogPriorAccumulator{T}}, acc::LogPriorAccumulator) where {T}
     return LogPriorAccumulator(convert(T, acc.logp))
@@ -126,15 +138,19 @@ function Base.convert(
     return LogLikelihoodAccumulator(convert(T, acc.logp))
 end
 function Base.convert(
-    ::Type{NumProduceAccumulator{T}}, acc::NumProduceAccumulator
-) where {T}
-    return NumProduceAccumulator(convert(T, acc.num))
+    ::Type{VariableOrderAccumulator{ElType, VnType}}, acc::VariableOrderAccumulator
+) where {ElType, VnType}
+    order = OrderedDict{VnType, ElType}()
+    for (k, v) in acc.order
+        order[convert(VnType, k)] = convert(ElType, v)
+    end
+    return VariableOrderAccumulator(convert(ElType, acc.num_produce), order)
 end
 
 # TODO(mhauru)
-# We ignore the convert_eltype calls for NumProduceAccumulator, by letting them fallback on
+# We ignore the convert_eltype calls for VariableOrderAccumulator, by letting them fallback on
 # convert_eltype(::AbstractAccumulator, ::Type). This is because they are only used to
-# deal with dual number types of AD backends, which shouldn't concern NumProduceAccumulator. This is
+# deal with dual number types of AD backends, which shouldn't concern VariableOrderAccumulator. This is
 # horribly hacky and should be fixed. See also comment in `unflatten` in `src/varinfo.jl`.
 function convert_eltype(::Type{T}, acc::LogPriorAccumulator) where {T}
     return LogPriorAccumulator(convert(T, acc.logp))
@@ -149,6 +165,6 @@ function default_accumulators(
     return AccumulatorTuple(
         LogPriorAccumulator{FloatT}(),
         LogLikelihoodAccumulator{FloatT}(),
-        NumProduceAccumulator{IntT}(),
+        VariableOrderAccumulator{IntT}(),
     )
 end

src/simple_varinfo.jl

@@ -125,15 +125,15 @@ Evaluation in transformed space of course also works:
 
 ```jldoctest simplevarinfo-general
 julia> vi = DynamicPPL.settrans!!(SimpleVarInfo((x = -1.0,)), true)
-Transformed SimpleVarInfo((x = -1.0,), (LogPrior = LogPriorAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0), NumProduce = NumProduceAccumulator(0)))
+Transformed SimpleVarInfo((x = -1.0,), (LogPrior = LogPriorAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0), VariableOrder = VariableOrderAccumulator(0, OrderedDict{VarName, Int64}())))
 
 julia> # (✓) Positive probability mass on negative numbers!
        getlogjoint(last(DynamicPPL.evaluate!!(m, vi, DynamicPPL.DefaultContext())))
 -1.3678794411714423
 
 julia> # While if we forget to indicate that it's transformed:
        vi = DynamicPPL.settrans!!(SimpleVarInfo((x = -1.0,)), false)
-SimpleVarInfo((x = -1.0,), (LogPrior = LogPriorAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0), NumProduce = NumProduceAccumulator(0)))
+SimpleVarInfo((x = -1.0,), (LogPrior = LogPriorAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0), VariableOrder = VariableOrderAccumulator(0, OrderedDict{VarName, Int64}())))
 
 julia> # (✓) No probability mass on negative numbers!
        getlogjoint(last(DynamicPPL.evaluate!!(m, vi, DynamicPPL.DefaultContext())))

src/threadsafe.jl

@@ -78,7 +78,7 @@ end
 
 syms(vi::ThreadSafeVarInfo) = syms(vi.varinfo)
 
-setorder!(vi::ThreadSafeVarInfo, vn::VarName, index::Int) = setorder!(vi.varinfo, vn, index)
+setorder!!(vi::ThreadSafeVarInfo, vn::VarName, index::Int) = ThreadSafeVarInfo(setorder!!(vi.varinfo, vn, index), vi.accs_by_thread)
 setval!(vi::ThreadSafeVarInfo, val, vn::VarName) = setval!(vi.varinfo, val, vn)
 
 keys(vi::ThreadSafeVarInfo) = keys(vi.varinfo)