Implement inner using BP

README.md

@@ -39,7 +39,7 @@ julia> using ITensorNetworks: ITensorNetwork, siteinds
 julia> using NamedGraphs: named_grid, subgraph
 
 julia> tn = ITensorNetwork(named_grid(4); link_space=2)
-ITensorNetworks.ITensorNetwork{Int64} with 4 vertices:
+ITensorNetwork{Int64} with 4 vertices:
 4-element Vector{Int64}:
  1
  2
@@ -90,7 +90,7 @@ and here is a similar example for making a tensor network on a grid (a tensor pr
 
 ```julia
 julia> tn = ITensorNetwork(named_grid((2, 2)); link_space=2)
-ITensorNetworks.ITensorNetwork{Tuple{Int64, Int64}} with 4 vertices:
+ITensorNetwork{Tuple{Int64, Int64}} with 4 vertices:
 4-element Vector{Tuple{Int64, Int64}}:
  (1, 1)
  (2, 1)
@@ -125,7 +125,7 @@ julia> neighbors(tn, (1, 2))
  (2, 2)
 
 julia> tn_1 = subgraph(v -> v[1] == 1, tn)
-ITensorNetworks.ITensorNetwork{Tuple{Int64, Int64}} with 2 vertices:
+ITensorNetwork{Tuple{Int64, Int64}} with 2 vertices:
 2-element Vector{Tuple{Int64, Int64}}:
  (1, 1)
  (1, 2)
@@ -139,7 +139,7 @@ with vertex data:
  (1, 2) │ ((dim=2|id=598|"1×1,1×2"), (dim=2|id=683|"1×2,2×2"))
 
 julia> tn_2 = subgraph(v -> v[1] == 2, tn)
-ITensorNetworks.ITensorNetwork{Tuple{Int64, Int64}} with 2 vertices:
+ITensorNetwork{Tuple{Int64, Int64}} with 2 vertices:
 2-element Vector{Tuple{Int64, Int64}}:
  (2, 1)
  (2, 2)
@@ -184,7 +184,7 @@ and edge data:
 0-element Dictionaries.Dictionary{NamedGraphs.NamedEdge{Int64}, Vector{ITensors.Index}}
 
 julia> tn1 = ITensorNetwork(s; link_space=2)
-ITensorNetworks.ITensorNetwork{Int64} with 3 vertices:
+ITensorNetwork{Int64} with 3 vertices:
 3-element Vector{Int64}:
  1
  2
@@ -201,7 +201,7 @@ with vertex data:
  3 │ ((dim=2|id=558|"S=1/2,Site,n=3"), (dim=2|id=430|"2,3"))
 
 julia> tn2 = ITensorNetwork(s; link_space=2)
-ITensorNetworks.ITensorNetwork{Int64} with 3 vertices:
+ITensorNetwork{Int64} with 3 vertices:
 3-element Vector{Int64}:
  1
  2

src/ITensorNetworks.jl

@@ -13,7 +13,6 @@ include("opsum.jl")
 include("sitetype.jl")
 include("abstractitensornetwork.jl")
 include("contraction_sequences.jl")
-include("expect.jl")
 include("tebd.jl")
 include("itensornetwork.jl")
 include("mincut.jl")
@@ -64,6 +63,8 @@ include("solvers/contract.jl")
 include("solvers/linsolve.jl")
 include("solvers/sweep_plans/sweep_plans.jl")
 include("apply.jl")
+include("inner.jl")
+include("expect.jl")
 include("environment.jl")
 include("exports.jl")
 include("ModelHamiltonians/ModelHamiltonians.jl")

src/abstractitensornetwork.jl

@@ -728,66 +728,23 @@ function flatten_networks(
   return flatten_networks(flatten_networks(tn1, tn2; kwargs...), tn3, tn_tail...; kwargs...)
 end
 
-#Ideally this will dispatch to inner_network but this is a temporary fast version for now
-function norm_network(tn::AbstractITensorNetwork)
-  tnbra = rename_vertices(v -> (v, 1), data_graph(tn))
-  tndag = copy(tn)
-  for v in vertices(tndag)
-    setindex_preserve_graph!(tndag, dag(tndag[v]), v)
-  end
-  tnket = rename_vertices(v -> (v, 2), data_graph(prime(tndag; sites=[])))
-  tntn = ITensorNetwork(union(tnbra, tnket))
-  for v in vertices(tn)
-    if !isempty(commoninds(tntn[(v, 1)], tntn[(v, 2)]))
-      add_edge!(tntn, (v, 1) => (v, 2))
-    end
-  end
-  return tntn
+function inner_network(x::AbstractITensorNetwork, y::AbstractITensorNetwork; kwargs...)
+  return BilinearFormNetwork(x, y; kwargs...)
 end
 
-# TODO: Use or replace with `flatten_networks`
 function inner_network(
-  tn1::AbstractITensorNetwork,
-  tn2::AbstractITensorNetwork;
-  map_bra_linkinds=sim,
-  combine_linkinds=false,
-  flatten=combine_linkinds,
-  kwargs...,
+  x::AbstractITensorNetwork, A::AbstractITensorNetwork, y::AbstractITensorNetwork; kwargs...
 )
-  @assert issetequal(vertices(tn1), vertices(tn2))
-  tn1 = map_bra_linkinds(tn1; sites=[])
-  inner_net = ⊗(dag(tn1), tn2; kwargs...)
-  if flatten
-    for v in vertices(tn1)
-      inner_net = contract(inner_net, (v, 2) => (v, 1); merged_vertex=v)
-    end
-  end
-  if combine_linkinds
-    inner_net = ITensorNetworks.combine_linkinds(inner_net)
-  end
-  return inner_net
+  return BilinearFormNetwork(A, x, y; kwargs...)
 end
 
-# TODO: Rename `inner`.
-function contract_inner(
-  ϕ::AbstractITensorNetwork,
-  ψ::AbstractITensorNetwork;
-  sequence=nothing,
-  contraction_sequence_kwargs=(;),
-)
-  tn = inner_network(ϕ, ψ; combine_linkinds=true)
-  if isnothing(sequence)
-    sequence = contraction_sequence(tn; contraction_sequence_kwargs...)
-  end
-  return contract(tn; sequence)[]
+# TODO: We should make this pass to inner_network and then to BiLinearForm. 
+# Parts of the code (tests relying on norm_sqr being two layer and the gauging code
+#  which relies on specific message tensors) currently would break in that case so we need to resolve
+function norm_sqr_network(ψ::AbstractITensorNetwork)
+  return disjoint_union("bra" => dag(prime(ψ; sites=[])), "ket" => ψ)
 end
 
-# TODO: rename `sqnorm` to match https://github.com/JuliaStats/Distances.jl,
-# or `norm_sqr` to match `LinearAlgebra.norm_sqr`
-norm_sqr(ψ::AbstractITensorNetwork; sequence) = contract_inner(ψ, ψ; sequence)
-
-norm_sqr_network(ψ::AbstractITensorNetwork; kwargs...) = inner_network(ψ, ψ; kwargs...)
-
 #
 # Printing
 #
@@ -933,7 +890,7 @@ function ITensorMPS.add(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork
 
   #Create vertices of tn12 as direct sum of tn1[v] and tn2[v]. Work out the matching indices by matching edges. Make index tags those of tn1[v]
   for v in vertices(tn1)
-    @assert siteinds(tn1, v) == siteinds(tn2, v)
+    @assert issetequal(siteinds(tn1, v), siteinds(tn2, v))
 
     e1_v = filter(x -> src(x) == v || dst(x) == v, edges_tn1)
     e2_v = filter(x -> src(x) == v || dst(x) == v, edges_tn2)
@@ -957,3 +914,5 @@ function ITensorMPS.add(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork
 end
 
 Base.:+(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork) = add(tn1, tn2)
+
+ITensors.hasqns(tn::AbstractITensorNetwork) = any(v -> hasqns(tn[v]), vertices(tn))

src/approx_itensornetwork/binary_tree_partition.jl

@@ -130,6 +130,6 @@ function _partition(
   return rename_vertices(par, name_map)
 end
 
-function _partition(tn::ITensorNetwork, inds_btree::DataGraph; alg::String)
+function _partition(tn::ITensorNetwork, inds_btree::DataGraph; alg)
   return _partition(Algorithm(alg), tn, inds_btree)
 end

src/caches/beliefpropagationcache.jl

@@ -6,10 +6,11 @@ using NamedGraphs: PartitionVertex
 using LinearAlgebra: diag
 using ITensors: dir
 using ITensors.ITensorMPS: ITensorMPS
-using NamedGraphs: boundary_partitionedges
+using NamedGraphs: boundary_partitionedges, partitionvertices, partitionedges
 
 default_message(inds_e) = ITensor[denseblocks(delta(inds_e))]
 default_messages(ptn::PartitionedGraph) = Dictionary()
+default_message_norm(m::ITensor) = norm(m)
 function default_message_update(contract_list::Vector{ITensor}; kwargs...)
   sequence = optimal_contraction_sequence(contract_list)
   updated_messages = contract(contract_list; sequence, kwargs...)
@@ -21,12 +22,20 @@ end
   return default_bp_maxiter(undirected_graph(underlying_graph(g)))
 end
 default_partitioned_vertices(ψ::AbstractITensorNetwork) = group(v -> v, vertices(ψ))
+function default_partitioned_vertices(f::AbstractFormNetwork)
+  return group(v -> original_state_vertex(f, v), vertices(f))
+end
 default_cache_update_kwargs(cache) = (; maxiter=20, tol=1e-5)
+function default_cache_construction_kwargs(alg::Algorithm"bp", ψ::AbstractITensorNetwork)
+  return (; partitioned_vertices=default_partitioned_vertices(ψ))
+end
 
-function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor})
+function message_diff(
+  message_a::Vector{ITensor}, message_b::Vector{ITensor}; message_norm=default_message_norm
+)
   lhs, rhs = contract(message_a), contract(message_b)
-  return 0.5 *
-         norm((denseblocks(lhs) / sum(diag(lhs))) - (denseblocks(rhs) / sum(diag(rhs))))
+  norm_lhs, norm_rhs = message_norm(lhs), message_norm(rhs)
+  return 0.5 * norm((denseblocks(lhs) / norm_lhs) - (denseblocks(rhs) / norm_rhs))
 end
 
 struct BeliefPropagationCache{PTN,MTS,DM}
@@ -47,8 +56,14 @@ function BeliefPropagationCache(tn, partitioned_vertices; kwargs...)
   return BeliefPropagationCache(ptn; kwargs...)
 end
 
-function BeliefPropagationCache(tn; kwargs...)
-  return BeliefPropagationCache(tn, default_partitioning(tn); kwargs...)
+function BeliefPropagationCache(
+  tn; partitioned_vertices=default_partitioned_vertices(tn), kwargs...
+)
+  return BeliefPropagationCache(tn, partitioned_vertices; kwargs...)
+end
+
+function cache(alg::Algorithm"bp", tn; kwargs...)
+  return BeliefPropagationCache(tn; kwargs...)
 end
 
 function partitioned_tensornetwork(bp_cache::BeliefPropagationCache)
@@ -118,7 +133,7 @@ function environment(
 )
   bpes = boundary_partitionedges(bp_cache, partition_vertices; dir=:in)
   ms = messages(bp_cache, setdiff(bpes, ignore_edges))
-  return reduce(vcat, ms; init=[])
+  return reduce(vcat, ms; init=ITensor[])
 end
 
 function environment(
@@ -216,11 +231,11 @@ function update(
   kwargs...,
 )
   compute_error = !isnothing(tol)
-  diff = compute_error ? Ref(0.0) : nothing
   if isnothing(maxiter)
     error("You need to specify a number of iterations for BP!")
   end
   for i in 1:maxiter
+    diff = compute_error ? Ref(0.0) : nothing
     bp_cache = update(bp_cache, edges; (update_diff!)=diff, kwargs...)
     if compute_error && (diff.x / length(edges)) <= tol
       if verbose
@@ -251,3 +266,31 @@ end
 function update_factor(bp_cache, vertex, factor)
   return update_factors(bp_cache, [vertex], ITensor[factor])
 end
+
+function region_scalar(bp_cache::BeliefPropagationCache, pv::PartitionVertex)
+  incoming_mts = environment(bp_cache, [pv])
+  local_state = factor(bp_cache, pv)
+  return scalar(vcat(incoming_mts, local_state))
+end
+
+function region_scalar(bp_cache::BeliefPropagationCache, pe::PartitionEdge)
+  return scalar(vcat(message(bp_cache, pe), message(bp_cache, reverse(pe))))
+end
+
+function vertex_scalars(
+  bp_cache::BeliefPropagationCache,
+  pvs::Vector=partitionvertices(partitioned_tensornetwork(bp_cache)),
+)
+  return [region_scalar(bp_cache, pv) for pv in pvs]
+end
+
+function edge_scalars(
+  bp_cache::BeliefPropagationCache,
+  pes::Vector=partitionedges(partitioned_tensornetwork(bp_cache)),
+)
+  return [region_scalar(bp_cache, pe) for pe in pes]
+end
+
+function scalar_factors_quotient(bp_cache::BeliefPropagationCache)
+  return vertex_scalars(bp_cache), edge_scalars(bp_cache)
+end

src/contract.jl

@@ -1,10 +1,10 @@
 using NamedGraphs: vertex_to_parent_vertex
-using ITensors: ITensor
+using ITensors: ITensor, scalar
 using ITensors.ContractionSequenceOptimization: deepmap
 using ITensors.NDTensors: NDTensors, Algorithm, @Algorithm_str, contract
 using LinearAlgebra: normalize!
 
-function NDTensors.contract(tn::AbstractITensorNetwork; alg::String="exact", kwargs...)
+function NDTensors.contract(tn::AbstractITensorNetwork; alg="exact", kwargs...)
   return contract(Algorithm(alg), tn; kwargs...)
 end
 
@@ -15,6 +15,10 @@ function NDTensors.contract(
   return contract(Vector{ITensor}(tn); sequence=sequence_linear_index, kwargs...)
 end
 
+function NDTensors.contract(alg::Algorithm"exact", tensors::Vector{ITensor}; kwargs...)
+  return contract(tensors; kwargs...)
+end
+
 function NDTensors.contract(
   alg::Union{Algorithm"density_matrix",Algorithm"ttn_svd"},
   tn::AbstractITensorNetwork;
@@ -36,3 +40,61 @@ function contract_density_matrix(
   end
   return out
 end
+
+function ITensors.scalar(
+  alg::Algorithm, tn::Union{AbstractITensorNetwork,Vector{ITensor}}; kwargs...
+)
+  return contract(alg, tn; kwargs...)[]
+end
+
+function ITensors.scalar(
+  tn::Union{AbstractITensorNetwork,Vector{ITensor}}; alg="exact", kwargs...
+)
+  return scalar(Algorithm(alg), tn; kwargs...)
+end
+
+function logscalar(
+  tn::Union{AbstractITensorNetwork,Vector{ITensor}}; alg="exact", kwargs...
+)
+  return logscalar(Algorithm(alg), tn; kwargs...)
+end
+
+function logscalar(
+  alg::Algorithm"exact", tn::Union{AbstractITensorNetwork,Vector{ITensor}}; kwargs...
+)
+  s = scalar(alg, tn; kwargs...)
+  s = real(s) < 0 ? complex(s) : s
+  return log(s)
+end
+
+function logscalar(
+  alg::Algorithm,
+  tn::AbstractITensorNetwork;
+  (cache!)=nothing,
+  cache_construction_kwargs=default_cache_construction_kwargs(alg, tn),
+  update_cache=isnothing(cache!),
+  cache_update_kwargs=default_cache_update_kwargs(cache!),
+)
+  if isnothing(cache!)
+    cache! = Ref(cache(alg, tn; cache_construction_kwargs...))
+  end
+
+  if update_cache
+    cache![] = update(cache![]; cache_update_kwargs...)
+  end
+
+  numerator_terms, denominator_terms = scalar_factors_quotient(cache![])
+  numerator_terms =
+    any(t -> real(t) < 0, numerator_terms) ? complex.(numerator_terms) : numerator_terms
+  denominator_terms = if any(t -> real(t) < 0, denominator_terms)
+    complex.(denominator_terms)
+  else
+    denominator_terms
+  end
+
+  return sum(log.(numerator_terms)) - sum(log.((denominator_terms)))
+end
+
+function ITensors.scalar(alg::Algorithm"bp", tn::AbstractITensorNetwork; kwargs...)
+  return exp(logscalar(alg, tn; kwargs...))
+end

src/expect.jl

@@ -14,13 +14,13 @@ function ITensorMPS.expect(
   # ElT = ishermitian(ITensors.op(op, s[vertices[1]])) ? real(ElT) : ElT
   res = Dictionary(vertices, Vector{ElT}(undef, length(vertices)))
   if isnothing(sequence)
-    sequence = contraction_sequence(inner_network(ψ, ψ; flatten=true))
+    sequence = contraction_sequence(inner_network(ψ, ψ))
   end
-  normψ² = norm_sqr(ψ; sequence)
+  normψ² = norm_sqr(ψ; alg="exact", sequence)
   for v in vertices
     O = ITensor(Op(op, v), s)
     Oψ = apply(O, ψ; cutoff, maxdim, ortho)
-    res[v] = contract_inner(ψ, Oψ; sequence) / normψ²
+    res[v] = inner(ψ, Oψ; alg="exact", sequence) / normψ²
   end
   return res
 end
@@ -36,12 +36,12 @@ function ITensorMPS.expect(
   s = siteinds(ψ)
   # h⃗ = Vector{ITensor}(ℋ, s)
   if isnothing(sequence)
-    sequence = contraction_sequence(inner_network(ψ, ψ; flatten=true))
+    sequence = contraction_sequence(inner_network(ψ, ψ))
   end
   h⃗ψ = [apply(hᵢ, ψ; cutoff, maxdim, ortho) for hᵢ in ITensors.terms(ℋ)]
-  ψhᵢψ = [contract_inner(ψ, hᵢψ; sequence) for hᵢψ in h⃗ψ]
+  ψhᵢψ = [inner(ψ, hᵢψ; alg="exact", sequence) for hᵢψ in h⃗ψ]
   ψh⃗ψ = sum(ψhᵢψ)
-  ψψ = norm_sqr(ψ; sequence)
+  ψψ = norm_sqr(ψ; alg="exact", sequence)
   return ψh⃗ψ / ψψ
 end