Add comparison methods

Author: mhauru

src/accumulators.jl

@@ -136,6 +136,8 @@ function Base.haskey(at::AccumulatorTuple, ::Val{accname}) where {accname}
     @inline return haskey(at.nt, accname)
 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)
 
 function Base.convert(::Type{AccumulatorTuple{N,T}}, accs::AccumulatorTuple{N}) where {N,T}
     return AccumulatorTuple(convert(T, accs.nt))

src/default_accumulators.jl

@@ -57,15 +57,16 @@ struct VariableOrderAccumulator{Eltype<:Integer,VNType<:VarName} <: AbstractAccu
     "the number of observations"
     num_produce::Eltype
     "mapping of variable names to their order in the model"
-    order::OrderedDict{VNType, Eltype}
+    order::OrderedDict{VNType,Eltype}
 end
 
 """
     VariableOrderAccumulator{T<:Integer}(n=zero(T))
 
 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{T}(n=zero(T)) where {T<:Integer} =
+    VariableOrderAccumulator(convert(T, n), OrderedDict{VarName,T}())
 VariableOrderAccumulator(n) = VariableOrderAccumulator{typeof(n)}(n)
 VariableOrderAccumulator() = VariableOrderAccumulator{Int}()
 
@@ -76,7 +77,38 @@ function Base.show(io::IO, acc::LogLikelihoodAccumulator)
     return print(io, "LogLikelihoodAccumulator($(repr(acc.logp)))")
 end
 function Base.show(io::IO, acc::VariableOrderAccumulator)
-    return print(io, "VariableOrderAccumulator($(repr(acc.num_produce)), $(repr(acc.order)))")
+    return print(
+        io, "VariableOrderAccumulator($(repr(acc.num_produce)), $(repr(acc.order)))"
+    )
+end
+
+# Note that == and isequal are different, and equality under the latter should imply
+# equality of hashes. Both of the below implementations are also different from the default
+# implementation for structs.
+Base.:(==)(acc1::LogPriorAccumulator, acc2::LogPriorAccumulator) = acc1.logp == acc2.logp
+function Base.:(==)(acc1::LogLikelihoodAccumulator, acc2::LogLikelihoodAccumulator)
+    return acc1.logp == acc2.logp
+end
+function Base.:(==)(acc1::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
+    return acc1.num_produce == acc2.num_produce && acc1.order == acc2.order
+end
+
+function Base.isequal(acc1::LogPriorAccumulator, acc2::LogPriorAccumulator)
+    return isequal(acc1.logp, acc2.logp)
+end
+function Base.isequal(acc1::LogLikelihoodAccumulator, acc2::LogLikelihoodAccumulator)
+    return isequal(acc1.logp, acc2.logp)
+end
+function Base.isequal(acc1::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
+    return isequal(acc1.num_produce, acc2.num_produce) && isequal(acc1.order, acc2.order)
+end
+
+Base.hash(acc::LogPriorAccumulator, h::UInt) = hash((LogPriorAccumulator, acc.logp), h)
+function Base.hash(acc::LogLikelihoodAccumulator, h::UInt)
+    return hash((LogLikelihoodAccumulator, acc.logp), h)
+end
+function Base.hash(acc::VariableOrderAccumulator, h::UInt)
+    return hash((VariableOrderAccumulator, acc.num_produce, acc.order), h)
 end
 
 accumulator_name(::Type{<:LogPriorAccumulator}) = :LogPrior
@@ -96,7 +128,9 @@ end
 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))
+    return VariableOrderAccumulator(
+        max(acc.num_produce, acc2.num_produce), merge(acc.order, acc2.order)
+    )
 end
 
 function Base.:+(acc1::LogPriorAccumulator, acc2::LogPriorAccumulator)
@@ -105,7 +139,9 @@ end
 function Base.:+(acc1::LogLikelihoodAccumulator, acc2::LogLikelihoodAccumulator)
     return LogLikelihoodAccumulator(acc1.logp + acc2.logp)
 end
-increment(acc::VariableOrderAccumulator) = VariableOrderAccumulator(acc.num_produce + oneunit(acc.num_produce), acc.order)
+function increment(acc::VariableOrderAccumulator)
+    return VariableOrderAccumulator(acc.num_produce + oneunit(acc.num_produce), acc.order)
+end
 
 Base.zero(acc::LogPriorAccumulator) = LogPriorAccumulator(zero(acc.logp))
 Base.zero(acc::LogLikelihoodAccumulator) = LogLikelihoodAccumulator(zero(acc.logp))
@@ -138,9 +174,9 @@ function Base.convert(
     return LogLikelihoodAccumulator(convert(T, acc.logp))
 end
 function Base.convert(
-    ::Type{VariableOrderAccumulator{ElType, VnType}}, acc::VariableOrderAccumulator
-) where {ElType, VnType}
-    order = OrderedDict{VnType, ElType}()
+    ::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

src/simple_varinfo.jl

@@ -198,6 +198,12 @@ struct SimpleVarInfo{NT,Accs<:AccumulatorTuple where {N},C<:AbstractTransformati
     transformation::C
 end
 
+function Base.:(==)(vi1::SimpleVarInfo, vi2::SimpleVarInfo)
+    return vi1.values == vi2.values &&
+           vi1.accs == vi2.accs &&
+           vi1.transformation == vi2.transformation
+end
+
 transformation(vi::SimpleVarInfo) = vi.transformation
 
 function SimpleVarInfo(values, accs)

src/threadsafe.jl

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

src/varinfo.jl

@@ -60,6 +60,17 @@ struct Metadata{
     flags::Dict{String,BitVector}
 end
 
+function Base.:(==)(md1::Metadata, md2::Metadata)
+    return (
+        md1.idcs == md2.idcs &&
+        md1.vns == md2.vns &&
+        md1.ranges == md2.ranges &&
+        md1.vals == md2.vals &&
+        md1.dists == md2.dists &&
+        md1.flags == md2.flags
+    )
+end
+
 ###########
 # VarInfo #
 ###########
@@ -155,6 +166,10 @@ const VarInfoOrThreadSafeVarInfo{Tmeta} = Union{
     VarInfo{Tmeta},ThreadSafeVarInfo{<:VarInfo{Tmeta}}
 }
 
+function Base.:(==)(vi1::VarInfo, vi2::VarInfo)
+    return (vi1.metadata == vi2.metadata && vi1.accs == vi2.accs)
+end
+
 # NOTE: This is kind of weird, but it effectively preserves the "old"
 # behavior where we're allowed to call `link!` on the same `VarInfo`
 # multiple times.
@@ -275,9 +290,7 @@ function typed_varinfo(vi::UntypedVarInfo)
 
         push!(
             new_metas,
-            Metadata(
-                sym_idcs, sym_vns, sym_ranges, sym_vals, sym_dists, sym_flags
-            ),
+            Metadata(sym_idcs, sym_vns, sym_ranges, sym_vals, sym_dists, sym_flags),
         )
     end
     nt = NamedTuple{syms_tuple}(Tuple(new_metas))
@@ -597,14 +610,7 @@ function subset(metadata::Metadata, vns_given::AbstractVector{VN}) where {VN<:Va
     end
 
     flags = Dict(k => v[indices_for_vns] for (k, v) in metadata.flags)
-    return Metadata(
-        indices,
-        vns,
-        ranges,
-        vals,
-        metadata.dists[indices_for_vns],
-        flags,
-    )
+    return Metadata(indices, vns, ranges, vals, metadata.dists[indices_for_vns], flags)
 end
 
 function Base.merge(varinfo_left::VarInfo, varinfo_right::VarInfo)

test/accumulators.jl

@@ -71,8 +71,9 @@ using DynamicPPL:
             @test convert(
                 LogLikelihoodAccumulator{Float32}, LogLikelihoodAccumulator(1.0)
             ) == LogLikelihoodAccumulator{Float32}(1.0f0)
-            @test convert(VariableOrderAccumulator{UInt8,VarName}, VariableOrderAccumulator(1)) ==
-                VariableOrderAccumulator{UInt8}(1)
+            @test convert(
+                VariableOrderAccumulator{UInt8,VarName}, VariableOrderAccumulator(1)
+            ) == VariableOrderAccumulator{UInt8}(1)
 
             @test convert_eltype(Float32, LogPriorAccumulator(1.0)) ==
                 LogPriorAccumulator{Float32}(1.0f0)