Express BP in terms of SweepIterator interface

Introduce `BeliefPropagationProblem` wrapper to hold the cache and the
error `diff` field.

Also simplifies some kwargs wrangling.

Author: Jack Dunham

Project.toml

@@ -10,6 +10,7 @@ BackendSelection = "680c2d7c-f67a-4cc9-ae9c-da132b1447a5"
 DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
@@ -33,6 +34,7 @@ BackendSelection = "0.1.6"
 DataGraphs = "0.2.7"
 Dictionaries = "0.4.5"
 Graphs = "1.13.1"
+ITensorBase = "0.2.14"
 LinearAlgebra = "1.10"
 MacroTools = "0.5.16"
 NamedDimsArrays = "0.8"

src/ITensorNetworksNext.jl

@@ -9,5 +9,6 @@ include("iterators.jl")
 
 include("beliefpropagation/abstractbeliefpropagationcache.jl")
 include("beliefpropagation/beliefpropagationcache.jl")
+include("beliefpropagation/beliefpropagationproblem.jl")
 
 end

src/beliefpropagation/beliefpropagationcache.jl

@@ -1,9 +1,7 @@
-using DiagonalArrays: delta
 using Dictionaries: Dictionary, set!, delete!
 using Graphs: AbstractGraph, is_tree, connected_components
 using NamedGraphs.GraphsExtensions: default_root_vertex, forest_cover, post_order_dfs_edges
 using ITensorBase: ITensor, dim
-using TypeParameterAccessors: unwrap_array_type, unwrap_array, parenttype
 
 struct BeliefPropagationCache{V, N <: AbstractDataGraph{V}} <:
     AbstractBeliefPropagationCache{V}
@@ -13,9 +11,8 @@ end
 
 messages(bp_cache::BeliefPropagationCache) = bp_cache.messages
 network(bp_cache::BeliefPropagationCache) = bp_cache.network
-default_messages() = Dictionary()
 
-BeliefPropagationCache(network) = BeliefPropagationCache(network, default_messages())
+BeliefPropagationCache(network) = BeliefPropagationCache(network, Dictionary())
 
 function Base.copy(bp_cache::BeliefPropagationCache)
     return BeliefPropagationCache(copy(network(bp_cache)), copy(messages(bp_cache)))
@@ -33,16 +30,15 @@ function setmessage!(bp_cache::BeliefPropagationCache, e::AbstractEdge, message)
     return bp_cache
 end
 
-function message(bp_cache::AbstractBeliefPropagationCache, edge::AbstractEdge; kwargs...)
+function message(bp_cache::BeliefPropagationCache, edge::AbstractEdge; kwargs...)
     ms = messages(bp_cache)
     return get(() -> default_message(bp_cache, edge; kwargs...), ms, edge)
 end
 
-function messages(bp_cache::AbstractBeliefPropagationCache, edges::Vector{<:AbstractEdge})
+function messages(bp_cache::BeliefPropagationCache, edges::Vector{<:AbstractEdge})
     return [message(bp_cache, e) for e in edges]
 end
 
-default_bp_maxiter(g::AbstractGraph) = is_tree(g) ? 1 : nothing
 #Forward onto the network
 for f in [
         :(Graphs.vertices),
@@ -62,11 +58,6 @@ for f in [
     end
 end
 
-#TODO: Get subgraph working on an ITensorNetwork to overload this directly
-function default_bp_edge_sequence(bp_cache::BeliefPropagationCache)
-    return forest_cover_edge_sequence(underlying_graph(bp_cache))
-end
-
 function factors(tn::AbstractTensorNetwork, vertex)
     return [tn[vertex]]
 end
@@ -91,33 +82,6 @@ function default_message(tn::AbstractTensorNetwork, edge::AbstractEdge)
     return t
 end
 
-#Algorithmic defaults
-default_update_alg(bp_cache::BeliefPropagationCache) = "bp"
-default_message_update_alg(bp_cache::BeliefPropagationCache) = "contract"
-default_normalize(::Algorithm"contract") = true
-default_sequence_alg(::Algorithm"contract") = "optimal"
-function set_default_kwargs(alg::Algorithm"contract")
-    normalize = get(alg, :normalize, default_normalize(alg))
-    sequence_alg = get(alg, :sequence_alg, default_sequence_alg(alg))
-    return Algorithm("contract"; normalize, sequence_alg)
-end
-function set_default_kwargs(alg::Algorithm"adapt_update")
-    _alg = set_default_kwargs(get(alg, :alg, Algorithm("contract")))
-    return Algorithm("adapt_update"; adapt = alg.adapt, alg = _alg)
-end
-default_verbose(::Algorithm"bp") = false
-default_tol(::Algorithm"bp") = nothing
-function set_default_kwargs(alg::Algorithm"bp", bp_cache::BeliefPropagationCache)
-    verbose = get(alg, :verbose, default_verbose(alg))
-    maxiter = get(alg, :maxiter, default_bp_maxiter(bp_cache))
-    edge_sequence = get(alg, :edge_sequence, default_bp_edge_sequence(bp_cache))
-    tol = get(alg, :tol, default_tol(alg))
-    message_update_alg = set_default_kwargs(
-        get(alg, :message_update_alg, Algorithm(default_message_update_alg(bp_cache)))
-    )
-    return Algorithm("bp"; verbose, maxiter, edge_sequence, tol, message_update_alg)
-end
-
 #TODO: Update message etc should go here...
 function updated_message(
         alg::Algorithm"contract", bp_cache::BeliefPropagationCache, edge::AbstractEdge
@@ -141,85 +105,21 @@ function updated_message(
     return updated_message
 end
 
-function updated_message(
-        bp_cache::BeliefPropagationCache,
-        edge::AbstractEdge;
-        alg = default_message_update_alg(bpc),
-        kwargs...,
+function default_algorithm(
+        ::Type{<:Algorithm"contract"}; normalize = true, sequence_alg = "optimal"
     )
-    return updated_message(set_default_kwargs(Algorithm(alg; kwargs...)), bp_cache, edge)
+    return Algorithm("contract"; normalize, sequence_alg)
 end
-
-function update_message!(
-        message_update_alg::Algorithm, bp_cache::BeliefPropagationCache, edge::AbstractEdge
+function default_algorithm(
+        ::Type{<:Algorithm"adapt_update"}; adapt, alg = default_algorithm(Algorithm"contract")
     )
-    return setmessage!(bp_cache, edge, updated_message(message_update_alg, bp_cache, edge))
+    return Algorithm("adapt_update"; adapt, alg)
 end
 
-"""
-Do a sequential update of the message tensors on `edges`
-"""
-function update_iteration(
-        alg::Algorithm"bp",
-        bpc::AbstractBeliefPropagationCache,
-        edges::Vector;
-        (update_diff!) = nothing,
-    )
-    bpc = copy(bpc)
-    for e in edges
-        prev_message = !isnothing(update_diff!) ? message(bpc, e) : nothing
-        update_message!(alg.message_update_alg, bpc, e)
-        if !isnothing(update_diff!)
-            update_diff![] += message_diff(message(bpc, e), prev_message)
-        end
-    end
-    return bpc
-end
-
-"""
-Do parallel updates between groups of edges of all message tensors
-Currently we send the full message tensor data struct to update for each edge_group. But really we only need the
-mts relevant to that group.
-"""
-function update_iteration(
-        alg::Algorithm"bp",
-        bpc::AbstractBeliefPropagationCache,
-        edge_groups::Vector{<:Vector{<:AbstractEdge}};
-        (update_diff!) = nothing,
+function update_message!(
+        message_update_alg::Algorithm, bpc::BeliefPropagationCache, edge::AbstractEdge
     )
-    new_mts = empty(messages(bpc))
-    for edges in edge_groups
-        bpc_t = update_iteration(alg.kwargs.message_update_alg, bpc, edges; (update_diff!))
-        for e in edges
-            set!(new_mts, e, message(bpc_t, e))
-        end
-    end
-    return set_messages(bpc, new_mts)
-end
-
-"""
-More generic interface for update, with default params
-"""
-function update(alg::Algorithm"bp", bpc::AbstractBeliefPropagationCache)
-    compute_error = !isnothing(alg.tol)
-    if isnothing(alg.maxiter)
-        error("You need to specify a number of iterations for BP!")
-    end
-    for i in 1:alg.maxiter
-        diff = compute_error ? Ref(0.0) : nothing
-        bpc = update_iteration(alg, bpc, alg.edge_sequence; (update_diff!) = diff)
-        if compute_error && (diff.x / length(alg.edge_sequence)) <= alg.tol
-            if alg.verbose
-                println("BP converged to desired precision after $i iterations.")
-            end
-            break
-        end
-    end
-    return bpc
-end
-
-function update(bpc::AbstractBeliefPropagationCache; alg = default_update_alg(bpc), kwargs...)
-    return update(set_default_kwargs(Algorithm(alg; kwargs...), bpc), bpc)
+    return setmessage!(bpc, edge, updated_message(message_update_alg, bpc, edge))
 end
 
 #Edge sequence stuff
@@ -234,4 +134,4 @@ function forest_cover_edge_sequence(g::AbstractGraph; root_vertex = default_root
         end
     end
     return edges
-end
+end

src/beliefpropagation/beliefpropagationproblem.jl

@@ -0,0 +1,85 @@
+mutable struct BeliefPropagationProblem{V, Cache <: AbstractBeliefPropagationCache{V}} <:
+    AbstractProblem
+    const cache::Cache
+    diff::Union{Nothing, Float64}
+end
+
+function default_algorithm(
+        ::Type{<:Algorithm"bp"},
+        bpc::BeliefPropagationCache;
+        verbose = false,
+        tol = nothing,
+        edge_sequence = forest_cover_edge_sequence(underlying_graph(bpc)),
+        message_update_alg = default_algorithm(Algorithm"contract"),
+        maxiter = is_tree(bpc) ? 1 : nothing,
+    )
+    return Algorithm("bp"; verbose, tol, edge_sequence, message_update_alg, maxiter)
+end
+
+function compute!(iter::RegionIterator{<:BeliefPropagationProblem})
+    prob = iter.problem
+
+    edge_group, kwargs = current_region_plan(iter)
+
+    new_message_tensors = map(edge_group) do edge
+        old_message = message(prob.cache, edge)
+
+        new_message = updated_message(kwargs.message_update_alg, prob.cache, edge)
+
+        if !isnothing(prob.diff)
+            # TODO: Define `message_diff`
+            prob.diff += message_diff(new_message, old_message)
+        end
+
+        return new_message
+    end
+
+    foreach(edge_group, new_message_tensors) do edge, new_message
+        setmessage!(prob.cache, edge, new_message)
+    end
+
+    return iter
+end
+
+function region_plan(
+        prob::BeliefPropagationProblem; root_vertex = default_root_vertex, sweep_kwargs...
+    )
+    edges = forest_cover_edge_sequence(underlying_graph(prob.cache); root_vertex)
+
+    plan = map(edges) do e
+        return [e] => (; sweep_kwargs...)
+    end
+
+    return plan
+end
+
+function update(bpc::AbstractBeliefPropagationCache; kwargs...)
+    return update(default_algorithm(Algorithm"bp", bpc; kwargs...), bpc)
+end
+function update(alg::Algorithm"bp", bpc)
+    compute_error = !isnothing(alg.tol)
+
+    diff = compute_error ? 0.0 : nothing
+
+    prob = BeliefPropagationProblem(bpc, diff)
+
+    iter = SweepIterator(prob, alg.maxiter; compute_error, getfield(alg, :kwargs)...)
+
+    for _ in iter
+        if compute_error && prob.diff <= alg.tol
+            break
+        end
+    end
+
+    if alg.verbose && compute_error
+        if prob.diff <= alg.tol
+            println("BP converged to desired precision after $(iter.which_sweep) iterations.")
+        else
+            println(
+                "BP failed to converge to precision $(alg.tol), got $(prob.diff) after $(iter.which_sweep) iterations",
+            )
+        end
+    end
+
+    return bpc
+end