Simplify accs with LogProbAccumulator

Author: mhauru

src/default_accumulators.jl

@@ -1,5 +1,84 @@
 """
-    LogPriorAccumulator{T<:Real} <: AbstractAccumulator
+    LogProbAccumulator{T} <: AbstractAccumulator
+
+An abstract type for accumulators that hold a single scalar log probability value.
+
+Every subtype of `LogProbAccumulator` must implement
+* A method for `logp` that returns the scalar log probability value that defines it.
+* A single-argument constructor that takes a `logp` value.
+* `accumulator_name`, `accumulate_assume!!`, and `accumulate_observe!!` methods like any
+    other accumulator.
+
+`LogProbAccumulator` provides implementations for other common functions, like convenience
+constructors, `copy`, `show`, `==`, `isequal`, `hash`, `split`, and `combine`.
+
+This type has no great conceptual significance, it just reduces code duplication between
+types like LogPriorAccumulator, LogJacobianAccumulator, and LogLikelihoodAccumulator.
+"""
+abstract type LogProbAccumulator{T<:Real} <: AbstractAccumulator end
+
+# The first of the below methods sets AccType{T}() = AccType(zero(T)) for any
+# AccType <: LogProbAccumulator{T}. The second one sets LogProbType as the default eltype T
+# when calling AccType().
+"""
+    LogProbAccumulator{T}()
+
+Create a new `LogProbAccumulator` accumulator with the log prior initialized to zero.
+"""
+(::Type{AccType})() where {T<:Real,AccType<:LogProbAccumulator{T}} = AccType(zero(T))
+(::Type{AccType})() where {AccType<:LogProbAccumulator} = AccType{LogProbType}()
+
+Base.copy(acc::LogProbAccumulator) = acc
+
+function Base.show(io::IO, acc::LogProbAccumulator)
+    return print(io, "$(repr(accumulator_name(acc)))($(repr(logp(acc)))))")
+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.
+function Base.:(==)(acc1::LogProbAccumulator, acc2::LogProbAccumulator)
+    return accumulator_name(acc1) === accumulator_name(acc2) && logp(acc1) == logp(acc2)
+end
+
+function Base.isequal(acc1::LogProbAccumulator, acc2::LogProbAccumulator)
+    return basetypeof(acc1) === basetypeof(acc2) && isequal(logp(acc1), logp(acc2))
+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))
+
+function combine(acc::LogProbAccumulator, acc2::LogProbAccumulator)
+    if basetypeof(acc) !== basetypeof(acc2)
+        msg = "Cannot combine accumulators of different types: $(basetypeof(acc)) and $(basetypeof(acc2))"
+        throw(ArgumentError(msg))
+    end
+    return basetypeof(acc)(logp(acc) + logp(acc2))
+end
+
+function Base.:+(acc1::LogProbAccumulator, acc2::LogProbAccumulator)
+    if basetypeof(acc1) !== basetypeof(acc2)
+        msg = "Cannot add accumulators of different types: $(basetypeof(acc1)) and $(basetypeof(acc2))"
+        throw(ArgumentError(msg))
+    end
+    return basetypeof(acc1)(logp(acc1) + logp(acc2))
+end
+
+Base.zero(acc::T) where {T<:LogProbAccumulator} = T(zero(logp(acc)))
+
+function Base.convert(
+    ::Type{AccType}, acc::LogProbAccumulator
+) where {T,AccType<:LogProbAccumulator{T}}
+    return AccType(convert(T, logp(acc)))
+end
+
+function convert_eltype(::Type{T}, acc::LogProbAccumulator) where {T}
+    return basetypeof(acc)(convert(T, logp(acc)))
+end
+
+"""
+    LogPriorAccumulator{T<:Real} <: LogProbAccumulator{T}
 
 An accumulator that tracks the cumulative log prior during model execution.
 
@@ -10,21 +89,22 @@ linked or not.
 # Fields
 $(TYPEDFIELDS)
 """
-struct LogPriorAccumulator{T<:Real} <: AbstractAccumulator
+struct LogPriorAccumulator{T<:Real} <: LogProbAccumulator{T}
     "the scalar log prior value"
     logp::T
 end
 
-"""
-    LogPriorAccumulator{T}()
+logp(acc::LogPriorAccumulator) = acc.logp
 
-Create a new `LogPriorAccumulator` accumulator with the log prior initialized to zero.
-"""
-LogPriorAccumulator{T}() where {T<:Real} = LogPriorAccumulator(zero(T))
-LogPriorAccumulator() = LogPriorAccumulator{LogProbType}()
+accumulator_name(::Type{<:LogPriorAccumulator}) = :LogPrior
+
+function accumulate_assume!!(acc::LogPriorAccumulator, val, logjac, vn, right)
+    return acc + LogPriorAccumulator(logpdf(right, val))
+end
+accumulate_observe!!(acc::LogPriorAccumulator, right, left, vn) = acc
 
 """
-    LogJacobianAccumulator{T<:Real} <: AbstractAccumulator
+    LogJacobianAccumulator{T<:Real} <: LogProbAccumulator{T}
 
 An accumulator that tracks the cumulative log Jacobian (technically,
 log(abs(det(J)))) during model execution. Specifically, J refers to the
@@ -53,39 +133,44 @@ distribution to unconstrained space.
 # Fields
 $(TYPEDFIELDS)
 """
-struct LogJacobianAccumulator{T<:Real} <: AbstractAccumulator
+struct LogJacobianAccumulator{T<:Real} <: LogProbAccumulator{T}
     "the logabsdet of the link transform Jacobian"
     logJ::T
 end
 
-"""
-    LogJacobianAccumulator{T}()
+logp(acc::LogJacobianAccumulator) = acc.logJ
 
-Create a new `LogJacobianAccumulator` accumulator with the log Jacobian initialized to zero.
-"""
-LogJacobianAccumulator{T}() where {T<:Real} = LogJacobianAccumulator(zero(T))
-LogJacobianAccumulator() = LogJacobianAccumulator{LogProbType}()
+accumulator_name(::Type{<:LogJacobianAccumulator}) = :LogJacobian
+
+function accumulate_assume!!(acc::LogJacobianAccumulator, val, logjac, vn, right)
+    return acc + LogJacobianAccumulator(logjac)
+end
+accumulate_observe!!(acc::LogJacobianAccumulator, right, left, vn) = acc
 
 """
-    LogLikelihoodAccumulator{T<:Real} <: AbstractAccumulator
+    LogLikelihoodAccumulator{T<:Real} <: LogProbAccumulator{T}
 
 An accumulator that tracks the cumulative log likelihood during model execution.
 
 # Fields
 $(TYPEDFIELDS)
 """
-struct LogLikelihoodAccumulator{T<:Real} <: AbstractAccumulator
+struct LogLikelihoodAccumulator{T<:Real} <: LogProbAccumulator{T}
     "the scalar log likelihood value"
     logp::T
 end
 
-"""
-    LogLikelihoodAccumulator{T}()
+logp(acc::LogLikelihoodAccumulator) = acc.logp
 
-Create a new `LogLikelihoodAccumulator` accumulator with the log likelihood initialized to zero.
-"""
-LogLikelihoodAccumulator{T}() where {T<:Real} = LogLikelihoodAccumulator(zero(T))
-LogLikelihoodAccumulator() = LogLikelihoodAccumulator{LogProbType}()
+accumulator_name(::Type{<:LogLikelihoodAccumulator}) = :LogLikelihood
+
+accumulate_assume!!(acc::LogLikelihoodAccumulator, val, logjac, vn, right) = acc
+function accumulate_observe!!(acc::LogLikelihoodAccumulator, right, left, vn)
+    # Note that it's important to use the loglikelihood function here, not logpdf, because
+    # they handle vectors differently:
+    # https://github.com/JuliaStats/Distributions.jl/issues/1972
+    return acc + LogLikelihoodAccumulator(Distributions.loglikelihood(right, left))
+end
 
 """
     VariableOrderAccumulator{T} <: AbstractAccumulator
@@ -117,85 +202,32 @@ VariableOrderAccumulator{T}(n=zero(T)) where {T<:Integer} =
 VariableOrderAccumulator(n) = VariableOrderAccumulator{typeof(n)}(n)
 VariableOrderAccumulator() = VariableOrderAccumulator{Int}()
 
-Base.copy(acc::LogPriorAccumulator) = acc
-Base.copy(acc::LogJacobianAccumulator) = 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
-function Base.show(io::IO, acc::LogJacobianAccumulator)
-    return print(io, "LogJacobianAccumulator($(repr(acc.logJ)))")
-end
-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)))"
     )
 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::LogJacobianAccumulator, acc2::LogJacobianAccumulator)
-    return acc1.logJ == acc2.logJ
-end
-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::LogJacobianAccumulator, acc2::LogJacobianAccumulator)
-    return isequal(acc1.logJ, acc2.logJ)
-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::LogJacobianAccumulator, h::UInt)
-    return hash((LogJacobianAccumulator, acc.logJ), h)
-end
-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
-accumulator_name(::Type{<:LogJacobianAccumulator}) = :LogJacobian
-accumulator_name(::Type{<:LogLikelihoodAccumulator}) = :LogLikelihood
 accumulator_name(::Type{<:VariableOrderAccumulator}) = :VariableOrder
 
-split(::LogPriorAccumulator{T}) where {T} = LogPriorAccumulator(zero(T))
-split(::LogJacobianAccumulator{T}) where {T} = LogJacobianAccumulator(zero(T))
-split(::LogLikelihoodAccumulator{T}) where {T} = LogLikelihoodAccumulator(zero(T))
 split(acc::VariableOrderAccumulator) = copy(acc)
 
-function combine(acc::LogPriorAccumulator, acc2::LogPriorAccumulator)
-    return LogPriorAccumulator(acc.logp + acc2.logp)
-end
-function combine(acc::LogJacobianAccumulator, acc2::LogJacobianAccumulator)
-    return LogJacobianAccumulator(acc.logJ + acc2.logJ)
-end
-function combine(acc::LogLikelihoodAccumulator, acc2::LogLikelihoodAccumulator)
-    return LogLikelihoodAccumulator(acc.logp + acc2.logp)
-end
 function combine(acc::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
     # Note that assumptions are not allowed in parallelised blocks, and thus the
     # dictionaries should be identical.
@@ -204,60 +236,16 @@ function combine(acc::VariableOrderAccumulator, acc2::VariableOrderAccumulator)
     )
 end
 
-function Base.:+(acc1::LogPriorAccumulator, acc2::LogPriorAccumulator)
-    return LogPriorAccumulator(acc1.logp + acc2.logp)
-end
-function Base.:+(acc1::LogJacobianAccumulator, acc2::LogJacobianAccumulator)
-    return LogJacobianAccumulator(acc1.logJ + acc2.logJ)
-end
-function Base.:+(acc1::LogLikelihoodAccumulator, acc2::LogLikelihoodAccumulator)
-    return LogLikelihoodAccumulator(acc1.logp + acc2.logp)
-end
 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::LogJacobianAccumulator) = LogJacobianAccumulator(zero(acc.logJ))
-Base.zero(acc::LogLikelihoodAccumulator) = LogLikelihoodAccumulator(zero(acc.logp))
-
-function accumulate_assume!!(acc::LogPriorAccumulator, val, logjac, vn, right)
-    return acc + LogPriorAccumulator(logpdf(right, val))
-end
-accumulate_observe!!(acc::LogPriorAccumulator, right, left, vn) = acc
-
-function accumulate_assume!!(acc::LogJacobianAccumulator, val, logjac, vn, right)
-    return acc + LogJacobianAccumulator(logjac)
-end
-accumulate_observe!!(acc::LogJacobianAccumulator, right, left, vn) = acc
-
-accumulate_assume!!(acc::LogLikelihoodAccumulator, val, logjac, vn, right) = acc
-function accumulate_observe!!(acc::LogLikelihoodAccumulator, right, left, vn)
-    # Note that it's important to use the loglikelihood function here, not logpdf, because
-    # they handle vectors differently:
-    # https://github.com/JuliaStats/Distributions.jl/issues/1972
-    return acc + LogLikelihoodAccumulator(Distributions.loglikelihood(right, left))
-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{LogPriorAccumulator{T}}, acc::LogPriorAccumulator) where {T}
-    return LogPriorAccumulator(convert(T, acc.logp))
-end
-function Base.convert(
-    ::Type{LogJacobianAccumulator{T}}, acc::LogJacobianAccumulator
-) where {T}
-    return LogJacobianAccumulator(convert(T, acc.logJ))
-end
-function Base.convert(
-    ::Type{LogLikelihoodAccumulator{T}}, acc::LogLikelihoodAccumulator
-) where {T}
-    return LogLikelihoodAccumulator(convert(T, acc.logp))
-end
 function Base.convert(
     ::Type{VariableOrderAccumulator{ElType,VnType}}, acc::VariableOrderAccumulator
 ) where {ElType,VnType}
@@ -273,15 +261,6 @@ end
 # 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 convert_eltype(::Type{T}, acc::LogPriorAccumulator) where {T}
-    return LogPriorAccumulator(convert(T, acc.logp))
-end
-function convert_eltype(::Type{T}, acc::LogJacobianAccumulator) where {T}
-    return LogJacobianAccumulator(convert(T, acc.logJ))
-end
-function convert_eltype(::Type{T}, acc::LogLikelihoodAccumulator) where {T}
-    return LogLikelihoodAccumulator(convert(T, acc.logp))
-end
 
 function default_accumulators(
     ::Type{FloatT}=LogProbType, ::Type{IntT}=Int

src/utils.jl

@@ -1332,3 +1332,10 @@ function group_varnames_by_symbol(vns::VarNameTuple)
     elements = map(collect, tuple((filter(vn -> getsym(vn) == s, vns) for s in syms)...))
     return NamedTuple{syms}(elements)
 end
+
+"""
+    basetypeof(x)
+
+Return `typeof(x)` stripped of its type parameters.
+"""
+basetypeof(x::T) where {T} = Base.typename(T).wrapper