Generalize ProjTTNSum, replace ProjTTNApply with ProjOuterProdTTN

src/ITensorNetworks.jl

@@ -116,7 +116,7 @@ include(joinpath("treetensornetworks", "opsum_to_ttn.jl"))
 include(joinpath("treetensornetworks", "projttns", "abstractprojttn.jl"))
 include(joinpath("treetensornetworks", "projttns", "projttn.jl"))
 include(joinpath("treetensornetworks", "projttns", "projttnsum.jl"))
-include(joinpath("treetensornetworks", "projttns", "projttn_apply.jl"))
+include(joinpath("treetensornetworks", "projttns", "projouterprodttn.jl"))
 include(joinpath("treetensornetworks", "solvers", "solver_utils.jl"))
 include(joinpath("treetensornetworks", "solvers", "update_step.jl"))
 include(joinpath("treetensornetworks", "solvers", "alternating_update.jl"))

src/exports.jl

@@ -76,7 +76,7 @@ export AbstractITensorNetwork,
   random_ttn,
   ProjTTN,
   ProjTTNSum,
-  ProjTTNApply,
+  ProjOuterProdTTN,
   set_nsite,
   position,
   finite_state_machine,

src/imports.jl

@@ -96,6 +96,10 @@ import ITensors:
   scalartype,
   #adding
   add
+
+import ITensors.LazyApply:
+  # extracting terms from a sum
+  terms
 #Algorithm
 Algorithm
 

src/solvers/contract.jl

@@ -9,11 +9,6 @@ function contract_updater(
   region_kwargs,
   updater_kwargs,
 )
-  v = ITensor(true)
-  projected_operator = projected_operator![]
-  for j in sites(projected_operator)
-    v *= init_state(projected_operator)[j]
-  end
-  vp = contract(projected_operator, v)
-  return vp, (;)
+  P = projected_operator![]
+  return contract_ket(P, ITensor(true)), (;)
 end

src/treetensornetworks/projttns/abstractprojttn.jl

@@ -65,30 +65,13 @@ function _separate_first_two(V::Vector)
   return frst, scnd, rst
 end
 
-function contract(P::AbstractProjTTN, v::ITensor)::ITensor
-  environments = ITensor[environment(P, edge) for edge in incident_edges(P)]
-  # manual heuristic for contraction order fixing: for each site in ProjTTN, apply up to
-  # two environments, then TTN tensor, then other environments
-  if on_edge(P)
-    itensor_map = environments
-  else
-    itensor_map = Union{ITensor,OneITensor}[] # TODO: will a Hamiltonian TTN tensor ever be a OneITensor?
-    for s in sites(P)
-      site_envs = filter(hascommoninds(operator(P)[s]), environments)
-      frst, scnd, rst = _separate_first_two(site_envs)
-      site_tensors = vcat(frst, scnd, operator(P)[s], rst)
-      append!(itensor_map, site_tensors)
-    end
-  end
-  # TODO: actually use optimal contraction sequence here
-  Hv = v
-  for it in itensor_map
-    Hv *= it
-  end
-  return Hv
+projected_operator_tensors(P::AbstractProjTTN) = error("Not implemented.")
+
+function contract(P::AbstractProjTTN, v::ITensor)
+  return foldl(*, projected_operator_tensors(P); init=v)
 end
 
-function product(P::AbstractProjTTN, v::ITensor)::ITensor
+function product(P::AbstractProjTTN, v::ITensor)
   Pv = contract(P, v)
   if order(Pv) != order(v)
     error(

src/treetensornetworks/projttns/projouterprodttn.jl

@@ -0,0 +1,115 @@
+struct ProjOuterProdTTN{V} <: AbstractProjTTN{V}
+  pos::Union{Vector{<:V},NamedEdge{V}}
+  internal_state::TTN{V}
+  operator::TTN{V}
+  environments::Dictionary{NamedEdge{V},ITensor}
+end
+
+environments(p::ProjOuterProdTTN) = p.environments
+operator(p::ProjOuterProdTTN) = p.operator
+underlying_graph(p::ProjOuterProdTTN) = underlying_graph(operator(p))
+pos(p::ProjOuterProdTTN) = p.pos
+internal_state(p::ProjOuterProdTTN) = p.internal_state
+
+function ProjOuterProdTTN(internal_state::AbstractTTN, operator::AbstractTTN)
+  return ProjOuterProdTTN(
+    vertextype(operator)[],
+    internal_state,
+    operator,
+    Dictionary{edgetype(operator),ITensor}(),
+  )
+end
+
+function copy(P::ProjOuterProdTTN)
+  return ProjOuterProdTTN(
+    P.pos, copy(internal_state(P)), copy(operator(P)), copy(environments(P))
+  )
+end
+
+function set_nsite(P::ProjOuterProdTTN, nsite)
+  return P
+end
+
+function shift_position(P::ProjOuterProdTTN, pos)
+  return ProjOuterProdTTN(pos, internal_state(P), operator(P), environments(P))
+end
+
+function set_environments(p::ProjOuterProdTTN, environments)
+  return ProjOuterProdTTN(pos(p), internal_state(p), operator(p), environments)
+end
+
+set_environment(p::ProjOuterProdTTN, edge, env) = set_environment!(copy(p), edge, env)
+function set_environment!(p::ProjOuterProdTTN, edge, env)
+  set!(environments(p), edge, env)
+  return p
+end
+
+function make_environment(P::ProjOuterProdTTN, state::AbstractTTN, e::AbstractEdge)
+  # invalidate environment for opposite edge direction if necessary
+  reverse(e) ∈ incident_edges(P) || (P = invalidate_environment(P, reverse(e)))
+  # do nothing if valid environment already present
+  if !haskey(environments(P), e)
+    if is_leaf(underlying_graph(P), src(e))
+      # leaves are easy
+      env = internal_state(P)[src(e)] * operator(P)[src(e)] * dag(state[src(e)])
+    else
+      # construct by contracting neighbors
+      neighbor_envs = ITensor[]
+      for n in setdiff(neighbors(underlying_graph(P), src(e)), [dst(e)])
+        P = make_environment(P, state, edgetype(P)(n, src(e)))
+        push!(neighbor_envs, environment(P, edgetype(P)(n, src(e))))
+      end
+      # manually heuristic for contraction order: two environments, site tensors, then
+      # other environments
+      frst, scnd, rst = _separate_first_two(neighbor_envs)
+      itensor_map = vcat(
+        internal_state(P)[src(e)], frst, scnd, operator(P)[src(e)], dag(state[src(e)]), rst
+      ) # no prime here in comparison to the same routine for Projttn
+      # TODO: actually use optimal contraction sequence here
+      env = reduce(*, itensor_map)
+    end
+    P = set_environment(P, e, env)
+  end
+  @assert(
+    hascommoninds(environment(P, e), state[src(e)]),
+    "Something went wrong, probably re-orthogonalized this edge in the same direction twice!"
+  )
+  return P
+end
+
+function projected_operator_tensors(P::ProjOuterProdTTN)
+  environments = ITensor[environment(P, edge) for edge in incident_edges(P)]
+  # manual heuristic for contraction order fixing: for each site in ProjTTN, apply up to
+  # two environments, then TTN tensor, then other environments
+  itensor_map = Union{ITensor,OneITensor}[] # TODO: will a Hamiltonian TTN tensor ever be a OneITensor?
+  for j in sites(P)
+    push!(itensor_map, internal_state(P)[j])
+  end
+  if on_edge(P)
+    append!(itensor_map, environments)
+  else
+    for s in sites(P)
+      site_envs = filter(hascommoninds(operator(P)[s]), environments)
+      frst, scnd, rst = _separate_first_two(site_envs)
+      site_tensors = vcat(frst, scnd, operator(P)[s], rst)
+      append!(itensor_map, site_tensors)
+    end
+  end
+  return itensor_map
+end
+
+function contract_ket(P::ProjOuterProdTTN, v::ITensor)
+  itensor_map = projected_operator_tensors(P)
+  for t in itensor_map
+    v *= t
+  end
+  return v
+end
+
+function contract(P::ProjOuterProdTTN, x::ITensor)
+  return conj(contract_ket(P, x)) * contract_ket(P, ITensor(true))
+end
+
+function product(P::ProjOuterProdTTN, v::ITensor)
+  return noprime(contract(P, v))
+end

src/treetensornetworks/projttns/projttn.jl

@@ -68,3 +68,21 @@ function make_environment(P::ProjTTN, state::AbstractTTN, e::AbstractEdge)
   )
   return P
 end
+
+function projected_operator_tensors(P::ProjTTN)
+  environments = ITensor[environment(P, edge) for edge in incident_edges(P)]
+  # manual heuristic for contraction order fixing: for each site in ProjTTN, apply up to
+  # two environments, then TTN tensor, then other environments
+  if on_edge(P)
+    itensor_map = environments
+  else
+    itensor_map = Union{ITensor,OneITensor}[] # TODO: will a Hamiltonian TTN tensor ever be a OneITensor?
+    for s in sites(P)
+      site_envs = filter(hascommoninds(operator(P)[s]), environments)
+      frst, scnd, rst = _separate_first_two(site_envs)
+      site_tensors = vcat(frst, scnd, operator(P)[s], rst)
+      append!(itensor_map, site_tensors)
+    end
+  end
+  return itensor_map
+end

src/treetensornetworks/projttns/projttnsum.jl

@@ -1,59 +1,51 @@
 """
 ProjTTNSum
 """
-struct ProjTTNSum{V}
-  pm::Vector{ProjTTN{V}}
-  function ProjTTNSum(pm::Vector{ProjTTN{V}}) where {V}
-    return new{V}(pm)
+struct ProjTTNSum{V} <: AbstractProjTTN{V}
+  terms::Vector{AbstractProjTTN{V}}
+  function ProjTTNSum(terms::Vector{<:AbstractProjTTN{V}}) where {V}
+    return new{V}(terms)
   end
 end
 
-copy(P::ProjTTNSum) = ProjTTNSum(copy.(P.pm))
+terms(P::ProjTTNSum) = P.terms
 
-ProjTTNSum(ttnos::Vector{<:TTN}) = ProjTTNSum([ProjTTN(M) for M in ttnos])
+copy(P::ProjTTNSum) = ProjTTNSum(copy.(terms(P)))
 
-ProjTTNSum(Ms::TTN{V}...) where {V} = ProjTTNSum([Ms...])
+ProjTTNSum(operators::Vector{<:AbstractTTN}) = ProjTTNSum(ProjTTN.(operators))
 
-on_edge(P::ProjTTNSum) = on_edge(P.pm[1])
+on_edge(P::ProjTTNSum) = on_edge(terms(P)[1])
 
-nsite(P::ProjTTNSum) = nsite(P.pm[1])
+nsite(P::ProjTTNSum) = nsite(terms(P)[1])
 
 function set_nsite(Ps::ProjTTNSum, nsite)
-  return ProjTTNSum(map(M -> set_nsite(M, nsite), Ps.pm))
+  return ProjTTNSum(map(p -> set_nsite(p, nsite), terms(Ps)))
 end
 
-underlying_graph(P::ProjTTNSum) = underlying_graph(P.pm[1])
+underlying_graph(P::ProjTTNSum) = underlying_graph(terms(P)[1])
 
-Base.length(P::ProjTTNSum) = length(P.pm[1])
+Base.length(P::ProjTTNSum) = length(terms(P)[1])
 
-sites(P::ProjTTNSum) = sites(P.pm[1])
+sites(P::ProjTTNSum) = sites(terms(P)[1])
 
-incident_edges(P::ProjTTNSum) = incident_edges(P.pm[1])
+incident_edges(P::ProjTTNSum) = incident_edges(terms(P)[1])
 
-internal_edges(P::ProjTTNSum) = internal_edges(P.pm[1])
+internal_edges(P::ProjTTNSum) = internal_edges(terms(P)[1])
 
-function product(P::ProjTTNSum, v::ITensor)::ITensor
-  Pv = product(P.pm[1], v)
-  for n in 2:length(P.pm)
-    Pv += product(P.pm[n], v)
-  end
-  return Pv
-end
+product(P::ProjTTNSum, v::ITensor) = noprime(contract(P, v))
+
+contract(P::ProjTTNSum, v::ITensor) = sum(p -> contract(p, v), terms(P))
+
+contract_ket(P::ProjTTNSum, v::ITensor) = sum(p -> contract_ket(p, v), terms(P))
 
 function Base.eltype(P::ProjTTNSum)
-  elT = eltype(P.pm[1])
-  for n in 2:length(P.pm)
-    elT = promote_type(elT, eltype(P.pm[n]))
-  end
-  return elT
+  return mapreduce(eltype, promote_type, terms(P))
 end
 
 (P::ProjTTNSum)(v::ITensor) = product(P, v)
 
-Base.size(P::ProjTTNSum) = size(P.pm[1])
+Base.size(P::ProjTTNSum) = size(terms(P)[1])
 
-function position(
-  P::ProjTTNSum{V}, psi::TTN{V}, pos::Union{Vector{<:V},NamedEdge{V}}
-) where {V}
-  return ProjTTNSum(map(M -> position(M, psi, pos), P.pm))
+function position(P::ProjTTNSum, psi::AbstractTTN, pos)
+  return ProjTTNSum(map(M -> position(M, psi, pos), terms(P)))
 end