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

examples/contract_bp.jl

@@ -0,0 +1,40 @@
+# https://docs.julialang.org/en/v1/base/base/#Base.names
+function imported_names(M::Module)
+  return setdiff(names(M; imported=true), names(M; imported=false))
+end
+
+function imported_functions(M::Module)
+  return map(f -> getfield(M, f), imported_names(M))
+end
+
+using FileUtils: replace_in_files
+
+function prepend_imports(M::Module)
+  fs = imported_functions(M)
+  # foreach(f -> @show((parentmodule(f), f)), fs)
+  src_path = joinpath(pkgdir(M), "src")
+  @show src_path
+  for f in fs
+    function_name = last(split(string(f), "."))
+    module_name = last(split(string(parentmodule(f)), "."))
+    annotated_function_name = module_name * "." * function_name
+    @show function_name, annotated_function_name
+    replace_in_files(
+      src_path,
+      "function $(function_name)(" => "function $(annotated_function_name)(";
+      recursive=true,
+      ignore_dirs=[".git"],
+      showdiffs=false,
+    )
+    replace_in_files(
+      src_path,
+      "\n$(function_name)(" => "\n$(annotated_function_name)(";
+      recursive=true,
+      ignore_dirs=[".git"],
+      showdiffs=false,
+    )
+  end
+end
+
+using ITensorNetworks: ITensorNetworks
+prepend_imports(ITensorNetworks)

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,21 @@ 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)[]
+inner_network(x::AbstractITensorNetwork; kwargs...) = inner_network(x, x; kwargs...)
+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 +888,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 +912,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,7 +6,7 @@ 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()
@@ -21,12 +21,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
 
+#TODO: Define a version of this that works for QN supporting tensors
 function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor})
   lhs, rhs = contract(message_a), contract(message_b)
-  return 0.5 *
-         norm((denseblocks(lhs) / sum(diag(lhs))) - (denseblocks(rhs) / sum(diag(rhs))))
+  tr_lhs = length(inds(lhs)) == 1 ? sum(lhs) : sum(diag(lhs))
+  tr_rhs = length(inds(rhs)) == 1 ? sum(rhs) : sum(diag(rhs))
+  return 0.5 * norm((denseblocks(lhs) / tr_lhs) - (denseblocks(rhs) / tr_rhs))
 end
 
 struct BeliefPropagationCache{PTN,MTS,DM}
@@ -47,8 +55,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 +132,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(
@@ -215,7 +229,7 @@ function update(
   verbose=false,
   kwargs...,
 )
-  compute_error = !isnothing(tol)
+  compute_error = !isnothing(tol) && !hasqns(tensornetwork(bp_cache))
   diff = compute_error ? Ref(0.0) : nothing
   if isnothing(maxiter)
     error("You need to specify a number of iterations for BP!")
@@ -251,3 +265,27 @@ 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)
+  return [region_scalar(bp_cache, pv) for pv in pvs]
+end
+
+function edge_scalars(bp_cache::BeliefPropagationCache, pes::Vector)
+  return [region_scalar(bp_cache, pe) for pe in pes]
+end
+
+function scalar_factors_quotient(bp_cache::BeliefPropagationCache)
+  pvs = partitionvertices(partitioned_tensornetwork(bp_cache))
+  pes = partitionedges(partitioned_tensornetwork(bp_cache))
+  return vertex_scalars(bp_cache, pvs), edge_scalars(bp_cache, pes)
+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