Remove VariableOrderAcc and merge/subset on accs

Author: mhauru

HISTORY.md

@@ -72,6 +72,17 @@ The other case where one might use `PriorContext` was to use `@addlogprob!` to a
 Previously, this was accomplished by manually checking `__context__ isa DynamicPPL.PriorContext`.
 Now, you can write `@addlogprob (; logprior=x, loglikelihood=y)` to add `x` to the log-prior and `y` to the log-likelihood.
 
+### Removal of `order` and `num_produce`
+
+The `VarInfo` type used to carry with it information about the order in which variables had been evaluated, and a variable called `num_produce` for where in this evaluation the `VarInfo` was at. This was used in particle samplers in Turing.jl. The particle sampler code has been simplified and no longer needs this functionality, and thus we remove it from DynamicPPL. The following exported functions are now gone:
+
+  - `get_num_produce`
+  - `set_num_produce!!`
+  - `reset_num_produce!!`
+  - `increment_num_produce!!`
+  - `set_retained_vns_del!`
+  - `setorder!!`
+
 **Internals**
 
 ### Accumulators
@@ -83,7 +94,6 @@ This release overhauls how VarInfo objects track variables such as the log joint
   - `tilde_observe` and `observe` have been removed. `tilde_observe!!` still exists, and any contexts should modify its behaviour. We may further rework the call stack under `tilde_observe!!` in the near future.
   - `tilde_assume` no longer returns the log density of the current assumption as its second return value. We may further rework the `tilde_assume!!` call stack as well.
   - 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 `acclogp!!` method with a single scalar value has been deprecated and falls back on `accloglikelihood!!`, and the single scalar version of `setlogp!!` has been removed. Corresponding setter/accumulator functions exist for the log prior as well.

docs/src/api.md

@@ -334,17 +334,6 @@ unset_flag!
 is_flagged
 ```
 
-The following functions were used for sequential Monte Carlo methods.
-
-```@docs
-get_num_produce
-set_num_produce!!
-increment_num_produce!!
-reset_num_produce!!
-setorder!!
-set_retained_vns_del!
-```
-
 ```@docs
 Base.empty!
 ```

src/DynamicPPL.jl

@@ -51,7 +51,6 @@ export AbstractVarInfo,
     LogLikelihoodAccumulator,
     LogPriorAccumulator,
     LogJacobianAccumulator,
-    VariableOrderAccumulator,
     push!!,
     empty!!,
     subset,
@@ -72,15 +71,9 @@ export AbstractVarInfo,
     acclogprior!!,
     accloglikelihood!!,
     resetlogp!!,
-    get_num_produce,
-    set_num_produce!!,
-    reset_num_produce!!,
-    increment_num_produce!!,
-    set_retained_vns_del!,
     is_flagged,
     set_flag!,
     unset_flag!,
-    setorder!!,
     istrans,
     link,
     link!!,

src/abstract_varinfo.jl

@@ -441,24 +441,6 @@ 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)
@@ -509,8 +491,7 @@ function getindex_internal end
 @doc """
     empty!!(vi::AbstractVarInfo)
 
-Empty the fields of `vi.metadata` and reset `vi.logp[]` and `vi.num_produce[]` to
-zeros.
+Empty `vi` of variables and reset any `logp` accumulators zeros.
 
 This is useful when using a sampling algorithm that assumes an empty `vi`, e.g. `SMC`.
 """ BangBang.empty!!
@@ -1068,43 +1049,6 @@ function invlink_with_logpdf(vi::AbstractVarInfo, vn::VarName, dist, y)
     return x, logpdf(dist, x) + logjac
 end
 
-"""
-    get_num_produce(vi::AbstractVarInfo)
-
-Return the `num_produce` of `vi`.
-"""
-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`.
-"""
-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(:VariableOrder))
-
-"""
-    reset_num_produce!!(vi::AbstractVarInfo)
-
-Reset the value of `num_produce` in `vi` to 0.
-"""
-reset_num_produce!!(vi::AbstractVarInfo) = set_num_produce!!(vi, zero(get_num_produce(vi)))
-
 """
     from_internal_transform(varinfo::AbstractVarInfo, vn::VarName[, dist])
 

src/accumulators.jl

@@ -30,13 +30,6 @@ To be able to work with multi-threading, it should also implement:
 - `split(acc::T)`
 - `combine(acc::T, acc2::T)`
 
-If two accumulators of the same type should be merged in some non-trivial way, other than
-always keeping the second one over the first, `merge(acc1::T, acc2::T)` should be defined.
-
-If limiting the accumulator to a subset of `VarName`s is a meaningful operation and should
-do something other than copy the original accumulator, then
-`subset(acc::T, vns::AbstractVector{<:VarnName})` should be defined.`
-
 See the documentation for each of these functions for more details.
 """
 abstract type AbstractAccumulator end
@@ -120,24 +113,6 @@ used by various AD backends, should implement a method for this function.
 """
 convert_eltype(::Type, acc::AbstractAccumulator) = acc
 
-"""
-    subset(acc::AbstractAccumulator, vns::AbstractVector{<:VarName})
-
-Return a new accumulator that only contains the information for the `VarName`s in `vns`.
-
-By default returns a copy of `acc`. Subtypes should override this behaviour as needed.
-"""
-subset(acc::AbstractAccumulator, ::AbstractVector{<:VarName}) = copy(acc)
-
-"""
-    merge(acc1::AbstractAccumulator, acc2::AbstractAccumulator)
-
-Merge two accumulators of the same type. Returns a new accumulator of the same type.
-
-By default returns a copy of `acc2`. Subtypes should override this behaviour as needed.
-"""
-Base.merge(acc1::AbstractAccumulator, acc2::AbstractAccumulator) = copy(acc2)
-
 """
     AccumulatorTuple{N,T<:NamedTuple}
 
@@ -183,50 +158,6 @@ function Base.convert(::Type{AccumulatorTuple{N,T}}, accs::AccumulatorTuple{N})
     return AccumulatorTuple(convert(T, accs.nt))
 end
 
-"""
-    subset(at::AccumulatorTuple, vns::AbstractVector{<:VarName})
-
-Replace each accumulator `acc` in `at` with `subset(acc, vns)`.
-"""
-function subset(at::AccumulatorTuple, vns::AbstractVector{<:VarName})
-    return AccumulatorTuple(map(Base.Fix2(subset, vns), at.nt))
-end
-
-"""
-    _joint_keys(nt1::NamedTuple, nt2::NamedTuple)
-
-A helper function that returns three tuples of keys given two `NamedTuple`s:
-The keys only in `nt1`, only in `nt2`, and in both, and in that order.
-
-Implemented as a generated function to enable constant propagation of the result in `merge`.
-"""
-@generated function _joint_keys(
-    nt1::NamedTuple{names1}, nt2::NamedTuple{names2}
-) where {names1,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
-
-"""
-    merge(at1::AccumulatorTuple, at2::AccumulatorTuple)
-
-Merge two `AccumulatorTuple`s.
-
-For any `accumulator_name` that exists in both `at1` and `at2`, we call `merge` on the two
-accumulators themselves. Other accumulators are copied.
-"""
-function Base.merge(at1::AccumulatorTuple, at2::AccumulatorTuple)
-    keys_in_at1, keys_in_at2, keys_in_both = _joint_keys(at1.nt, at2.nt)
-    accs_in_at1 = (getfield(at1.nt, key) for key in keys_in_at1)
-    accs_in_at2 = (getfield(at2.nt, key) for key in keys_in_at2)
-    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_both..., accs_in_at2...)
-end
-
 """
     setacc!!(at::AccumulatorTuple, acc::AbstractAccumulator)
 

src/default_accumulators.jl

@@ -166,119 +166,12 @@ function accumulate_observe!!(acc::LogLikelihoodAccumulator, right, left, vn)
     return acclogp(acc, Distributions.loglikelihood(right, left))
 end
 
-"""
-    VariableOrderAccumulator{T} <: AbstractAccumulator
-
-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 VariableOrderAccumulator{Eltype<:Integer,VNType<:VarName} <: AbstractAccumulator
-    "the number of observations"
-    num_produce::Eltype
-    "mapping of variable names to their order in the model"
-    order::Dict{VNType,Eltype}
-end
-
-"""
-    VariableOrderAccumulator{T<:Integer}(n=zero(T))
-
-Create a new `VariableOrderAccumulator` with the number of observations set to `n`.
-"""
-VariableOrderAccumulator{T}(n=zero(T)) where {T<:Integer} =
-    VariableOrderAccumulator(convert(T, n), Dict{VarName,T}())
-VariableOrderAccumulator(n) = VariableOrderAccumulator{typeof(n)}(n)
-VariableOrderAccumulator() = VariableOrderAccumulator{Int}()
-
-function Base.copy(acc::VariableOrderAccumulator)
-    return VariableOrderAccumulator(acc.num_produce, copy(acc.order))
-end
-
-function Base.show(io::IO, acc::VariableOrderAccumulator)
-    return print(
-        io, "VariableOrderAccumulator($(string(acc.num_produce)), $(repr(acc.order)))"
-    )
-end
-
-function Base.:(==)(acc1::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
-    return acc1.num_produce == acc2.num_produce && acc1.order == acc2.order
-end
-
-function Base.isequal(acc1::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
-    return isequal(acc1.num_produce, acc2.num_produce) && isequal(acc1.order, acc2.order)
-end
-
-function Base.hash(acc::VariableOrderAccumulator, h::UInt)
-    return hash((VariableOrderAccumulator, acc.num_produce, acc.order), h)
-end
-
-accumulator_name(::Type{<:VariableOrderAccumulator}) = :VariableOrder
-
-split(acc::VariableOrderAccumulator) = copy(acc)
-
-function combine(acc::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
-    # Note that assumptions are not allowed 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 increment(acc::VariableOrderAccumulator)
-    return VariableOrderAccumulator(acc.num_produce + oneunit(acc.num_produce), acc.order)
-end
-
-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{VariableOrderAccumulator{ElType,VnType}}, acc::VariableOrderAccumulator
-) where {ElType,VnType}
-    order = Dict{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 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 VariableOrderAccumulator. This is
-# horribly hacky and should be fixed. See also comment in `unflatten` in `src/varinfo.jl`.
-
 function default_accumulators(
     ::Type{FloatT}=LogProbType, ::Type{IntT}=Int
 ) where {FloatT,IntT}
     return AccumulatorTuple(
         LogPriorAccumulator{FloatT}(),
         LogJacobianAccumulator{FloatT}(),
         LogLikelihoodAccumulator{FloatT}(),
-        VariableOrderAccumulator{IntT}(),
     )
 end
-
-function subset(acc::VariableOrderAccumulator, vns::AbstractVector{<:VarName})
-    order = filter(pair -> any(subsumes(vn, first(pair)) for vn in vns), acc.order)
-    return VariableOrderAccumulator(acc.num_produce, order)
-end
-
-"""
-    merge(acc1::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
-
-Merge two `VariableOrderAccumulator` instances.
-
-The `num_produce` field of the return value is the `num_produce` of `acc2`.
-"""
-function Base.merge(acc1::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
-    return VariableOrderAccumulator(acc2.num_produce, merge(acc1.order, acc2.order))
-end

src/simple_varinfo.jl

@@ -122,15 +122,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), LogJacobian = LogJacobianAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0), VariableOrder = VariableOrderAccumulator(0, Dict{VarName, Int64}())))
+Transformed SimpleVarInfo((x = -1.0,), (LogPrior = LogPriorAccumulator(0.0), LogJacobian = LogJacobianAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0)))
 
 julia> # (✓) Positive probability mass on negative numbers!
        getlogjoint_internal(last(DynamicPPL.evaluate!!(m, vi)))
 -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), LogJacobian = LogJacobianAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0), VariableOrder = VariableOrderAccumulator(0, Dict{VarName, Int64}())))
+SimpleVarInfo((x = -1.0,), (LogPrior = LogPriorAccumulator(0.0), LogJacobian = LogJacobianAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0)))
 
 julia> # (✓) No probability mass on negative numbers!
        getlogjoint_internal(last(DynamicPPL.evaluate!!(m, vi)))
@@ -418,7 +418,7 @@ Base.eltype(::SimpleOrThreadSafeSimple{<:Any,V}) where {V} = V
 # `subset`
 function subset(varinfo::SimpleVarInfo, vns::AbstractVector{<:VarName})
     return SimpleVarInfo(
-        _subset(varinfo.values, vns), subset(getaccs(varinfo), vns), varinfo.transformation
+        _subset(varinfo.values, vns), map(copy, getaccs(varinfo)), varinfo.transformation
     )
 end
 
@@ -456,7 +456,7 @@ _subset(x::VarNamedVector, vns) = subset(x, vns)
 # `merge`
 function Base.merge(varinfo_left::SimpleVarInfo, varinfo_right::SimpleVarInfo)
     values = merge(varinfo_left.values, varinfo_right.values)
-    accs = merge(getaccs(varinfo_left), getaccs(varinfo_right))
+    accs = map(copy, getaccs(varinfo_right))
     transformation = merge_transformations(
         varinfo_left.transformation, varinfo_right.transformation
     )