BP Caching overhauls

Author: Jack Dunham

src/beliefpropagation/abstractbeliefpropagationcache.jl

@@ -1,117 +1,124 @@
-abstract type AbstractBeliefPropagationCache{V} <: AbstractGraph{V} end
+using Graphs: AbstractGraph, AbstractEdge
+using DataGraphs: AbstractDataGraph, edge_data, vertex_data, edge_data_eltype
+using NamedGraphs.GraphsExtensions: boundary_edges
+using NamedGraphs.PartitionedGraphs: QuotientView, QuotientEdge, parent
 
-#Interface
-factor(bp_cache::AbstractBeliefPropagationCache, vertex) = not_implemented()
-setfactor!(bp_cache::AbstractBeliefPropagationCache, vertex, factor) = not_implemented()
-messages(bp_cache::AbstractBeliefPropagationCache) = not_implemented()
-message(bp_cache::AbstractBeliefPropagationCache, edge::AbstractEdge) = not_implemented()
-function default_message(bp_cache::AbstractBeliefPropagationCache, edge::AbstractEdge)
-    return not_implemented()
-end
-default_messages(bp_cache::AbstractBeliefPropagationCache) = not_implemented()
-function setmessage!(bp_cache::AbstractBeliefPropagationCache, edge::AbstractEdge, message)
-    return not_implemented()
+messages(::AbstractGraph) = not_implemented()
+messages(bp_cache::AbstractDataGraph) = edge_data(bp_cache)
+messages(bp_cache::AbstractGraph, edges) = [message(bp_cache, e) for e in edges]
+
+message(bp_cache::AbstractGraph, edge::AbstractEdge) = messages(bp_cache)[edge]
+
+deletemessage!(bp_cache::AbstractGraph, edge) = not_implemented()
+function deletemessage!(bp_cache::AbstractDataGraph, edge)
+    ms = messages(bp_cache)
+    delete!(ms, edge)
+    return bp_cache
 end
-function deletemessage!(bp_cache::AbstractBeliefPropagationCache, edge::AbstractEdge)
-    return not_implemented()
+
+function deletemessages!(bp_cache::AbstractGraph, edges = edges(bp_cache))
+    for e in edges
+        deletemessage!(bp_cache, e)
+    end
+    return bp_cache
 end
-function rescale_messages(
-        bp_cache::AbstractBeliefPropagationCache, edges::Vector{<:AbstractEdge}; kwargs...
-    )
-    return not_implemented()
+
+setmessage!(bp_cache::AbstractGraph, edge, message) = not_implemented()
+function setmessage!(bp_cache::AbstractDataGraph, edge, message)
+    ms = messages(bp_cache)
+    set!(ms, edge, message)
+    return bp_cache
 end
-function rescale_vertices(
-        bp_cache::AbstractBeliefPropagationCache, vertices::Vector; kwargs...
-    )
-    return not_implemented()
+function setmessage!(bp_cache::QuotientView, edge, message)
+    setmessages!(parent(bp_cache), QuotientEdge(edge), message)
+    return bp_cache
 end
 
-function vertex_scalar(bp_cache::AbstractBeliefPropagationCache, vertex; kwargs...)
-    return not_implemented()
+function setmessages!(bp_cache::AbstractGraph, edge::QuotientEdge, message)
+    for e in edges(bp_cache, edge)
+        setmessage!(parent(bp_cache), e, message[e])
+    end
+    return bp_cache
 end
-function edge_scalar(
-        bp_cache::AbstractBeliefPropagationCache, edge::AbstractEdge; kwargs...
-    )
-    return not_implemented()
+function setmessages!(bpc_dst::AbstractGraph, bpc_src::AbstractGraph, edges)
+    for e in edges
+        setmessage!(bpc_dst, e, message(bpc_src, e))
+    end
+    return bpc_dst
 end
 
-#Graph functionality needed
-Graphs.vertices(bp_cache::AbstractBeliefPropagationCache) = not_implemented()
-Graphs.edges(bp_cache::AbstractBeliefPropagationCache) = not_implemented()
-function NamedGraphs.GraphsExtensions.boundary_edges(
-        bp_cache::AbstractBeliefPropagationCache, vertices; kwargs...
-    )
-    return not_implemented()
+factors(bpc::AbstractGraph) = vertex_data(bpc)
+factors(bpc::AbstractGraph, vertices::Vector) = [factor(bpc, v) for v in vertices]
+factors(bpc::AbstractGraph{V}, vertex::V) where {V} = factors(bpc, V[vertex])
+
+factor(bpc::AbstractGraph, vertex) = factors(bpc)[vertex]
+
+setfactor!(bpc::AbstractGraph, vertex, factor) = not_implemented()
+function setfactor!(bpc::AbstractDataGraph, vertex, factor)
+    fs = factors(bpc)
+    set!(fs, vertex, factor)
+    return bpc
 end
 
-#Functions derived from the interface
-function setmessages!(bp_cache::AbstractBeliefPropagationCache, edges, messages)
-    for (e, m) in zip(edges)
-        setmessage!(bp_cache, e, m)
-    end
-    return
+function region_scalar(bp_cache::AbstractGraph, edge::AbstractEdge)
+    return message(bp_cache, edge) * message(bp_cache, reverse(edge))
 end
 
-function deletemessages!(
-        bp_cache::AbstractBeliefPropagationCache, edges::Vector{<:AbstractEdge} = edges(bp_cache)
-    )
-    for e in edges
-        deletemessage!(bp_cache, e)
-    end
-    return bp_cache
+function region_scalar(bp_cache::AbstractGraph, vertex)
+
+    messages = incoming_messages(bp_cache, vertex)
+    state = factors(bp_cache, vertex)
+
+    return reduce(*, messages) * reduce(*, state)
 end
 
-function vertex_scalars(
-        bp_cache::AbstractBeliefPropagationCache, vertices = Graphs.vertices(bp_cache); kwargs...
-    )
-    return map(v -> region_scalar(bp_cache, v; kwargs...), vertices)
+message_type(bpc::AbstractGraph) = message_type(typeof(bpc))
+message_type(G::Type{<:AbstractGraph}) = eltype(Base.promote_op(messages, G))
+message_type(type::Type{<:AbstractDataGraph}) = edge_data_eltype(type)
+
+function vertex_scalars(bp_cache::AbstractGraph, vertices = vertices(bp_cache))
+    return map(v -> region_scalar(bp_cache, v), vertices)
 end
 
-function edge_scalars(
-        bp_cache::AbstractBeliefPropagationCache, edges = Graphs.edges(bp_cache); kwargs...
-    )
-    return map(e -> region_scalar(bp_cache, e; kwargs...), edges)
+function edge_scalars(bp_cache::AbstractGraph, edges = edges(bp_cache))
+    return map(e -> region_scalar(bp_cache, e), edges)
 end
 
-function scalar_factors_quotient(bp_cache::AbstractBeliefPropagationCache)
+function scalar_factors_quotient(bp_cache::AbstractGraph)
     return vertex_scalars(bp_cache), edge_scalars(bp_cache)
 end
 
-function incoming_messages(
-        bp_cache::AbstractBeliefPropagationCache, vertices::Vector{<:Any}; ignore_edges = []
-    )
-    b_edges = NamedGraphs.GraphsExtensions.boundary_edges(bp_cache, vertices; dir = :in)
+function incoming_messages(bp_cache::AbstractGraph, vertices; ignore_edges = [])
+    b_edges = boundary_edges(bp_cache, [vertices;]; dir = :in)
     b_edges = !isempty(ignore_edges) ? setdiff(b_edges, ignore_edges) : b_edges
     return messages(bp_cache, b_edges)
 end
 
-function incoming_messages(bp_cache::AbstractBeliefPropagationCache, vertex; kwargs...)
-    return incoming_messages(bp_cache, [vertex]; kwargs...)
-end
+default_messages(::AbstractGraph) = not_implemented()
 
 #Adapt interface for changing device
-function map_messages(f, bp_cache::AbstractBeliefPropagationCache, es = edges(bp_cache))
-    bp_cache = copy(bp_cache)
+map_messages(f, bp_cache, es = edges(bp_cache)) = map_messages!(f, copy(bp_cache), es)
+function map_messages!(f, bp_cache, es = edges(bp_cache))
     for e in es
         setmessage!(bp_cache, e, f(message(bp_cache, e)))
     end
     return bp_cache
 end
-function map_factors(f, bp_cache::AbstractBeliefPropagationCache, vs = vertices(bp_cache))
-    bp_cache = copy(bp_cache)
+
+map_factors(f, bp_cache, vs = vertices(bp_cache)) = map_factors!(f, copy(bp_cache), vs)
+function map_factors!(f, bp_cache, vs = vertices(bp_cache))
     for v in vs
         setfactor!(bp_cache, v, f(factor(bp_cache, v)))
     end
     return bp_cache
 end
-function adapt_messages(to, bp_cache::AbstractBeliefPropagationCache, args...)
-    return map_messages(adapt(to), bp_cache, args...)
-end
-function adapt_factors(to, bp_cache::AbstractBeliefPropagationCache, args...)
-    return map_factors(adapt(to), bp_cache, args...)
-end
 
-function freenergy(bp_cache::AbstractBeliefPropagationCache)
+adapt_messages(to, bp_cache, es = edges(bp_cache)) = map_messages(adapt(to), bp_cache, es)
+adapt_factors(to, bp_cache, vs = vertices(bp_cache)) = map_factors(adapt(to), bp_cache, vs)
+
+abstract type AbstractBeliefPropagationCache{V, ED} <: AbstractDataGraph{V, Nothing, ED} end
+
+function free_energy(bp_cache::AbstractBeliefPropagationCache)
     numerator_terms, denominator_terms = scalar_factors_quotient(bp_cache)
     if any(t -> real(t) < 0, numerator_terms)
         numerator_terms = complex.(numerator_terms)
@@ -123,29 +130,4 @@ function freenergy(bp_cache::AbstractBeliefPropagationCache)
     any(iszero, denominator_terms) && return -Inf
     return sum(log.(numerator_terms)) - sum(log.((denominator_terms)))
 end
-
-function partitionfunction(bp_cache::AbstractBeliefPropagationCache)
-    return exp(freenergy(bp_cache))
-end
-
-function rescale_messages(bp_cache::AbstractBeliefPropagationCache, edge::AbstractEdge)
-    return rescale_messages(bp_cache, [edge])
-end
-
-function rescale_messages(bp_cache::AbstractBeliefPropagationCache)
-    return rescale_messages(bp_cache, edges(bp_cache))
-end
-
-function rescale_vertices(bpc::AbstractBeliefPropagationCache; kwargs...)
-    return rescale_vertices(bpc, collect(vertices(bpc)); kwargs...)
-end
-
-function rescale_vertex(bpc::AbstractBeliefPropagationCache, vertex; kwargs...)
-    return rescale_vertices(bpc, [vertex]; kwargs...)
-end
-
-function rescale(bpc::AbstractBeliefPropagationCache, args...; kwargs...)
-    bpc = rescale_messages(bpc)
-    bpc = rescale_partitions(bpc, args...; kwargs...)
-    return bpc
-end
+partitionfunction(bp_cache::AbstractBeliefPropagationCache) = exp(free_energy(bp_cache))