Fix mutation behavior of position

src/treetensornetworks/projttns/abstractprojttn.jl

@@ -131,13 +131,43 @@ end
 function position(
   P::AbstractProjTTN{V}, psi::TTN{V}, pos::Union{Vector{<:V},NamedEdge{V}}
 ) where {V}
-  # shift position
   P = shift_position(P, pos)
-  # remove internal edges (out of place)
+  P = invalidate_environments(P)
+  P = make_environments(P, psi)
+  return P
+end
+
+function position!(
+  P::AbstractProjTTN{V}, psi::TTN{V}, pos::Union{Vector{<:V},NamedEdge{V}}
+) where {V}
+  P = shift_position(P, pos)
+  P = invalidate_environments(P)
+  make_environments!(P, psi, e)
+  return P
+end
+
+function invalidate_environments(P::AbstractProjTTN)
   ie = internal_edges(P)
   newenvskeys = filter(!in(ie), keys(P.environments))
-  P = ProjTTN(pos, P.H, getindices(P.environments, newenvskeys))
-  # make all environments surrounding new position
+  P = ProjTTN(P.pos, P.H, getindices(P.environments, newenvskeys))
+  return P
+end
+
+function invalidate_environment(P::AbstractProjTTN{V}, e::NamedEdge{V}) where {V}
+  newenvskeys = filter(isequal(e), keys(P.environments))
+  P = ProjTTN(P.pos, P.H, getindices(P.environments, newenvskeys))
+  return P
+end
+
+function make_environments(P::AbstractProjTTN{V}, psi::TTN{V}) where {V}
+  P = copy(P)
+  for e in incident_edges(P)
+    set!(P.environments, e, make_environment(P, psi, e))
+  end
+  return P
+end
+
+function make_environments!(P::AbstractProjTTN{V}, psi::TTN{V}) where {V}
   for e in incident_edges(P)
     make_environment!(P, psi, e)
   end

src/treetensornetworks/projttns/projttn.jl

@@ -11,7 +11,7 @@ function ProjTTN(H::TTN)
   return ProjTTN(vertices(H), H, Dictionary{edgetype(H),ITensor}())
 end
 
-copy(P::ProjTTN) = ProjTTN(P.pos, copy(P.H), copy(P.environments))
+copy(P::ProjTTN) = ProjTTN(P.pos, copy(P.H), NamedGraphs.copy_keys_values(P.environments))
 
 # trivial if we choose to specify position as above; only kept to allow using alongside
 # ProjMPO
@@ -30,23 +30,7 @@ function make_environment!(P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V})::ITensor
   if haskey(P.environments, e)
     env = environment(P, e)
   else
-    if is_leaf(underlying_graph(P), src(e))
-      # leaves are easy
-      env = psi[src(e)] * P.H[src(e)] * dag(prime(psi[src(e)]))
-    else
-      # construct by contracting neighbors
-      neighbor_envs = ITensor[]
-      for n in setdiff(neighbors(underlying_graph(P), src(e)), [dst(e)])
-        push!(neighbor_envs, make_environment!(P, psi, 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(psi[src(e)], frst, scnd, P.H[src(e)], dag(prime(psi[src(e)])), rst)
-      # TODO: actually use optimal contraction sequence here
-      env = reduce(*, itensor_map)
-    end
-    # cache
+    env = _compute_environment(P, psi, e)
     set!(P.environments, e, env)
   end
   @assert(
@@ -55,3 +39,42 @@ function make_environment!(P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V})::ITensor
   )
   return env
 end
+
+function make_environment(P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V})::ITensor where {V}
+  # invalidate environment for opposite edge direction if necessary
+  #P=copy(P)
+  reverse(e) ∈ incident_edges(P) || (P = invalidate_environment(P, reverse(e)))
+  # do nothing if valid environment already present
+  if haskey(P.environments, e)
+    env = environment(P, e)
+  else
+    env = _compute_environment(P, psi, e)
+  end
+  @assert(
+    hascommoninds(env, psi[src(e)]),
+    "Something went wrong, probably re-orthogonalized this edge in the same direction twice!"
+  )
+  return env
+end
+
+function _compute_environment(
+  P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V}
+)::ITensor where {V}
+  if is_leaf(underlying_graph(P), src(e))
+    # leaves are easy
+    env = psi[src(e)] * P.H[src(e)] * dag(prime(psi[src(e)]))
+  else
+    # construct by contracting neighbors
+    neighbor_envs = ITensor[]
+    for n in setdiff(neighbors(underlying_graph(P), src(e)), [dst(e)])
+      push!(neighbor_envs, make_environment!(P, psi, 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(psi[src(e)], frst, scnd, P.H[src(e)], dag(prime(psi[src(e)])), rst)
+    # TODO: actually use optimal contraction sequence here
+    env = reduce(*, itensor_map)
+  end
+  return env
+end

test/test_treetensornetworks/test_position.jl

@@ -41,3 +41,4 @@ using Test
     ITensors.disable_auto_fermion()
   end
 end
+nothing