Use copy rather than deepcopy for accumulators

Author: mhauru

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, right, left, vn)`
 - `accumulate_assume!!(acc::T, val, logjac, vn, right)`
+- `Base.copy(acc::T)`
 
 To be able to work with multi-threading, it should also implement:
 - `split(acc::T)`
@@ -138,6 +139,7 @@ end
 Base.keys(at::AccumulatorTuple) = keys(at.nt)
 Base.:(==)(at1::AccumulatorTuple, at2::AccumulatorTuple) = at1.nt == at2.nt
 Base.hash(at::AccumulatorTuple, h::UInt) = Base.hash((AccumulatorTuple, at.nt), h)
+Base.copy(at::AccumulatorTuple) = AccumulatorTuple(map(copy, at.nt))
 
 function Base.convert(::Type{AccumulatorTuple{N,T}}, accs::AccumulatorTuple{N}) where {N,T}
     return AccumulatorTuple(convert(T, accs.nt))

src/default_accumulators.jl

@@ -66,10 +66,16 @@ end
 Create a new `VariableOrderAccumulator` accumulator with the number of observations set to n
 """
 VariableOrderAccumulator{T}(n=zero(T)) where {T<:Integer} =
-    VariableOrderAccumulator(convert(T, n), OrderedDict{VarName,T}())
+    VariableOrderAccumulator(convert(T, n), Dict{VarName,T}())
 VariableOrderAccumulator(n) = VariableOrderAccumulator{typeof(n)}(n)
 VariableOrderAccumulator() = VariableOrderAccumulator{Int}()
 
+Base.copy(acc::LogPriorAccumulator) = acc
+Base.copy(acc::LogLikelihoodAccumulator) = acc
+function Base.copy(acc::VariableOrderAccumulator)
+    return VariableOrderAccumulator(acc.num_produce, copy(acc.order))
+end
+
 function Base.show(io::IO, acc::LogPriorAccumulator)
     return print(io, "LogPriorAccumulator($(repr(acc.logp)))")
 end

src/extract_priors.jl

@@ -4,6 +4,10 @@ end
 
 PriorDistributionAccumulator() = PriorDistributionAccumulator(OrderedDict{VarName,Any}())
 
+function Base.copy(acc::PriorDistributionAccumulator)
+    return PriorDistributionAccumulator(copy(acc.priors))
+end
+
 accumulator_name(::PriorDistributionAccumulator) = :PriorDistributionAccumulator
 
 split(acc::PriorDistributionAccumulator) = PriorDistributionAccumulator(empty(acc.priors))

src/pointwise_logdensities.jl

@@ -31,6 +31,10 @@ function PointwiseLogProbAccumulator{whichlogprob,KeyType}() where {whichlogprob
     return PointwiseLogProbAccumulator{whichlogprob,KeyType,typeof(logps)}(logps)
 end
 
+function Base.copy(acc::PointwiseLogProbAccumulator{whichlogprob}) where {whichlogprob}
+    return PointwiseLogProbAccumulator{whichlogprob}(copy(acc.logps))
+end
+
 function Base.push!(acc::PointwiseLogProbAccumulator, vn, logp)
     logps = acc.logps
     # The last(fieldtypes(eltype(...))) gets the type of the values, rather than the keys.

src/simple_varinfo.jl

@@ -248,7 +248,7 @@ end
 # Constructor from `VarInfo`.
 function SimpleVarInfo(vi::NTVarInfo, ::Type{D}) where {D}
     values = values_as(vi, D)
-    return SimpleVarInfo(values, deepcopy(getaccs(vi)))
+    return SimpleVarInfo(values, copy(getaccs(vi)))
 end
 function SimpleVarInfo{T}(vi::NTVarInfo, ::Type{D}) where {T<:Real,D}
     values = values_as(vi, D)
@@ -447,7 +447,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 = deepcopy(getaccs(varinfo_right))
+    accs = copy(getaccs(varinfo_right))
     transformation = merge_transformations(
         varinfo_left.transformation, varinfo_right.transformation
     )

src/values_as_in_model.jl

@@ -20,6 +20,10 @@ function ValuesAsInModelAccumulator(include_colon_eq)
     return ValuesAsInModelAccumulator(OrderedDict(), include_colon_eq)
 end
 
+function Base.copy(acc::ValuesAsInModelAccumulator)
+    return ValuesAsInModelAccumulator(copy(acc.values), acc.include_colon_eq)
+end
+
 accumulator_name(::Type{<:ValuesAsInModelAccumulator}) = :ValuesAsInModel
 
 function split(acc::ValuesAsInModelAccumulator)

src/varinfo.jl

@@ -294,7 +294,7 @@ function typed_varinfo(vi::UntypedVarInfo)
         )
     end
     nt = NamedTuple{syms_tuple}(Tuple(new_metas))
-    return VarInfo(nt, deepcopy(vi.accs))
+    return VarInfo(nt, copy(vi.accs))
 end
 function typed_varinfo(vi::NTVarInfo)
     # This function preserves the behaviour of typed_varinfo(vi) where vi is
@@ -355,7 +355,7 @@ single `VarNamedVector` as its metadata field.
 """
 function untyped_vector_varinfo(vi::UntypedVarInfo)
     md = metadata_to_varnamedvector(vi.metadata)
-    return VarInfo(md, deepcopy(vi.accs))
+    return VarInfo(md, copy(vi.accs))
 end
 function untyped_vector_varinfo(
     rng::Random.AbstractRNG,
@@ -398,12 +398,12 @@ NamedTuple of `VarNamedVector`s as its metadata field.
 """
 function typed_vector_varinfo(vi::NTVarInfo)
     md = map(metadata_to_varnamedvector, vi.metadata)
-    return VarInfo(md, deepcopy(vi.accs))
+    return VarInfo(md, copy(vi.accs))
 end
 function typed_vector_varinfo(vi::UntypedVectorVarInfo)
     new_metas = group_by_symbol(vi.metadata)
     nt = NamedTuple(new_metas)
-    return VarInfo(nt, deepcopy(vi.accs))
+    return VarInfo(nt, copy(vi.accs))
 end
 function typed_vector_varinfo(
     rng::Random.AbstractRNG,
@@ -455,8 +455,7 @@ function unflatten(vi::VarInfo, x::AbstractVector)
     # convert to into an intermediate variable makes this unstable (constant propagation)
     # fails. Take care when editing.
     accs = map(
-        acc -> convert_eltype(float_type_with_fallback(eltype(x)), acc),
-        deepcopy(getaccs(vi)),
+        acc -> convert_eltype(float_type_with_fallback(eltype(x)), acc), copy(getaccs(vi))
     )
     return VarInfo(md, accs)
 end
@@ -538,7 +537,7 @@ end
 
 function subset(varinfo::VarInfo, vns::AbstractVector{<:VarName})
     metadata = subset(varinfo.metadata, vns)
-    return VarInfo(metadata, deepcopy(varinfo.accs))
+    return VarInfo(metadata, copy(varinfo.accs))
 end
 
 function subset(metadata::NamedTuple, vns::AbstractVector{<:VarName})
@@ -619,7 +618,7 @@ end
 
 function _merge(varinfo_left::VarInfo, varinfo_right::VarInfo)
     metadata = merge_metadata(varinfo_left.metadata, varinfo_right.metadata)
-    return VarInfo(metadata, deepcopy(varinfo_right.accs))
+    return VarInfo(metadata, copy(varinfo_right.accs))
 end
 
 function merge_metadata(vnv_left::VarNamedVector, vnv_right::VarNamedVector)