Simplify interface, fix tests

Author: JoeyT1994

src/caches/abstractbeliefpropagationcache.jl

@@ -25,23 +25,6 @@ function data_graph_type(bpc::AbstractBeliefPropagationCache)
 end
 data_graph(bpc::AbstractBeliefPropagationCache) = data_graph(tensornetwork(bpc))
 
-function message_update(alg::Algorithm"contract", contract_list::Vector{ITensor})
-  sequence = contraction_sequence(contract_list; alg=alg.kwargs.sequence_alg)
-  updated_messages = contract(contract_list; sequence)
-  message_norm = norm(updated_messages)
-  if alg.kwargs.normalize && !iszero(message_norm)
-    updated_messages /= message_norm
-  end
-  return ITensor[updated_messages]
-end
-
-function message_update(alg::Algorithm"adapt_update", contract_list::Vector{ITensor})
-  adapted_contract_list = alg.kwargs.adapt.(contract_list)
-  updated_messages = message_update(alg.kwargs.alg, adapted_contract_list)
-  dtype = mapreduce(datatype, promote_type, contract_list)
-  return map(adapt(dtype), updated_messages)
-end
-
 #TODO: Take `dot` without precontracting the messages to allow scaling to more complex messages
 function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor})
   lhs, rhs = contract(message_a), contract(message_b)
@@ -144,22 +127,28 @@ function incoming_messages(
 end
 
 #Adapt interface for changing device
-function map_messages(f, bpc::AbstractBeliefPropagationCache)
+function map_messages(
+  f, bpc::AbstractBeliefPropagationCache, pes=collect(keys(messages(bpc)))
+)
   bpc = copy(bpc)
-  for pe in keys(messages(bpc))
+  for pe in pes
     set_message!(bpc, pe, f.(message(bpc, pe)))
   end
   return bpc
 end
-function map_factors(f, bpc::AbstractBeliefPropagationCache)
+function map_factors(f, bpc::AbstractBeliefPropagationCache, vs=vertices(bpc))
   bpc = copy(bpc)
-  for v in vertices(bpc)
+  for v in vs
     @preserve_graph bpc[v] = f(bpc[v])
   end
   return bpc
 end
-adapt_messages(to, bpc::AbstractBeliefPropagationCache) = map_messages(adapt(to), bpc)
-adapt_factors(to, bpc::AbstractBeliefPropagationCache) = map_factors(adapt(to), bpc)
+function adapt_messages(to, bpc::AbstractBeliefPropagationCache, args...)
+  map_messages(adapt(to), bpc, args...)
+end
+function adapt_factors(to, bpc::AbstractBeliefPropagationCache, args...)
+  map_factors(adapt(to), bpc, args...)
+end
 
 function Adapt.adapt_structure(to, bpc::AbstractBeliefPropagationCache)
   bpc = adapt_messages(to, bpc)
@@ -257,33 +246,49 @@ function delete_message(bpc::AbstractBeliefPropagationCache, pe::PartitionEdge)
   return delete_messages(bpc, [pe])
 end
 
-"""
-Compute message tensor as product of incoming mts and local state
-"""
 function updated_message(
-  bpc::AbstractBeliefPropagationCache,
-  edge::PartitionEdge;
-  message_update_alg=default_message_update_alg(bpc),
-  kwargs...,
+  alg::Algorithm"contract", bpc::AbstractBeliefPropagationCache, edge::PartitionEdge
 )
   vertex = src(edge)
   incoming_ms = incoming_messages(bpc, vertex; ignore_edges=PartitionEdge[reverse(edge)])
   state = factors(bpc, vertex)
+  contract_list = ITensor[incoming_ms; state]
+  sequence = contraction_sequence(contract_list; alg=alg.kwargs.sequence_alg)
+  updated_messages = contract(contract_list; sequence)
+  message_norm = norm(updated_messages)
+  if alg.kwargs.normalize && !iszero(message_norm)
+    updated_messages /= message_norm
+  end
+  return ITensor[updated_messages]
+end
 
-  return message_update(message_update_alg, ITensor[incoming_ms; state]; kwargs...)
+function updated_message(
+  alg::Algorithm"adapt_update", bpc::AbstractBeliefPropagationCache, edge::PartitionEdge
+)
+  incoming_pes = setdiff(
+    boundary_partitionedges(bpc, [src(edge)]; dir=:in), [reverse(edge)]
+  )
+  adapted_bpc = adapt_messages(alg.kwargs.adapt, bpc, incoming_pes)
+  adapted_bpc = adapt_factors(alg.kwargs.adapt, bpc, vertices(bpc, src(edge)))
+  updated_messages = updated_message(alg.kwargs.alg, adapted_bpc, edge)
+  dtype = mapreduce(datatype, promote_type, message(bpc, edge))
+  return map(adapt(dtype), updated_messages)
 end
 
-function update(
-  alg::Algorithm"bp", bpc::AbstractBeliefPropagationCache, edge::PartitionEdge; kwargs...
+function update_message(
+  message_update_alg::Algorithm,
+  bpc::AbstractBeliefPropagationCache,
+  edge::PartitionEdge;
+  kwargs...,
 )
-  return set_message(bpc, edge, updated_message(bpc, edge; kwargs...))
+  return set_message(bpc, edge, updated_message(message_update_alg, bpc, edge; kwargs...))
 end
 
 """
 Do a sequential update of the message tensors on `edges`
 """
-function update(
-  alg::Algorithm,
+function update_one_iteration(
+  alg::Algorithm"bp",
   bpc::AbstractBeliefPropagationCache,
   edges::Vector;
   (update_diff!)=nothing,
@@ -292,7 +297,7 @@ function update(
   bpc = copy(bpc)
   for e in edges
     prev_message = !isnothing(update_diff!) ? message(bpc, e) : nothing
-    bpc = update(alg, bpc, e; kwargs...)
+    bpc = update_message(alg.kwargs.message_update_alg, bpc, e; kwargs...)
     if !isnothing(update_diff!)
       update_diff![] += message_diff(message(bpc, e), prev_message)
     end
@@ -305,15 +310,15 @@ 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(
+function update_one_iteration(
   alg::Algorithm,
   bpc::AbstractBeliefPropagationCache,
   edge_groups::Vector{<:Vector{<:PartitionEdge}};
   kwargs...,
 )
   new_mts = empty(messages(bpc))
   for edges in edge_groups
-    bpc_t = update(alg, bpc, edges; kwargs...)
+    bpc_t = update_one_iteration(alg.kwargs.message_update_alg, bpc, edges; kwargs...)
     for e in edges
       set!(new_mts, e, message(bpc_t, e))
     end
@@ -324,27 +329,17 @@ end
 """
 More generic interface for update, with default params
 """
-function update(
-  alg::Algorithm,
-  bpc::AbstractBeliefPropagationCache;
-  message_update_alg=default_message_update_alg(bpc),
-  kwargs...,
-)
+function update(alg::Algorithm, bpc::AbstractBeliefPropagationCache; kwargs...)
   compute_error = !isnothing(alg.kwargs.tol)
   if isnothing(alg.kwargs.maxiter)
     error("You need to specify a number of iterations for BP!")
   end
   for i in 1:alg.kwargs.maxiter
     diff = compute_error ? Ref(0.0) : nothing
-    bpc = update(
-      alg,
-      bpc,
-      alg.kwargs.edge_sequence;
-      (update_diff!)=diff,
-      message_update_alg=set_default_kwargs(message_update_alg),
-      kwargs...,
+    bpc = update_one_iteration(
+      alg, bpc, alg.kwargs.edge_sequence; (update_diff!)=diff, kwargs...
     )
-    if compute_error && (diff.x / length(edges)) <= alg.kwargs.tol
+    if compute_error && (diff.x / length(alg.kwargs.edge_sequence)) <= alg.kwargs.tol
       if alg.kwargs.verbose
         println("BP converged to desired precision after $i iterations.")
       end
@@ -355,7 +350,7 @@ function update(
 end
 
 function update(bpc::AbstractBeliefPropagationCache; alg=default_update_alg(bpc), kwargs...)
-  return update(set_default_kwargs(alg, bpc), bpc; kwargs...)
+  return update(set_default_kwargs(Algorithm(alg; kwargs...), bpc), bpc)
 end
 
 function rescale_messages(

src/caches/beliefpropagationcache.jl

@@ -68,18 +68,8 @@ function Base.copy(bp_cache::BeliefPropagationCache)
   )
 end
 
-function default_update_alg(bp_cache::BeliefPropagationCache)
-  Algorithm(
-    "bp";
-    verbose=false,
-    maxiter=default_bp_maxiter(bp_cache),
-    edge_sequence=default_bp_edge_sequence(bp_cache),
-    tol=nothing,
-  )
-end
-function default_message_update_alg(bp_cache::BeliefPropagationCache)
-  return Algorithm("contract"; normalize=true, sequence_alg="optimal")
-end
+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")
@@ -88,9 +78,8 @@ function set_default_kwargs(alg::Algorithm"contract")
   return Algorithm("contract"; normalize, sequence_alg)
 end
 function set_default_kwargs(alg::Algorithm"adapt_update")
-  return Algorithm(
-    "adapt_update"; adapt=alg.kwargs.adapt, alg=set_default_kwargs(alg.kwargs.alg)
-  )
+  _alg = set_default_kwargs(get(alg.kwargs, :alg, Algorithm("contract")))
+  return Algorithm("adapt_update"; adapt=alg.kwargs.adapt, alg=_alg)
 end
 default_verbose(::Algorithm"bp") = false
 default_tol(::Algorithm"bp") = nothing
@@ -99,7 +88,10 @@ function set_default_kwargs(alg::Algorithm"bp", bp_cache::BeliefPropagationCache
   maxiter = get(alg.kwargs, :maxiter, default_bp_maxiter(bp_cache))
   edge_sequence = get(alg.kwargs, :edge_sequence, default_bp_edge_sequence(bp_cache))
   tol = get(alg.kwargs, :tol, default_tol(alg))
-  return Algorithm("bp"; verbose, maxiter, edge_sequence, tol)
+  message_update_alg = set_default_kwargs(
+    get(alg.kwargs, :message_update_alg, Algorithm(default_message_update_alg(bp_cache)))
+  )
+  return Algorithm("bp"; verbose, maxiter, edge_sequence, tol, message_update_alg)
 end
 
 function default_bp_maxiter(bp_cache::BeliefPropagationCache)

test/test_apply.jl

@@ -30,12 +30,12 @@ using Test: @test, @testset
   ψψ = norm_sqr_network(ψ)
   #Simple Belief Propagation Grouping
   bp_cache = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ)))
-  bp_cache = update(bp_cache; alg=Algorithm("bp"; maxiter=20))
+  bp_cache = update(bp_cache; maxiter=20)
   envsSBP = environment(bp_cache, [(v1, "bra"), (v1, "ket"), (v2, "bra"), (v2, "ket")])
-  ψv = VidalITensorNetwork(ψ)
+  ψv = VidalITensorNetwork(ψ; cache_update_kwargs=(; maxiter=20))
   #This grouping will correspond to calculating the environments exactly (each column of the grid is a partition)
   bp_cache = BeliefPropagationCache(ψψ, group(v -> v[1][1], vertices(ψψ)))
-  bp_cache = update(bp_cache; alg=Algorithm("bp"; maxiter=20))
+  bp_cache = update(bp_cache; maxiter=20)
   envsGBP = environment(bp_cache, [(v1, "bra"), (v1, "ket"), (v2, "bra"), (v2, "ket")])
   inner_alg = "exact"
   ngates = 5

test/test_belief_propagation.jl

@@ -25,7 +25,8 @@ using ITensorNetworks:
   update_factor,
   updated_message,
   message_diff
-using ITensors: ITensors, ITensor, combiner, dag, inds, inner, op, prime, random_itensor
+using ITensors:
+  ITensors, ITensor, Algorithm, combiner, dag, inds, inner, op, prime, random_itensor
 using ITensorNetworks.ModelNetworks: ModelNetworks
 using ITensors.NDTensors: array
 using LinearAlgebra: eigvals, tr
@@ -62,7 +63,7 @@ using Test: @test, @testset
     @test bpc[vket] == new_A
     @test bpc[vbra] == new_A_dag
 
-    bpc = update(bpc; alg=Algorithm("bp"; maxiter=25, tol=eps(real(elt))))
+    bpc = update(bpc; alg="bp", maxiter=25, tol=eps(real(elt)))
     #Test messages are converged
     for pe in partitionedges(bpc)
       @test message_diff(updated_message(bpc, pe), message(bpc, pe)) < 10 * eps(real(elt))

test/test_expect.jl

@@ -36,7 +36,9 @@ using Test: @test, @testset
   cache_construction_kwargs = (;
     partitioned_vertices=group(v -> first(first(v)), quadratic_form_vertices)
   )
-  sz_bp = expect(ψ, "Sz"; alg="bp", cache_construction_kwargs)
+  sz_bp = expect(
+    ψ, "Sz"; alg="bp", cache_construction_kwargs, cache_update_kwargs=(; maxiter=20)
+  )
   sz_exact = expect(ψ, "Sz"; alg="exact")
   @test sz_bp ≈ sz_exact
 
@@ -47,7 +49,7 @@ using Test: @test, @testset
 
   ψ = ITensorNetwork(v -> isodd(sum(v)) ? "↑" : "↓", s)
 
-  sz_bp = expect(ψ, "Sz"; alg="bp")
+  sz_bp = expect(ψ, "Sz"; alg="bp", cache_update_kwargs=(; maxiter=20))
   sz_exact = expect(ψ, "Sz"; alg="exact")
   @test sz_bp ≈ sz_exact
 end

test/test_gauging.jl

@@ -29,7 +29,7 @@ using Test: @test, @testset
 
   # Move directly to vidal gauge
   ψ_vidal = VidalITensorNetwork(
-    ψ; cache_update_kwargs=(; alg=Algorithm("bp"; maxiter=30, verbose=true))
+    ψ; cache_update_kwargs=(; alg="bp", maxiter=30, verbose=true)
   )
   @test gauge_error(ψ_vidal) < 1e-8
 
@@ -39,10 +39,12 @@ using Test: @test, @testset
   bp_cache = cache_ref[]
 
   # Test we just did a gauge transform and didn't change the overall network
-  @test inner(ψ_symm, ψ) / sqrt(inner(ψ_symm, ψ_symm) * inner(ψ, ψ)) ≈ 1.0 atol = 1e-8
+  @test inner(ψ_symm, ψ; alg="exact") /
+        sqrt(inner(ψ_symm, ψ_symm; alg="exact") * inner(ψ, ψ; alg="exact")) ≈ 1.0 atol =
+    1e-8
 
   #Test all message tensors are approximately diagonal even when we keep running BP
-  bp_cache = update(bp_cache; alg=Algorithm("bp"; maxiter=10))
+  bp_cache = update(bp_cache; maxiter=10)
   for m_e in values(messages(bp_cache))
     @test diag_itensor(vector(diag(only(m_e))), inds(only(m_e))) ≈ only(m_e) atol = 1e-8
   end

test/test_normalize.jl

@@ -29,7 +29,7 @@ using Test: @test, @testset
   tn_r = rescale(tn; alg="exact")
   @test scalar(tn_r; alg="exact") ≈ 1.0
 
-  tn_r = rescale(tn; alg="bp")
+  tn_r = rescale(tn; alg="bp", cache_update_kwargs=(; maxiter=20))
   @test scalar(tn_r; alg="exact") ≈ 1.0
 
   #Now a state on a loopy graph
@@ -45,10 +45,12 @@ using Test: @test, @testset
   @test scalar(norm_sqr_network(ψ); alg="exact") ≈ 1.0
 
   ψIψ_bpc = Ref(BeliefPropagationCache(QuadraticFormNetwork(x)))
-  ψ = normalize(x; alg="bp", (cache!)=ψIψ_bpc, update_cache=true)
+  ψ = normalize(
+    x; alg="bp", (cache!)=ψIψ_bpc, update_cache=true, cache_update_kwargs=(; maxiter=20)
+  )
   ψIψ_bpc = ψIψ_bpc[]
   @test all(x -> x ≈ 1.0, edge_scalars(ψIψ_bpc))
   @test all(x -> x ≈ 1.0, vertex_scalars(ψIψ_bpc))
-  @test scalar(QuadraticFormNetwork(ψ); alg="bp") ≈ 1.0
+  @test scalar(QuadraticFormNetwork(ψ); alg="bp", cache_update_kwargs=(; maxiter=20)) ≈ 1.0
 end
 end